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

    fname = "ZH_airshower/mysim.sry"

    ####
    fname_dir = path.dirname(fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)

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

    # run over all baselines
    if True:
        baselines = list(combinations(antennas,2))
    # use ref_ant
    else:
        ref_ant = antennas[0]
        baselines = list(zip_longest([], antennas, fillvalue=ref_ant))

    freq_names = antennas[0].beacon_info.keys()
    if len(freq_names) > 1:
        raise NotImplementedError

    freq_name = next(iter(freq_names))


    # Determine integer multiple of periods to shift
    integer_periods = np.empty( (len(baselines), 3) )
    for i, base in enumerate(baselines):
        # which traces to keep track of
        traces = [ base[0].E_AxB, base[1].E_AxB ]

        # how many samples do we need to shift
        sampling_dt = (base[1].t[1] - base[1].t[0]) # ns
        ks, maxima = lib.coherence_sum_maxima(traces[0], traces[1])
        max_idx = np.argmax(maxima)
        best_k = ks[max_idx]
        delta_t_coherence = sampling_dt*best_k # ns

        print("K", best_k, sampling_dt, '=', delta_t_coherence)


        # get the amount of periods to move
        f_beacon = base[0].beacon_info[freq_name]['freq']
        k_period, rest = np.divmod(delta_t_coherence, 1/f_beacon)

        # always keep the reference before traces[1]
        if rest < 0:
            k_period -= 1

        # save k_period with antenna names
        integer_periods[i] = [int(base[0].name), int(base[1].name), k_period]

        if i in [ 98, 99 ]:
            print('i',i,'k[T]',k_period, 'rest[ns]',rest, 'T[ns]',1/f_beacon)

            # Show correlation maxima plot
            if True:
                fig, ax = plt.subplots()
                ax.set_title(f"Correlation Maxima {i}")
                ax.set_xlabel("k")
                ax.set_ylabel("Maximum correlation")
                ax.plot(ks, maxima)
                ax.plot(best_k, maxima[max_idx], marker='X')

            # Delta between first timestamp from both antennas
            delta_t_antennas = base[0].t[0] - base[1].t[0]

            # Delta t due to the beacon
            try:
                true_phases = np.array([ant.beacon_info[freq_name]['true_phase'] for ant in base])
                delta_true_phases = lib.phase_mod(true_phases[0] - true_phases[1])
                delta_t_beacon = delta_true_phases/(2*np.pi*f_beacon)
            except e:
                # freq_name not found
                # simply continue and set it them 0
                print("No beacon")
                delta_true_phases = 0
                delta_t_beacon = 0

            print("t0[ns]", delta_t_antennas, "t_beacon[ns]", delta_t_beacon, "phase", delta_true_phases)
            fig, ax = plt.subplots()
            ax.set_xlabel('t')
            ax.plot(base[0].t, traces[0], label=f'Reference {base[0].name}', alpha=0.5)
            # plot vertical lines indicating f_beacon
            min_t, max_t = base[0].t[0], base[0].t[-1]
            N_lines = int( (max_t - min_t)*f_beacon) +1
            for i, t in enumerate(np.arange(N_lines)/f_beacon):
                ax.axvline( min_t + t, color='k', alpha=0.5, label=None if i!=0 else 'P_beacon')

            ax.plot(base[1].t + delta_t_antennas, traces[1], label=f'Original {base[1].name} (t0 removed)', alpha=0.4, marker='+', ms=5)
            ax.plot(base[1].t + delta_t_antennas + k_period/f_beacon + rest, traces[1], label='Coherence', alpha=0.3, marker='x', ms=5)
            ax.plot(base[1].t + delta_t_antennas + k_period/f_beacon + delta_t_beacon, traces[1], label='Beacon only + Periods', alpha=0.6)

            ax.legend()

    # Save integer periods to antennas
    with h5py.File(antennas_fname, 'a') as fp:
        group_name = 'beacon_ks'
        if group_name in fp:
            del fp[group_name]

        fp.create_dataset(group_name, data=integer_periods)

    plt.show()
    # Report back to CLI
    print("Period Multiples resolved in", antennas_fname)