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

import numpy as np
import h5py

import lib

import aa_generate_beacon as beacon

from lib import direct_fourier_transform

from numpy.polynomial import Polynomial

def find_beacon_in_traces(
        traces,
        t_trace,
        f_beacon_estimate = 50e6,
        frequency_fit = False,
        N_test_freqs = 5e2,
        f_beacon_estimate_band = 0.01,
        amp_cut = 0.8
        ):
    """
    f_beacon_band is inclusive

    traces is [trace, trace, trace, ..  ]
    """

    amplitudes = np.zeros(len(traces))
    phases = np.zeros(len(traces))
    frequencies = np.zeros(len(traces))

    if frequency_fit: # fit frequency
        test_freqs = f_beacon_estimate + f_beacon_estimate_band * np.linspace(-1, 1, int(N_test_freqs)+1)
        ft_amp_gen = direct_fourier_transform(test_freqs, t_trace, (x for x in traces))

        n_samples = len(t_trace)

        for i, ft_amp in enumerate(ft_amp_gen):
            real, imag = ft_amp
            amps = 1/n_samples * ( real**2 + imag**2)**0.5

            # find frequency peak and surrounding
            # bins valid for parabola fitting
            max_amp_idx = np.argmax(amps)
            max_amp = amps[max_amp_idx]

            if True:
                frequencies[i] = test_freqs[max_amp_idx]
                continue

            valid_mask = amps >= amp_cut*max_amp

            if True: # make sure not to use other peaks
                lower_mask = valid_mask[0:max_amp_idx]
                upper_mask = valid_mask[max_amp_idx:]

                if any(lower_mask):
                    lower_end = np.argmin(lower_mask[::-1])
                else:
                    lower_end = max_amp_idx

                if any(upper_mask):
                    upper_end = np.argmin(upper_mask)
                else:
                    upper_end = 0


                valid_mask[0:(max_amp_idx - lower_end)] = False
                valid_mask[(max_amp_idx + upper_end):] = False

            if all(~valid_mask):
                frequencies[i] = np.nan
                continue

            # fit Parabola
            parafit = Polynomial.fit(test_freqs[valid_mask], amps[valid_mask], 2)
            func = parafit.convert()

            # find frequency where derivative is 0
            deriv = func.deriv(1)
            freq = deriv.roots()[0]

            frequencies[i] = freq
    else:
        frequencies[:] = f_beacon_estimate

    # evaluate fourier transform at freq for each trace
    for i, freq in enumerate(frequencies):
        if freq is np.nan:
            phases[i] = np.nan
            amplitudes[i] = np.nan
            continue

        real, imag = direct_fourier_transform(freq, t_trace, traces[i])

        phases[i] = np.arctan2(real, imag)
        amplitudes[i] = 1/len(t_trace) * (real**2 + imag**2)**0.5

    return frequencies, phases, amplitudes

if __name__ == "__main__":
    from os import path
    import sys

    
    f_beacon_band = (49e-3,55e-3) #GHz
    allow_frequency_fitting = True
    read_frequency_from_file = 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']

            freqs, phases, amps = find_beacon_in_traces( 
                    traces[1:-1], traces[0],
                    f_beacon_estimate=f_beacon,
                    frequency_fit=allow_frequency_fitting,
                    f_beacon_estimate_band=f_beacon_estimate_band
                    )

            # only take Ex for now
            frequency = freqs[-1]
            phase = phases[-1]
            amplitude = amps[-1]

            print(frequency, phase, amplitude)

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

            found_data[i] = frequency, phase, amplitude

    # show histogram of found frequencies
    if True:
        import matplotlib.pyplot as plt
        if True or allow_frequency_fitting:
            fig, ax = plt.subplots()
            ax.set_xlabel("Frequency")
            ax.set_ylabel("Counts")
            ax.hist(found_data[:,0], bins='auto', density=False)

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

        plt.show()