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

__doc__ = \
"""
Show the beacon amplitude per antenna.
"""

import numpy as np
import h5py
import matplotlib.pyplot as plt

import aa_generate_beacon as beacon
import lib

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

    fname = "ZH_airshower/mysim.sry"

    plot_phase_field = True
    plot_tx = False
    
    ####
    fname_dir = path.dirname(fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)

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

    pretitle = ""
    if True:
        freq_name = list(antennas[0].beacon_info.keys())[0]
        beacon_frequencies = np.array([ant.beacon_info[freq_name]['freq'] for ant in antennas])
        beacon_amplitudes = np.array([ant.beacon_info[freq_name]['amplitude'] for ant in antennas])
        beacon_phases = lib.phase_mod(np.array([ant.beacon_info[freq_name]['phase'] for ant in antennas]))

        if 'true_phase' in antennas[0].beacon_info[freq_name]:
            beacon_true_phases = lib.phase_mod(np.array([ant.beacon_info[freq_name]['true_phase'] for ant in antennas]))
    else:
        subtitle = " Phases from t0"
        beacon_frequencies = np.array([ f_beacon for ant in antennas ])
        beacon_amplitudes = np.array([ 1 for ant in antennas ])
        beacon_phases = np.array([ lib.phase_mod(ant.attrs['t0']*beacon_frequencies[i]*2*np.pi) for i, ant in enumerate(antennas)])

    #####
    sizes = 64
    if True:
        vals = beacon_phases
        colorlabel = '$\\varphi$'
        sizes = 64*(beacon_amplitudes/np.max(beacon_amplitudes))**2
    elif True: # True Phases
        vals = beacon_true_phases
        colorlabel = '$\\sigma_\\varphi$'
        plot_phase_field = False
        plot_tx = False
    else:
        vals = beacon_amplitudes
        colorlabel = "[$\\mu$V/m]"

    x = [ a.x for a in antennas ]
    y = [ a.y for a in antennas ]

    #####

    fig, axs = plt.subplots()
    axs.set_title(pretitle + f"Beacon frequency at A0: {beacon_frequencies[0]:.3e}GHz")
    axs.set_aspect('equal', 'datalim')
    axs.set_xlabel('[m]')
    axs.set_ylabel('[m]')

    if True:
        vmax, vmin = None, None
        if plot_phase_field:
            # underlie a calculate phase field
            if True: # only fill for antennas
                xs = np.linspace( np.min(x), np.max(x), 4*np.ceil(len(antennas)**0.5) -1 )
                ys = np.linspace( np.min(y), np.max(y), 4*np.ceil(len(antennas)**0.5) -1)
            else: # make field from halfway the transmitter
                xs = np.linspace( (tx.x - np.min(x))/2, np.max(x), 500)
                ys = np.linspace( (tx.y - np.min(y))/2, np.max(y), 500)

            phases, (xs, ys) = lib.phase_field_from_tx(xs, ys, tx, f_beacon*1e9,return_meshgrid=False)

            vmax, vmin = max(np.max(phases), np.max(vals)), min(np.min(phases), np.min(vals))
            sc2 = axs.scatter(xs, ys, c=phases, alpha=0.5, zorder=-5, cmap='Spectral_r', vmin=vmin, vmax=vmax)

        if False:
            # do not share the same colours
            fig.colorbar(sc2, ax=axs)
            vmax, vmin = None, None
    
        sc = axs.scatter(x, y, c=vals, s=sizes, cmap='Spectral_r', edgecolors='k', marker='X', vmin=vmin, vmax=vmax)
   
        if plot_tx:
            axs.plot(tx.x, tx.y, marker='X', color='k')
    
        #for i, freq in enumerate(beacon_frequencies):
        #    axs.text(f"{freq:.2e}", (x[i], y[i]))
    
        fig.colorbar(sc, ax=axs, label=colorlabel)
    else:
        phases, (xs, ys) = lib.phase_field_from_tx(x, y, tx, f_beacon*1e9, return_meshgrid=False)

        phase_diffs = vals - lib.phase_mod(phases)
        phase_diffs = lib.phase_mod(phase_diffs)

        print(phases)

        sc = axs.scatter(xs, ys, c=phase_diffs, s=sizes, cmap="Spectral_r")
        axs.plot(tx.x, tx.y, marker='X', color='k')

        fig.colorbar(sc, ax=axs, label=colorlabel)

    if not True:
        fig.savefig(__file__ + ".pdf")
    plt.show()