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

import h5py
from itertools import combinations, zip_longest
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
import matplotlib as mpl
import numpy as np

import aa_generate_beacon as beacon
import lib


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()

    fname = "ZH_airshower/mysim.sry"

    show_plots = args.show_plots
    ref_ant_id = None # leave None for all baselines

    ####
    fname_dir = path.dirname(fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)
    time_diffs_fname = 'time_diffs.hdf5' if False else antennas_fname
    fig_dir = "./figures" # set None to disable saving

    basenames, time_diffs, f_beacons, true_phase_diffs, k_periods = beacon.read_baseline_time_diffs_hdf5(time_diffs_fname)

    f_beacon, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)

    # TODO: allow multiple frequencies
    if (f_beacon != f_beacons).any():
        raise NotImplementedError

    N_base = len(basenames)
    N_ant = len(antennas)

    # reshape time_diffs into N_ant x N_ant matrix
    sigma_phase_matrix = np.full( (N_ant, N_ant), np.nan, dtype=float)

    ## set i=i terms to 0
    for i in range(N_ant):
        sigma_phase_matrix[i,i] = 0

    ## fill matrix
    name2idx = lambda name: int(name)-1
    for i, b in enumerate(basenames):
        idx = (name2idx(b[0]), name2idx(b[1]))

        sigma_phase_matrix[(idx[0], idx[1])] = lib.phase_mod(true_phase_diffs[i])
        sigma_phase_matrix[(idx[1], idx[0])] = lib.phase_mod(-1*true_phase_diffs[i])

    mat_kwargs = dict(
        norm = Normalize(vmin=-np.pi, vmax=+np.pi),
        cmap = mpl.cm.get_cmap('Spectral_r')
    )

    # Show Matrix as figure
    if True:
        fig, ax = plt.subplots()
        ax.set_title("Measured phase differences Baseline i,j")
        ax.set_ylabel("Antenna no. i")
        ax.set_xlabel("Antenna no. j")

        im = ax.imshow(sigma_phase_matrix, interpolation='none', **mat_kwargs)
        fig.colorbar(im, ax=ax)

        if fig_dir:
            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".matrix.original.pdf"))

        plt.close(fig)


    # Modify the matrix to let each column represent multiple
    # measurements of the same baseline (j,0) phase difference
    if True:
        # for each row j subtract the 0,j element from the whole row
        # and apply phase_mod
        first_row = -1*(sigma_phase_matrix[0,:] * np.ones_like(sigma_phase_matrix)).T

        # Show subtraction Matrix as figure
        if True:
            fig, ax = plt.subplots()
            ax.set_title("Subtraction matrix i,j")
            ax.set_ylabel("Antenna no. i")
            ax.set_xlabel("Antenna no. j")

            im = ax.imshow(first_row, interpolation='none', **mat_kwargs)
            fig.colorbar(im, ax=ax)

            if fig_dir:
                fig.savefig(path.join(fig_dir, path.basename(__file__) + f".matrix.first_row.pdf"))

            plt.close(fig)

        sigma_phase_matrix = sigma_phase_matrix - first_row
        sigma_phase_matrix = lib.phase_mod(sigma_phase_matrix)

    # Except for the first row, these are all separate measurements
    # Condense into phase offset per antenna
    if True: # do not use the first row
        my_mask = np.arange(1, len(sigma_phase_matrix), dtype=int)
    else:
        my_mask = np.arange(0, len(sigma_phase_matrix), dtype=int)

    mean_sigma_phase = np.nanmean(sigma_phase_matrix[my_mask], axis=0)
    std_sigma_phase  = np.nanstd( sigma_phase_matrix[my_mask], axis=0)


    # Show resulting matrix as figure
    if True:
        fig, axs = plt.subplots(2,1, sharex=True)
        axs[0].set_title("Modified measured phase differences Baseline 0,j")
        axs[0].set_ylabel("Antenna no. 0")
        axs[-1].set_xlabel("Antenna no. j")

        im = axs[0].imshow(sigma_phase_matrix, interpolation='none', **mat_kwargs)
        fig.colorbar(im, ax=axs)
        axs[0].set_aspect('auto')

        colours = [mat_kwargs['cmap'](mat_kwargs['norm'](x)) for x in mean_sigma_phase]
        axs[1].set_ylabel("Mean Baseline Phase (0, j)[rad]")
        axs[1].errorbar(np.arange(N_ant), mean_sigma_phase, yerr=std_sigma_phase, ls='none')
        axs[1].scatter(np.arange(N_ant), mean_sigma_phase, c=colours,s=4)

        if fig_dir:
            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".matrix.modified.pdf"))

        plt.close(fig)

    # write into antenna hdf5
    with h5py.File(antennas_fname, 'a') as fp:
        group = fp['antennas']

        freq_name = None
        for i, ant in enumerate(antennas):
            h5ant = group[ant.name]

            h5beacon_info = h5ant['beacon_info']

            # find out freq_name
            if freq_name is None:
                freq_name = [ k for k in h5beacon_info.keys() if np.isclose(h5beacon_info[k].attrs['freq'], f_beacon)][0]

            h5attrs = h5beacon_info[freq_name].attrs

            idx = name2idx(ant.name)
            h5attrs['sigma_phase_mean'] = mean_sigma_phase[idx]
            h5attrs['sigma_phase_std']  =  std_sigma_phase[idx]


    ##############################
    # Compare actual time shifts #
    ##############################
    actual_antenna_time_shifts = { a.name: a.attrs['clock_offset'] for a in sorted(antennas, key=lambda a: int(a.name)) }

    if True:
        actual_antenna_phase_shifts = [ -1*lib.phase_mod(2*np.pi*f_beacon*v) for k,v in actual_antenna_time_shifts.items() ]
        antenna_names = [int(k)-1 for k,v in actual_antenna_time_shifts.items() ]

        # Make sure to shift all antennas by a global phase
        global_phase_shift = actual_antenna_phase_shifts[0] - mean_sigma_phase[0]
        actual_antenna_phase_shifts = lib.phase_mod(actual_antenna_phase_shifts - global_phase_shift )

        for i in range(2):
            plot_residuals = i == 1
            colors = ['blue', 'orange']

            fig, axs = plt.subplots(2, 1, sharex=True)

            if True:
                forward = lambda x: x/(2*np.pi*f_beacon)
                inverse = lambda x: 2*np.pi*x*f_beacon
                secax = axs[0].secondary_xaxis('top', functions=(forward, inverse))
                secax.set_xlabel('Time $\\Delta\\varphi/(2\\pi f_{beac})$ [ns]')

            if plot_residuals:
                phase_residuals = lib.phase_mod(mean_sigma_phase - actual_antenna_phase_shifts)
                fig.suptitle("Difference between Measured and Actual phases (minus global phase)\n for Antenna $i$")
                axs[-1].set_xlabel("Antenna Phase Residual $\\Delta_\\varphi$")
            else:
                fig.suptitle("Comparison Measured and Actual phases (minus global phase)\n for Antenna $i$")
                axs[-1].set_xlabel("Antenna Phase $\\varphi$")


            i=0
            axs[i].set_ylabel("#")
            if plot_residuals:
                axs[i].hist(phase_residuals, bins='sqrt', alpha=0.8, color=colors[0])
            else:
                axs[i].hist(mean_sigma_phase,    bins='sqrt', density=False, alpha=0.8, color=colors[0], ls='solid' , histtype='step', label='Measured')
                axs[i].hist(actual_antenna_phase_shifts, bins='sqrt', density=False, alpha=0.8, color=colors[1], ls='dashed', histtype='step', label='Actual')


            i=1
            axs[i].set_ylabel("Antenna no.")
            if plot_residuals:
                axs[i].plot(phase_residuals, np.arange(N_ant), alpha=0.6, ls='none', marker='x', color=colors[0])
            else:
                axs[i].errorbar(mean_sigma_phase, np.arange(N_ant), yerr=std_sigma_phase, marker='4', alpha=0.7, ls='none', color=colors[0], label='Measured')
                axs[i].plot(actual_antenna_phase_shifts,  antenna_names, ls='none', marker='3', alpha=0.8, color=colors[1], label='Actual')

                axs[i].legend()
            fig.tight_layout()

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

    ##########################
    ##########################
    ##########################

    actual_baseline_time_shifts = []
    for i,b in enumerate(basenames):
        actual_baseline_time_shift = actual_antenna_time_shifts[b[0]] - actual_antenna_time_shifts[b[1]]

        actual_baseline_time_shifts.append(actual_baseline_time_shift)

    # unpack mean_sigma_phase back into a list of time diffs
    measured_baseline_time_diffs = []
    for i,b in enumerate(basenames):
        phase0, phase1 = mean_sigma_phase[name2idx(b[0])], mean_sigma_phase[name2idx(b[1])]
        measured_baseline_time_diffs.append(lib.phase_mod(phase1 - phase0)/(2*np.pi*f_beacon))

    # Make a plot
    if True:
        for i in range(2):
            fig, ax = plt.subplots()
            ax.set_title("Baseline Time difference reconstruction" + ( '' if i == 0 else ' (wrapped time)'))
            ax.set_xlabel("Baseline no.")
            ax.set_ylabel("Time $\\Delta t$ [ns]")
            if True:
                forward = lambda x: x/(2*np.pi*f_beacon)
                inverse = lambda x: 2*np.pi*x*f_beacon
                secax = ax.secondary_yaxis('right', functions=(inverse, forward))
                secax.set_ylabel('Phase $\\Delta \\varphi$ [rad]')

            if True: # indicate single beacon period span
                ax.plot((-1, -1), (-1/(2*f_beacon), 1/(2*f_beacon)), marker='3', ms=10, label='1/f_beacon')
            if i == 0:
                ax.plot(np.arange(N_base), actual_baseline_time_shifts, ls='none', marker='h', alpha=0.8, label='actual time shifts')
            else:
                ax.plot(np.arange(N_base), (actual_baseline_time_shifts+1/(2*f_beacon))%(1/f_beacon) - 1/(2*f_beacon), ls='none', marker='h', label='actual time shifts', alpha=0.8)
            ax.plot(np.arange(N_base), measured_baseline_time_diffs, ls='none', alpha=0.8, marker='x', label='calculated')

            ax.legend()
            if fig_dir:
                extra_name = ''
                if i == 1:
                    extra_name = '.wrapped'

                old_lims = (ax.get_xlim(), ax.get_ylim())
                for j in range(2):
                    if j == 1:
                        ax.set_xlim(-5, 50)
                        extra_name += '.n50'
                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f".time_comparison{extra_name}.pdf"))

                ax.set_xlim(*old_lims[0])
                ax.set_ylim(*old_lims[1])

    if show_plots:
        plt.show()