#!/usr/bin/env python3


import numpy as np
import matplotlib.pyplot as plt
import json

from scipy import special


# Mimic import aa_generate_beacon as beacon
beacon = lambda: None

def read_snr_file(fname):
    with open(fname, 'r') as fp:
        return json.load(fp)

def read_tx_file(fname):
    with open(fname, 'r') as fp:
        data = json.load(fp)

    f_beacon = data['f_beacon']
    tx = data['tx']

    del data['f_beacon']
    del data['tx']

    return tx, f_beacon, data


beacon.snr_fname = 'snr.json'
beacon.tx_fname = 'tx.json'
beacon.read_snr_file = read_snr_file
beacon.read_tx_file = read_tx_file

def read_snr_directories(data_directories, tx_fname=beacon.tx_fname, snr_fname=beacon.snr_fname, phase_res_fname='phase_time_residuals.json'):
    # Read in snr
    snrs = np.zeros(len(data_directories), dtype=float)
    time_residuals = np.zeros( (len(snrs), 2), dtype=float)


    tx, f_beacon, _ = beacon.read_tx_file(path.join(data_directories[0], tx_fname))
    for i, data_dir in enumerate(data_directories):
        # Read SNR file
        snr_data = read_snr_file(path.join(data_dir, snr_fname))

        # Open antennas file
        with open(path.join(data_dir, phase_res_fname), 'r') as fp:
            time_res_data = json.load(fp)

        snrs[i] = snr_data['mean']

        time_residuals[i] = time_res_data['mean'], time_res_data['std']

    return f_beacon, snrs, time_residuals


## Math
def expectation(x, pdfx):
    dx = x[1] - x[0]
    return np.sum(x*pdfx*dx)

def variance(x, pdfx):
    mu = expectation(x, pdfx)
    dx = x[1] - x[0]
    return np.sum(x**2 *pdfx*dx) - mu**2

def phase_distribution(theta, sigma, s):
    k = s/sigma
    ct = np.cos(theta)
    st = np.sin(theta)
    pipi = 2*np.pi

    return (np.exp(-k**2/2)/pipi) \
            + (
                (pipi**-0.5)*k*np.exp(-(k*st)**2/2)
                * ((1.+special.erf(k*ct*2**-0.5))*ct/2)
            )

def phase_distribution_gauss(theta, sigma, s):
    k = s/sigma
    return 1/np.sqrt(2*np.pi) * k *np.exp( -(k*theta)**2/2)



if __name__ == "__main__":
    from os import path
    import sys

    import os
    import matplotlib
    if os.name == 'posix' and "DISPLAY" not in os.environ:
        matplotlib.use('Agg')

    from argparse import ArgumentParser
    parser = ArgumentParser()

    figsize = (12,8)

    # Multiple Data Directories
    parser.add_argument("-d", dest='data_directories', default=[], nargs='*')

    # Whether to show plots
    group1 = parser.add_mutually_exclusive_group(required=False)
    group1.add_argument('--show-plots', action='store_true', default=True, help='Default: %(default)s')
    group1.add_argument('--no-show-plots', dest='show-plots', action='store_false')
 
    # Figures directory
    parser.add_argument('--fig-dir', type=str, default='./figures', help='Set None to disable saving figures. (Default: %(default)s)')

    args = parser.parse_args()
    
    show_plots = args.show_plots

    if not args.data_directories:
        parser.error("At least one data directory should be specified.")

    f_beacon, snrs, time_residuals = read_snr_directories(args.data_directories)

    phase2time = lambda x: x/(2*np.pi*f_beacon)
    time2phase = lambda x: 2*np.pi*x*f_beacon

    fig, ax = plt.subplots(figsize=(8,6))
    ax.set_title("Beacon ({:.2f}MHz) Simulation".format(f_beacon*1e3))
    ax.set_xlabel('Signal to Noise ratio')
    ax.set_ylabel('Time Accuracy $\\sigma(t)$ [ns]')

    # group per directories per N
    if True:
        import re

        dirdict = {}
        N_re = re.compile(r'_N(\d+)_')

        for d in args.data_directories:
            m = N_re.findall(d)[0]

            if m not in dirdict:
                dirdict[m] = []

            dirdict[m].append(d)

    
        dirlist = dirdict.values()

    # plot data directories as one group
    else:
        dirlist = [args.data_directories]
 
    for dirlist in dirlist:
        f_beacon, snrs, time_residuals = read_snr_directories(dirlist)

        # plot data
        ax.plot(snrs*np.sqrt(np.pi/2), phase2time(time_residuals[:,1]), ls='none', marker='o')

    # Add secondary axis
    if True:
        if False and secondary_axis == 'time':
            functions = (phase2time, time2phase)
            label = 'Time Accuracy $\\sigma_t\\varphi/(2\\pi f_{beac})$ [ns]'
        else:
            functions = (time2phase, phase2time)
            label = 'Phase Accuracy $\\sigma_\\varphi$ [rad]'

        secax = ax.secondary_yaxis('right', functions=functions)
        secax.set_ylabel(label)
 
    # logscales
    if True:
        ax.set_xscale('log')
        ax.set_yscale('log')

    # plot phasor snr
    if True:
        thetas = np.linspace(-np.pi, np.pi, 500)
        sigma = 1
        _snr_min = min(10, min(snrs))
        _snr_max = min(100, max(snrs))

        if ax.get_xscale() == 'log': #log
            _snrs = np.logspace(np.log10(_snr_min), np.log10(_snr_max))
        else: #linear
            _snrs = np.linspace(_snr_min, _snr_max)

        # Phasor Rice
        phasor_pdfs = np.array([phase_distribution(thetas, sigma, s) for s in _snrs ])
        phasor_sigma = np.sqrt(np.array([variance(thetas, pdf) for pdf in phasor_pdfs]))

        ax.plot(_snrs, phase2time(phasor_sigma), label='Rice')

        if True: # plot a pdf
            fig2, ax2 = plt.subplots()
            for idx in [0, len(_snrs)//4, len(_snrs)//2, -1]:
                ax2.plot(thetas, phasor_pdfs[idx], label='s = {:.1f}'.format(_snrs[idx]))
            ax2.set_xlabel('$\\theta$')
            ax2.set_ylabel('$p(\\theta)$')
            ax2.legend()

        # Gauss Phasor
        phasor_pdfs = np.array([phase_distribution_gauss(thetas, sigma, s) for s in _snrs ])
        phasor_sigma = np.sqrt(np.array([variance(thetas, pdf) for pdf in phasor_pdfs]))

        ax.plot(_snrs, phase2time(phasor_sigma), label='Gauss')
        ax.legend()

    # Set limit on x values
    if not True or ax.get_xscale() != 'log':
        ax.set_xlim(0, 100)

    if args.fig_dir:
        fig.tight_layout()
        fig.savefig(path.join(args.fig_dir, "time_res_vs_snr.pdf"))

    if args.show_plots:
        plt.show()


