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

import h5py
from itertools import combinations, zip_longest
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()
    args = parser.parse_args()

    figsize = (12,8)
    c_light = lib.c_light
    show_plots = args.show_plots

    remove_absolute_phase_offset_first_antenna = True # takes precedence
    remove_absolute_phase_offset_minimum = True

    ####
    fname_dir = args.data_dir
    antennas_fname = path.join(fname_dir, beacon.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)

    # Make sure at least one beacon has been identified
    if not hasattr(antennas[0], 'beacon_info') or len(antennas[0].beacon_info) == 0:
        print(f"No analysed beacon found for {antennas[0].name}, try running the beacon phase analysis script first.")
        sys.exit(1)

    #
    N_beacon_freqs = len(antennas[0].beacon_info)

    for freq_name in antennas[0].beacon_info.keys():
        beacon_phases = np.empty( (len(antennas)) )
        for i, ant in enumerate(antennas):
            beacon_phases[i] = ant.beacon_info[freq_name]['beacon_phase']

        f_beacon = antennas[0].beacon_info[freq_name]['freq']

        clock_phases = lib.remove_antenna_geometry_phase(tx, antennas, f_beacon, beacon_phases, c_light=c_light)

        # Remove the phase from one antenna
        # this is a free parameter
        # (only required for absolute timing)
        if remove_absolute_phase_offset_first_antenna or remove_absolute_phase_offset_minimum:
            if remove_absolute_phase_offset_first_antenna: # just take the first phase
                minimum_phase = clock_phases[0]
            else: # take the minimum
                minimum_phase = np.min(clock_phases, axis=-1)

            clock_phases -= minimum_phase
            clock_phases = lib.phase_mod(clock_phases)

        # Save to antennas in file
        with h5py.File(antennas_fname, 'a') as fp:
            h5group = fp['antennas']

            for i, ant in enumerate(antennas):
                h5ant = fp['antennas'][ant.name]

                h5beacon_freq = h5ant['beacon_info'][freq_name]
                h5beacon_freq.attrs['clock_phase'] = clock_phases[i]

        # Plot True Phases at their locations
        if show_plots or fig_dir:
            actual_clock_phases = lib.phase_mod(np.array([ -2*np.pi*a.attrs['clock_offset']*f_beacon for a in antennas ]))
            for i in range(2):
                plot_residuals = i == 1
                spatial_unit='m'

                antenna_locs = list(zip(*[(ant.x, ant.y) for ant in antennas]))

                scatter_kwargs = {}
                scatter_kwargs['cmap'] = 'inferno'

                # Measurements
                if not plot_residuals:
                    title='Clock phases'
                    color_label='$\\varphi(\\sigma_t)$ [rad]'

                    fname_extra='measured.'

                    #scatter_kwargs['vmin'] = -np.pi
                    #scatter_kwargs['vmax'] = +np.pi

                # Plot Clock Phases - True Clock Phases at their location
                else:
                    title='Clock phase Residuals'
                    color_label='$\\Delta\\varphi(\\sigma_t) = \\varphi_{meas} - \\varphi_{true}$ [rad]'

                    fname_extra='residuals.'

                    # Modify actual_clock_phases, the same way as clock_phases
                    # was modified
                    if remove_absolute_phase_offset_first_antenna or remove_absolute_phase_offset_minimum:
                        if remove_absolute_phase_offset_first_antenna: # just take the first phase
                            minimum_phase = actual_clock_phases[0]
                        else: # take the minimum
                            minimum_phase = np.min(actual_clock_phases, axis=-1)

                        actual_clock_phases -= minimum_phase
                        actual_clock_phases = lib.phase_mod(actual_clock_phases)

                    clock_phase_residuals = lib.phase_mod(clock_phases - actual_clock_phases)

                if not plot_residuals:
                    loc_c = clock_phases
                else:
                    loc_c = clock_phase_residuals

                ##
                ## Geometrical Plot
                ##
                fig, ax = plt.subplots(figsize=figsize)

                ax.set_xlabel('x' if spatial_unit is None else 'x [{}]'.format(spatial_unit))
                ax.set_ylabel('y' if spatial_unit is None else 'y [{}]'.format(spatial_unit))

                fig.suptitle(title+'\nf_beacon= {:2.0f}MHz'.format(f_beacon*1e3))

                sc = ax.scatter(*antenna_locs, c=loc_c, **scatter_kwargs)
                fig.colorbar(sc, ax=ax, label=color_label)

                if False:
                    for i, ant in enumerate(antennas):
                        ax.text(ant.x, ant.y, ant.name)

                if not True:
                    ax.plot(tx.x, tx.y, 'X', color='k', markeredgecolor='white')

                if fig_dir:
                    fig.tight_layout()
                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f".geom.{fname_extra}F{freq_name}.pdf"))

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

                fig = figlib.phase_comparison_figure(
                        loc_c,
                        actual_clock_phases,
                        plot_residuals=plot_residuals,
                        f_beacon=f_beacon,
                        figsize=figsize,
                        fit_gaussian=plot_residuals,
                        )

                axs = fig.get_axes()
                if plot_residuals:
                    axs[0].set_title("Difference between Measured and True Clock phases")
                else:
                    axs[0].set_title("Comparison Measured and True Clock Phases")

                axs[-1].set_xlabel(f'Antenna {title} {color_label}')

                #
                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("Antenna no.")


                #
                fig.legend(title=snr_str)

                # Save figure
                if fig_dir:
                    fig.tight_layout()
                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f".{fname_extra}F{freq_name}.pdf"))

    print(f"True phases written to", antennas_fname)

    if show_plots:
        plt.show()