#!/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


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

    import matplotlib
    matplotlib.use('Agg')

    fname = "ZH_airshower/mysim.sry"
    c_light = 3e8*1e-9
    show_plots = not True
    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 not True else antennas_fname

    fig_dir = "./figures" # 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 ref_ant_id is None:
        print("Doing all baselines")
        baselines = list(combinations(antennas,2))
    # use ref_ant
    else:
        ref_ant = antennas[ref_ant_id]
        print(f"Doing all baselines with {ref_ant.name}")
        baselines = list(zip_longest([], antennas, fillvalue=ref_ant))

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

    # Get phase difference per baseline
    phase_diffs = np.empty( (len(baselines), 2) )
    for i, base in enumerate(baselines):
        if i%100==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:

            true_phases = np.array([ant.beacon_info[freq_name]['phase'] for ant in base])
            true_phases_diff = lib.phase_mod(lib.phase_mod(true_phases[1]) - lib.phase_mod(true_phases[0]))
        except IndexError:
            # true_phase not determined yet
            print(f"Missing true_phases for {freq_name} in baseline {base[0].name},{base[1].name}")
            true_phases_diff = np.nan

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

    # Read actual phases from antenna hdf5
    actual_antenna_phases = { a.name: 2*np.pi*a.attrs['clock_offset']*f_beacon for a in antennas }


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

        if True: # remove time axis phase difference
            t0s = np.array([ant.t_AxB[0] for ant in base])
            delta_t0 = t0s[1] - t0s[1]
            actual_phase_diff = lib.phase_mod( actual_phase_diff - lib.phase_mod(delta_t0*2*np.pi*f_beacon))
        my_phase_diffs.append(actual_phase_diff)

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

        phase_residuals = lib.phase_mod(phase_diffs[:,1] - my_phase_diffs)

        fig, axs = plt.subplots(2, 1, sharex=True)
        axs[0].set_title("Measured phase difference - Actual phase difference")
        axs[0].set_xlabel("Phase residuals $\\varphi$")

        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('$\\varphi/(2\\pi f_{beacon})$ [ns]')

        i=0
        axs[i].set_ylabel("#")
        axs[i].hist(phase_residuals, bins='sqrt', density=False)

        i=1
        axs[i].set_ylabel("Baseline combination #")
        if not True:
            axs[i].plot(my_phase_diffs, np.arange(N_base), ls='none', marker='+', label='actual time shifts')
            l = axs[i].plot(phase_diffs[:,1], np.arange(N_base), ls='none', marker='x', label='calculated')

            axs[i].legend()

        else:
            axs[i].plot(phase_residuals, np.arange(N_base), ls='none', marker='x')

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

    if show_plots:
        plt.show()