#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
import numpy.fft as ft

import aa_generate_beacon as beacon
from view_orig_ant0 import plot_antenna_geometry
import lib
from earsim import Antenna


if __name__ == "__main__":
    import os.path as path

    fname = "ZH_airshower/mysim.sry"

    plot_ft_amplitude = True
    plot_geometry = True

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

    f_beacon, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)

    if not True:
        idx = [0, 1, len(antennas)//2, len(antennas)//2+1, -2, -1]
    elif not True:
        idx = np.arange(1, 20, 2, dtype=int)
    elif True:
        # center 6 antennas
        names = [55, 56, 57, 65, 66, 45, 46]

        idx = [ i for i, ant in enumerate(antennas) if int(ant.name) in names ]

        [ print(antennas[i].name) for i in names ]

    fig1, axs = plt.subplots(1+plot_ft_amplitude*2)
    if not plot_ft_amplitude:
        axs = [axs]
    axs[0].set_xlabel('t [ns]')
    axs[0].set_ylabel('[$\mu$V/m]')
    if plot_ft_amplitude:
        axs[1].set_xlabel('f [GHz]')
        axs[1].set_ylabel('Power')

        axs[2].set_ylabel("Phase")
        axs[2].set_xlabel('f [GHz]')
        axs[2].set_ylim(-np.pi,+np.pi)

    colorlist = []
    for i in idx:
        ant = antennas[i]

        n_samples = len(ant.t)
        samplerate = (ant.t[-1] - ant.t[0])/n_samples

        axs[0].axvline(ant.t[0], color='k', alpha=0.5)

        if True: # total E field
            mydict = dict(AxB=ant.E_AxB)
        elif False: # polarisations
            mydict = dict(x=ant.Ex, y=ant.Ex, z=ant.Ez)
        else: # beacon
            mydict = dict(b=ant.beacon)

        for j, (direction, trace) in enumerate(mydict.items()):
            l = axs[0].plot(ant.t, trace, label=f"$E_{{{direction}}}$ {ant.name}", alpha=0.8)

            #if False and j == 0 and 't0' in ant.attrs:
            #    axs[0].axvline(ant.attrs['t0'], color=l[0].get_color(), alpha=0.5)

            colorlist.append(l[0].get_color())

            if not plot_ft_amplitude:
                continue

            freqs = ft.fftfreq(n_samples, 1/samplerate)[:n_samples//2]
            fft = 2*ft.fft(trace)[:n_samples//2]/n_samples

            #axs[1].plot(freqs, np.abs(fft)**2, color=l[0].get_color())

            if True:
                cft = lib.direct_fourier_transform(f_beacon, ant.t, trace)
                amp = 2*len(ant.t) * (cft[0]**2 + cft[1]**2)

                #axs[0].axhline(amp, color=l[0].get_color())

                print(amp)
                phase = np.arctan2(cft[0],cft[1])
                axs[1].plot(f_beacon, amp, color=l[0].get_color(), marker='3', alpha=0.8, ms=30)
                axs[2].plot(f_beacon, phase, color=l[0].get_color(), marker='3', alpha=0.8, ms=30)

    if plot_ft_amplitude:
        fig1.legend(loc='center right', ncol=min(2, len(idx)))
    else:
        fig1.legend(loc='upper right', ncol=min(3, len(idx)))

    if plot_geometry:
        if len(mydict) == 1:
            geom_colorlist = colorlist
        else:
            # only take the colour belonging to mydict[0]
            geom_colorlist = [ colorlist[len(mydict)*(i)] for i in range(len(colorlist)//len(mydict)) ]

        fig2, axs2 = plt.subplots(1)
        plot_antenna_geometry(antennas, ax=axs2, plot_max_values=False, color='grey', plot_names=False)
        plot_antenna_geometry([ antennas[i] for i in idx], ax=axs2, colors=geom_colorlist, plot_max_values=False)

        axs2.plot(tx.x, tx.y, marker='X', color='k')
        axs2.set_title("Geometry with selected antennas")

    #fig1.savefig('./fig1.png')
    plt.show()