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

"""
Find the best integer multiple to shift
antennas to be able to resolve
"""

import h5py
from itertools import combinations, zip_longest, product
import matplotlib.pyplot as plt
import numpy as np
from os import path
from scipy.interpolate import interp1d

from earsim import REvent
from atmocal import AtmoCal
import aa_generate_beacon as beacon
import lib
from lib import rit

def find_best_sample_shifts_summing_at_location(test_loc, antennas, allowed_sample_shifts, dt=None, sample_shift_first_trace=0, plot_iteration_with_shifted_trace=None, fig_dir=None, fig_distinguish=None):
    """
    Find the best sample_shift for each antenna by summing the antenna traces
    and seeing how to get the best alignment.
    """
    a_ = []
    t_ = []
    t_min = 1e9
    t_max = -1e9

    if dt is None:
        dt = antennas[0].t_AxB[1] - antennas[0].t_AxB[0]

    # propagate to test location
    for i, ant in enumerate(antennas):
        aloc = [ant.x, ant.y, ant.z]
        delta, dist = atm.light_travel_time(test_loc, aloc)
        delta = delta*1e9

        t__ = np.subtract(ant.t_AxB, delta)
        t_.append(t__)
        a_.append(ant.E_AxB)

        if t__[0] < t_min:
            t_min = t__[0]
        if t__[-1] > t_max:
            t_max = t__[-1]

    # Interpolate and find best sample shift
    max_neg_shift = 0 #np.min(allowed_sample_shifts) * dt
    max_pos_shift = 0 #np.max(allowed_sample_shifts) * dt

    t_sum = np.arange(t_min+max_neg_shift, t_max+max_pos_shift, dt)
    a_sum = np.zeros(len(t_sum))

    best_sample_shifts = np.zeros( (len(antennas)) ,dtype=int)
    for i, (t_r, E_) in enumerate(zip(t_, a_)):
        f = interp1d(t_r, E_, assume_sorted=True, bounds_error=False, fill_value=0)
        a_int = f(t_sum)

        if i == 0:
            a_sum += a_int
            best_sample_shifts[i] = sample_shift_first_trace
            continue

        # init figure
        if i == plot_iteration_with_shifted_trace:
            fig, ax = plt.subplots()
            ax.set_title("Traces at ({:.1f},{:.1f},{:.1f})".format(*test_loc))
            ax.set_xlabel("Time [ns]")
            ax.set_ylabel("Amplitude")
            ax.plot(t_sum, a_sum)

        shift_maxima = np.zeros( len(allowed_sample_shifts) )
        for j, shift in enumerate(allowed_sample_shifts):
            augmented_a = np.roll(a_int, shift)

            shift_maxima[j] = np.max(augmented_a + a_sum)

            if i == plot_iteration_with_shifted_trace:
                ax.plot(t_sum, augmented_a, label=f'{shift} shifted')

        # transform maximum into best_sample_shift
        best_idx = np.argmax(shift_maxima)
        best_sample_shifts[i] = allowed_sample_shifts[best_idx]
        a_sum += np.roll(a_int, best_sample_shifts[i])

        # cleanup figure
        if i == plot_iteration_with_shifted_trace:
            if fig_dir: # note this is a global variable here
                fig.savefig(path.join(fig_dir, __file__ + '.loc{:.1f}-{:.1f}-{:.1f}'.format(*test_loc) + f'.i{i}{fig_distinguish}.pdf'))
            plt.close(fig)

    # Return ks
    return best_sample_shifts, np.max(a_sum)

