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

"""
Do a reconstruction of airshower after correcting for the
clock offsets.
"""

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 pickle
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__":
    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()
    args = parser.parse_args()

    fname = "ZH_airshower/mysim.sry"

    fig_dir = args.fig_dir
    fig_subdir = path.join(fig_dir, 'reconstruction')
    show_plots = args.show_plots

    remove_beacon_from_traces = 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)
        if fig_subdir:
            os.makedirs(fig_subdir, 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)
    # .. 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']

    # Repair clock offsets with the measured offsets
    measured_repair_offsets = beacon.read_antenna_clock_repair_offsets(ev.antennas, mode='phases', freq_name=freq_name)
    for i, ant in enumerate(ev.antennas):
        # t_AxB will be set by the rit.set_pol_and_bp function
        ev.antennas[i].t += measured_repair_offsets[i]

        # .. and remove the beacon from the traces
        if remove_beacon_from_traces:
            clock_phase = measured_repair_offsets[i]*2*np.pi*f_beacon

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

            tx_amps = txdata['amplitudes']
            tx_amps_sum = np.sum(tx_amps)

            # 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

    N_X, Xlow, Xhigh = 23, 100, 1200
    with joblib.parallel_backend("loky"):
        res = rit.reconstruction(ev, outfile=fig_subdir+'/fig.pdf', slice_outdir=fig_subdir+'/', Xlow=Xlow, N_X=N_X, Xhigh=Xhigh)

    ## Save a pickle
    with open(pickle_fname, 'wb') as fp:
        pickle.dump(res,fp)

    if show_plots:
        plt.show()