#!/usr/bin/env python3
# vim: fdm=indent ts=4

import h5py
from itertools import combinations, product
import matplotlib.pyplot as plt
import numpy as np

import aa_generate_beacon as beacon
import lib
from lib import figlib


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 scriptlib import MyArgumentParser
    parser = MyArgumentParser()
    parser.add_argument('ref_ant_idx', default=None, nargs='*', type=int, help='Reference Antenna Indices for Baselines(ref_ant_idx, *). Leave empty to use all available antennas. (Default: %(default)s) ')
    args = parser.parse_args()

    figsize = (12,8)
    c_light = 3e8*1e-9
    show_plots = args.show_plots
    ref_ant_id = args.ref_ant_idx # leave None for all baselines

    ####
    fname_dir = args.data_dir
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)
    time_diffs_fname = 'time_diffs.hdf5' if False else antennas_fname
    snr_fname = path.join(fname_dir, beacon.snr_fname)

    fig_dir = args.fig_dir # set None to disable saving

    if not path.isfile(antennas_fname):
        print("Antenna file cannot be found, did you try generating a beacon?")
        sys.exit(1)

    # Read in antennas from file
    f_beacon, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)

    # run over all baselines
    if not ref_ant_id:
        print("Doing all baselines")
        baselines = combinations(antennas,2)

    # use ref_ant
    else:
        ref_ants = [antennas[i] for i in ref_ant_id]
        print("Doing all baselines with {}".format([int(a.name) for a in ref_ants]))
        baselines = product(ref_ants, antennas)

    # For now, only one beacon_frequency is supported
    freq_names = antennas[0].beacon_info.keys()
    if len(freq_names) > 1:
        raise NotImplementedError

    freq_name = next(iter(freq_names))

    # Collect baselines from optional generators
    baselines = list(baselines)

    # Get phase difference per baseline

    phase_diffs = np.empty( (len(baselines), 2) )
    for i, base in enumerate(baselines):
        if i%1000==0:
            print(i, "out of", len(baselines))

        # read f_beacon from the first antenna
        f_beacon = base[0].beacon_info[freq_name]['freq']

        # Get true phase diffs
        try:
            clock_phases = np.array([ant.beacon_info[freq_name]['clock_phase'] for ant in base])
            clock_phases_diff = lib.phase_mod(lib.phase_mod(clock_phases[1]) - lib.phase_mod(clock_phases[0]))
        except IndexError:
            # clock_phase not determined yet
            print(f"Missing clock_phases for {freq_name} in baseline {base[0].name},{base[1].name}")
            clock_phases_diff = np.nan

        # save phase difference with antenna names
        phase_diffs[i] = [f_beacon, clock_phases_diff]

    beacon.write_baseline_time_diffs_hdf5(time_diffs_fname, baselines, phase_diffs[:,1], [0]*len(phase_diffs), phase_diffs[:,0])

    ##############################
    # Compare actual time shifts #
    ##############################
    snrs = beacon.read_snr_file(snr_fname)
    snr_str = f"$\\langle SNR \\rangle$ = {snrs['mean']: .1e}"

    actual_antenna_clock_phases = { a.name: -2*np.pi*a.attrs['clock_offset']*f_beacon for a in sorted(antennas, key=lambda a: int(a.name)) }

    # Compare actual time shifts
    my_phase_diffs = []
    for i,b in enumerate(baselines):
        actual_clock_phase_diff = lib.phase_mod( lib.phase_mod(actual_antenna_clock_phases[b[1].name]) - lib.phase_mod(actual_antenna_clock_phases[b[0].name]))

        this_actual_clock_phase_diff = lib.phase_mod( actual_clock_phase_diff )
        my_phase_diffs.append(this_actual_clock_phase_diff)

    # Make a plot
    if True:
        N_base = len(baselines)
        N_ant = len(antennas)

        for i in range(2):
            plot_residuals = i == 1

            true_phases = my_phase_diffs
            measured_phases = phase_diffs[:,1]

            hist_kwargs = {}
            if plot_residuals:
                measured_phases = lib.phase_mod(measured_phases - true_phases)
                hist_kwargs['histtype'] = 'stepfilled'

            fig = figlib.phase_comparison_figure(
                    measured_phases,
                    true_phases,
                    plot_residuals=plot_residuals,
                    f_beacon=f_beacon,
                    figsize=figsize,
                    hist_kwargs=hist_kwargs,
                    fit_gaussian=plot_residuals,
                    )

            axs = fig.get_axes()

            axs[0].legend(title=snr_str)

            if plot_residuals:
                fig.suptitle("Difference between Measured and Actual phase difference\n for Baselines (i,j" + (')' if not ref_ant_id else '='+str([ int(a.name) for a in ref_ants])+')'))
                axs[-1].set_xlabel("Baseline Phase Residual $\\Delta\\varphi_{ij_{meas}} - \\Delta\\varphi_{ij_{true}}$ [rad]")
            else:
                fig.suptitle("Comparison Measured and Actual phase difference\n for Baselines (i,j"  + (')' if not ref_ant_id else '='+str([ int(a.name) for a in ref_ants])+')'))
                axs[-1].set_xlabel("Baseline Phase $\\Delta\\varphi_{ij}$ [rad]")

            #
            i=0
            secax = axs[i].child_axes[0]
            secax.set_xlabel('Time $\\Delta\\varphi/(2\\pi f_{beac})$ [ns]')

            #
            i=1
            axs[i].set_ylabel("Baseline no.")

            if fig_dir:
                extra_name = "measured"
                if plot_residuals:
                    extra_name = "residuals"
                fig.savefig(path.join(fig_dir, path.basename(__file__) + f".{extra_name}.F{freq_name}.pdf"))

    if show_plots:
        plt.show()