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

"""
Show how the Power changes when incorporating the
various clock offsets by plotting on a grid.
"""

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # required for projection='3d' on old matplotliblib versions
import numpy as np
from os import path
import joblib

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


if __name__ == "__main__":
    valid_cases = ['no_offset', 'repair_none', 'repair_phases', 'repair_all']

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

    atm = AtmoCal()

    from scriptlib import MyArgumentParser
    parser = MyArgumentParser()

    group = parser.add_argument_group('figures')
    for case in valid_cases:
        group.add_argument('--'+case.replace('_','-'), dest='figures', action='append_const', const=case)

    args = parser.parse_args()

    wanted_cases = args.figures

    if not wanted_cases or 'all' in wanted_cases:
        wanted_cases = valid_cases

    fname = "ZH_airshower/mysim.sry"

    fig_dir = args.fig_dir
    show_plots = args.show_plots

    remove_beacon_from_traces = True
    apply_signal_window_from_max = True

    ####
    fname_dir = path.dirname(fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)
    pickle_fname = path.join(fname_dir, 'res.pkl')
    tx_fname = path.join(fname_dir, beacon.tx_fname)

    # create fig_dir
    if fig_dir:
        os.makedirs(fig_dir, exist_ok=True)

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

    ##
    ## Setup grid
    ##
    X = 400
    zgr = 0 #not exact
    dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),750,zgr+ev.core[2])
    scale2d = dXref*np.tan(np.deg2rad(2.))
    scale4d = dXref*np.tan(np.deg2rad(4.))
    scale02d = dXref*np.tan(np.deg2rad(0.2))

    Nx, Ny = 21, 21
    scales = {
            'scale2d': scale2d,
            'scale4d': scale4d,
            'scale02d': scale02d,
            }

    plot_titling = {
            'no_offset': "no clock offset",
            'repair_none': "unrepaired clock offset",
            'repair_phases': "phase resolved clock offsets repaired",
            'repair_all': "final measured clock offsets repaired"
            }

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

    freq_name = next(iter(freq_names))

    # Pre remove the beacon from the traces
    #
    # We need to remove it here, so we do not shoot ourselves in
    # the foot when changing to the various clock offsets.
    #
    # Note that the bandpass filter is applied only after E_AxB is
    # reconstructed so we have to manipulate the original traces.
    if remove_beacon_from_traces:
        tx_amps = txdata['amplitudes']
        tx_amps_sum = np.sum(tx_amps)

        for i, ant in enumerate(ev.antennas):
            beacon_phase = ant.beacon_info[freq_name]['beacon_phase']
            f = ant.beacon_info[freq_name]['freq']
            ampl_AxB = ant.beacon_info[freq_name]['amplitude']
            calc_beacon = lib.sine_beacon(f, ev.antennas[i].t, amplitude=ampl_AxB, phase=beacon_phase)

            # Split up contribution to the various polarisations
            for j, amp in enumerate(tx_amps):
                if j == 0:
                    ev.antennas[i].Ex -= amp*(1/tx_amps_sum)*calc_beacon
                elif j == 1:
                    ev.antennas[i].Ey -= amp*(1/tx_amps_sum)*calc_beacon
                elif j == 2:
                    ev.antennas[i].Ez -= amp*(1/tx_amps_sum)*calc_beacon

            # 
            ev.antennas[i].E_AxB -= calc_beacon

    # Slice the traces to a small part around the peak
    if apply_signal_window_from_max:
        N_pre, N_post = 250, 250

        for i, ant in enumerate(ev.antennas):
            max_idx = np.argmax(ant.E_AxB)

            low_idx = max(0, max_idx-N_pre)
            high_idx = min(len(ant.t), max_idx+N_post)

            ev.antennas[i].t = ant.t[low_idx:high_idx]
            ev.antennas[i].t_AxB = ant.t_AxB[low_idx:high_idx]

            ev.antennas[i].Ex = ant.Ex[low_idx:high_idx]
            ev.antennas[i].Ey = ant.Ey[low_idx:high_idx]
            ev.antennas[i].Ez = ant.Ez[low_idx:high_idx]
            ev.antennas[i].E_AxB = ant.E_AxB[low_idx:high_idx]

    # backup antenna times
    backup_antenna_t = [ ant.t for ant in ev.antennas ]
    backup_antenna_t_AxB = [ ant.t_AxB for ant in ev.antennas ]

    ## Apply polarisation and bandpass filter
    rit.set_pol_and_bp(ev)

    with joblib.parallel_backend("loky"):
        for case in wanted_cases:
            print(f"Starting {case} figure")
            # Repair clock offsets with the measured offsets
            transl_modes = {'no_offset':'orig', 'repair_phases':'phases', 'repair_all':'all'}

            if case in transl_modes:
                transl_mode = transl_modes[case]
                measured_offsets = beacon.read_antenna_clock_repair_offsets(antennas, mode=transl_mode, freq_name=freq_name)
            else:
                measured_offsets = [0]*len(ev.antennas)

            for i, ant in enumerate(ev.antennas):
                total_clock_offset = measured_offsets[i]

                ev.antennas[i].t =     backup_antenna_t[i] + total_clock_offset
                ev.antennas[i].t_AxB = backup_antenna_t_AxB[i] + total_clock_offset

                if i == 0:
                    # Specifically compare the times
                    print(bak_ants[i].t[0], ev.antennas[i].t[0], ev.antennas[i].t[0], ev.antennas[i].attrs['clock_offset'], measured_offsets[i])

            #
            # Plot overlapping traces at 0,0,0
            #
            if True:
                P, t_, a_, a_sum, t_sum = rit.pow_and_time([0,0,0], ev, dt=1)

                fig, axs = plt.subplots()
                axs.set_title("Antenna traces" + "\n" + plot_titling[case])
                axs.set_xlabel("Time [ns]")
                axs.set_ylabel("Amplitude [$\\mu V/m$]")

                a_max = [ np.amax(ant.E_AxB) for ant in ev.antennas ]
                power_sort_idx = np.argsort(a_max)

                N_plot = 30
                for i, idx in enumerate(power_sort_idx):
                    if i > N_plot:
                        break
                    
                    alpha = max(0.7, 1/len(a_))

                    axs.plot(t_[idx], a_[idx], color='r', alpha=alpha)

                if fig_dir:
                    fig.tight_layout()
                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.trace_overlap.{case}.pdf'))
                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.trace_overlap.{case}.png'), transparent=True)

            # Measure power on grid
            for scalename, scale in scales.items():
                wx, wy = scale, scale
                print(f"Starting grid measurement for figure {case} with {scalename}")
                xx, yy, p, maxp = rit.shower_plane_slice(ev, X=X, Nx=Nx, Ny=Nx, wx=wx, wy=wy)

                fig, axs = rit.slice_figure(ev, X, xx, yy, p, mode='sp')

                suptitle = fig._suptitle.get_text()
                fig.suptitle("")
                axs.set_title(suptitle +"\n" +plot_titling[case])
                #axs.set_aspect('equal', 'datalim')

                if fig_dir:
                    fig.tight_layout()
                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.X{X}.{case}.{scalename}.pdf'))