#!/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"
    
    ####
    fname_dir = path.dirname(fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)

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

    subtitle = ""
    if True:
        beacon_frequencies = np.array([ant.attrs['beacon_freq'] for ant in antennas])
        beacon_amplitudes = np.array([ant.attrs['beacon_amplitude'] for ant in antennas])
        beacon_phases = np.array([lib.phase_mod(ant.attrs['beacon_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
    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(f"Amplitude at beacon frequency at each antenna\nf at A0: {beacon_frequencies[0]}GHz" + subtitle)
    axs.set_aspect('equal', 'datalim')
    axs.set_xlabel('[m]')
    axs.set_ylabel('[m]')

    if True:
        if True:
            # 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)
            sc2 = axs.scatter(xs, ys, c=phases, alpha=0.5, zorder=-5, cmap='Spectral_r')
            fig.colorbar(sc2, ax=axs)
    
        sc = axs.scatter(x, y, c=vals, s=sizes, cmap='Spectral_r', edgecolors='k', marker='X')
    
        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)


    fig.savefig(__file__ + ".pdf")
    plt.show()