if __name__ == "__main__":
    import sys
    import os
    import matplotlib
    if os.name == 'posix' and "DISPLAY" not in os.environ:
        matplotlib.use('Agg')

    atm = AtmoCal()

    fname = "ZH_airshower/mysim.sry"

    fig_dir = "./figures/periods_from_shower_figures/"
    fig_subdir = path.join(fig_dir, 'shifts/')

    allowed_ks = np.arange(-2, 3, 1)
    Xref = 400

    x_coarse = np.linspace(-20e3, 20e3, 10)
    y_coarse = np.linspace(-20e3, 20e3, 10)

    x_fine = np.linspace(-2e3, 2e3, 30)
    y_fine = np.linspace(-2e3, 2e3, 30)

    N_runs = 3

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

    # Read in antennas from file
    _, __, antennas = beacon.read_beacon_hdf5(antennas_fname)
    # Read original REvent
    ev = REvent(fname)
    # .. patch in our antennas
    ev.antennas = antennas

    # For now only implement using one freq_name
    freq_names = antennas[0].beacon_info.keys()
    if len(freq_names) > 1:
        raise NotImplementedError

    freq_name = next(iter(freq_names))
    f_beacon = ev.antennas[0].beacon_info[freq_name]['freq']

    # determine allowable ks per location
    dt = ev.antennas[0].t[1] - ev.antennas[0].t[0]
    allowed_sample_shifts = np.rint(allowed_ks/f_beacon /dt).astype(int)
    print("Checking:", allowed_ks, ": shifts :", allowed_sample_shifts)


    # Remove time due to true phase
    for i, ant in enumerate(ev.antennas):
        true_phase = ant.beacon_info[freq_name]['phase']
        true_phase_time = true_phase/(2*np.pi*f_beacon)

        ev.antennas[i].t -= true_phase_time

    # Prepare polarisation and passbands
    rit.set_pol_and_bp(ev, low=0.03, high=0.08)

    dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,0)
    scale2d = dXref*np.tan(np.deg2rad(2.))

    # Setup Plane grid to test
    for r in range(N_runs):
        xoff, yoff = 0,0
        if r == 0:
            x = x_coarse
            y = y_coarse
        else:
            old_ks_per_loc = ks_per_loc[best_idx]
            xoff = xx[best_idx]
            yoff = yy[best_idx]
            if r == 1:
                x = x_fine
                y = y_fine
            else:
                x /= 4
                y /= 4

        print(f"Testing grid run {r} centered on {xoff}, {yoff}")

        ks_per_loc = np.zeros( (len(x)*len(y), len(ev.antennas)) , dtype=int)
        maxima_per_loc = np.zeros( (len(x)*len(y)))

        ## Check each location on grid
        xx = []
        yy = []
        N_loc = len(maxima_per_loc)

        for i, (x_, y_) in enumerate(product(x,y)):
            tmp_fig_subdir = None
            if i % 10 ==0:
                print(f"Testing location {i} out of {N_loc}")
                tmp_fig_subdir = fig_subdir

            test_loc = (x_+xoff)* ev.uAxB + (y_+yoff)*ev.uAxAxB + dXref *ev.uA
            xx.append(x_+xoff)
            yy.append(y_+yoff)

            # Find best k for each antenna
            shifts, maximum = find_best_sample_shifts_summing_at_location(test_loc, ev.antennas, allowed_sample_shifts, dt=dt, fig_dir=tmp_fig_subdir, plot_iteration_with_shifted_trace=len(ev.antennas)-1, fig_distinguish=f".run{r}")

            # Translate sample shifts back into period multiple k
            ks = np.rint(shifts*f_beacon*dt)

            ks_per_loc[i] = ks
            maxima_per_loc[i] = maximum

        xx = np.array(xx)
        yy = np.array(yy)
        locs = list(zip(xx, yy))

        ## Save maxima to file
        np.savetxt(fig_dir + __file__+f'.maxima.run{r}.txt', np.column_stack((locs, maxima_per_loc, ks_per_loc)) )

        if True: #plot maximum at test locations
            fig, axs = plt.subplots()
            axs.set_title(f"Grid Run {r}")
            axs.set_ylabel("vxvxB [km]")
            axs.set_xlabel("  vxB [km]")
            sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, cmap='Spectral_r', alpha=0.6)
            fig.colorbar(sc, ax=axs)

            if fig_dir:
                fig.savefig(path.join(fig_dir, __file__+f'.maxima.run{r}.pdf'))

        ## Save ks to file
        best_idx = np.argmax(maxima_per_loc)
        np.savetxt(fig_dir + __file__+f'.bestk.run{r}.txt', ks_per_loc[best_idx] )
        print('Best k:', ks_per_loc[best_idx])

        # Abort if no improvement
        if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx]).all() ):
            print("No changes from previous grid, breaking")
            break

    # Save best ks to hdf5 antenna file
    with h5py.File(antennas_fname, 'a') as fp:
        group = fp.require_group('antennas')

        for i, ant in enumerate(antennas):
            h5ant = group[ant.name]
            h5ant.attrs['best_k'] = old_ks_per_loc[i]

    plt.show()