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

import h5py
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


    f_beacon_band = (49e-3,55e-3) #GHz
    allow_frequency_fitting = False
    read_frequency_from_file = True

    show_plots = True

    fname = "ZH_airshower/mysim.sry"

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

    if not path.isfile(antennas_fname):
        print("Antenna file cannot be found, did you try generating a beacon?")
        sys.exit(1)

    # read in antennas
    with h5py.File(antennas_fname, 'a') as fp:
        if 'antennas' not in fp.keys():
            print("Antenna file corrupted? no antennas")
            sys.exit(1)

        group = fp['antennas']

        f_beacon = None
        if read_frequency_from_file and 'tx' in fp:
            tx = fp['tx']
            if 'f_beacon' in tx.attrs:
                f_beacon = tx.attrs['f_beacon']
            else:
                print("No frequency found in file.")
                sys.exit(2)
            f_beacon_estimate_band = 0.01*f_beacon

        elif allow_frequency_fitting:
            f_beacon_estimate_band = (f_beacon_band[1] - f_beacon_band[0])/2
            f_beacon = f_beacon_band[1] - f_beacon_estimate_band
        else:
            print("Not allowed to fit frequency and no tx group found in file.")
            sys.exit(2)

        N_antennas = len(group.keys())
        # just for funzies
        found_data = np.zeros((N_antennas, 3))

        # Determine frequency and phase
        for i, name in enumerate(group.keys()):
            ant_group = group[name]

            if 'traces' not in ant_group.keys():
                print(f"Antenna file corrupted? no 'traces' in {name}")
                sys.exit(1)

            traces = ant_group['traces']
            t_trace = traces[0]

            if not True:
                # only take the Beacon trace
                test_traces = traces[4]
                orients = ['B']
            else:
                test_traces = traces[1:]
                orients = ['Ex', 'Ey', 'Ez', 'B']

            if True:
                freqs, phases, amps = lib.find_beacon_in_traces(
                        test_traces, t_trace,
                        f_beacon_estimate=f_beacon,
                        frequency_fit=allow_frequency_fitting,
                        f_beacon_estimate_band=f_beacon_estimate_band
                        )
            else:
                # Testing
                freqs = [f_beacon]
                t0 = ant_group.attrs['t0']

                phases = [ 2*np.pi*t0*f_beacon ]
                amps = [ 3e-7 ]

            # choose highest amp
            #idx = np.argmax(amps, axis=1)
            idx = 0

            frequency = freqs[idx]
            phase = phases[idx]
            amplitude = amps[idx]
            orientation = orients[idx]

            if show_plots and (i == 60 or i == 72):
                fig, ax = plt.subplots()
                trace_amp = max(traces[-1]) - min(traces[-1])

                myt = np.linspace(min(traces[0]), max(traces[0]), 10*len(traces[0]))
                ax.plot(t_trace, traces[-1], marker='.', label='trace')
                ax.plot(myt, lib.sine_beacon(frequency, myt, amplitude=amplitude, phase=phase), ls='dashed', label='simulated')
                ax.set_title(f"Beacon at antenna {ant_group}\nF:{frequency}, P:{phase}, A:{amplitude}")
                ax.legend()

            ant_group.attrs['beacon_freq'] = frequency
            ant_group.attrs['beacon_phase'] = phase
            ant_group.attrs['beacon_amplitude'] = amplitude
            ant_group.attrs['beacon_orientation'] = orientation

            found_data[i] = frequency, phase, amplitude

    print("Beacon Phases, Amplitudes and Frequencies written to", antennas_fname)

    # show histogram of found frequencies
    if show_plots:
        if True or allow_frequency_fitting:
            fig, ax = plt.subplots()
            ax.set_xlabel("Frequency")
            ax.set_ylabel("Counts")
            ax.axvline(f_beacon, ls='dashed', color='g')
            ax.hist(found_data[:,0], bins='sqrt', density=False)

        if True:
            fig, ax = plt.subplots()
            ax.set_xlabel("Amplitudes")
            ax.set_ylabel("Counts")
            ax.hist(found_data[:,2], bins='sqrt', density=False)

        plt.show()