#!/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)

    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([ant.attrs['beacon_phase'] for ant in 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("Amplitude at beacon frequency at each antenna")
    axs.set_aspect('equal', 'datalim')
    axs.set_xlabel('[m]')
    axs.set_ylabel('[m]')

    if True:
        # underlie a calculate phase field
        xs = np.linspace( np.min(x), np.max(x), 50)
        ys = np.linspace( np.min(y), np.max(y), 50)
        phases, (xs, ys) = lib.phase_field_from_tx(xs, ys, tx, f_beacon, return_meshgrid=False)
        sc2 = axs.scatter(xs, ys, c=phases, alpha=0.5, zorder=-5)
        fig.colorbar(sc2, ax=axs)

    sc = axs.scatter(x, y, c=vals, s=sizes)

    axs.plot(tx.x, tx.y, marker='X', color='k')

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

    plt.show()