#!/usr/bin/env python
# coding: utf-8
import numpy as np
import matplotlib.pyplot as plt
import argparse
import pandas

def save_fig(fig:plt.figure, figName:str):
    fig.savefig(figName + '.pdf')
    fig.savefig(figName + '.png')
    fig.savefig(figName + '.svg')

# Channel Frequency Response
def channel_est(dataPath):
    binfile = dataPath + 'chestF0_ext_s1.m.bin'
    data = np.fromfile(binfile,dtype=np.int16)
    dataCmp = data[0::2] + 1j*data[1::2]
    fig = plt.gcf()
    plt.plot(np.abs(dataCmp))
    save_fig(fig, 'channel_est')
    plt.close(fig)

# Constellation
def constellation(dataPath):
    binfile = dataPath + 'rxF_comp_s1.m.bin'
    data = np.fromfile(binfile,dtype=np.int16)
    dataCmp = data[0::2] + 1j*data[1::2]
    fig = plt.gcf()
    plt.scatter(np.real(dataCmp),np.imag(dataCmp))
    save_fig(fig, 'constellation')
    plt.close(fig)

# LLR
def llr(dataPath):
    binfile = dataPath + 'rxF_llr.m.bin'
    data = np.fromfile(binfile,dtype=np.int16)
    fig = plt.gcf()
    plt.plot(data,'*')
    save_fig(fig, 'llr')
    plt.close(fig)

# BER from data in CSV files generated by nr_dlsim
def ber(dataPath, csvFiles):
    d0 = pandas.read_csv(dataPath + csvFiles[0] + '.csv')
    ax = d0.plot(x='SNR',y='channel_ber_0',logy=True,label=csvFiles[0])
    plt.ylabel('BER')
    for i in range(1,len(csvFiles)):
        d = pandas.read_csv(dataPath + csvFiles[i] + '.csv')
        d.plot(ax=ax,x='SNR',y='channel_ber_0',logy=True,label=csvFiles[i],grid=True)

    fig = ax.figure
    save_fig(fig, 'ber_compare')
    plt.close(fig)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="DL plot generator")
    parser.add_argument('--data-path', help='Path to directory containing data', default='../../cmake_targets/ran_build/build/')
    parser.add_argument('--csv-data', nargs='+' ,help='CSV file name without extension (generated from nr_dlsim with -Z option)')
    parser.add_argument('--ber-compare', help='Compare BER from multiple CSV files', action='store_true')

    args = parser.parse_args()

    channel_est(args.data_path)
    constellation(args.data_path)
    llr(args.data_path)
    if args.ber_compare:
        ber(args.data_path, args.csv_data)