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

"""
Show Signal to noise for the original simulation signal,
the beacon signal and the combined signal for each antenna
"""

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

from earsim import REvent, block_filter
import aa_generate_beacon as beacon
import lib

if __name__ == "__main__":
    from os import path
    import sys
    import matplotlib
    import os
    if os.name == 'posix' and "DISPLAY" not in os.environ:
        matplotlib.use('Agg')

    from scriptlib import MyArgumentParser
    parser = MyArgumentParser()
    args = parser.parse_args()

    fname = "ZH_airshower/mysim.sry"

    fig_dir = args.fig_dir
    show_plots = args.show_plots

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

    # create fig_dir
    if fig_dir:
        os.makedirs(fig_dir, exist_ok=True)

    # Read in antennas from file
    f_beacon, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)
    _, __, txdata = beacon.read_tx_file(tx_fname)

    # general properties
    dt = antennas[0].t[1] - antennas[0].t[0] # ns
    pb = lib.passband(30e-3, 80e-3) # GHz
    beacon_pb = lib.passband(f_beacon-1e-3, f_beacon+1e-3) # GHz

    beacon_amp = np.max(txdata['amplitudes'])# mu V/m

    if True: # Apply filter to raw beacon/noise to compare with Filtered Traces
        myfilter = lambda x: block_filter(x, dt, pb[0], pb[1])
    else: # Compare raw beacon/noise with Filtered Traces
        myfilter = lambda x: x

    ##
    ## Beacon vs Noise SNR
    ##
    if True:
        beacon_snrs = [ lib.signal_to_noise(myfilter(beacon_amp*ant.beacon), myfilter(ant.noise), samplerate=1/dt, signal_band=beacon_pb) for ant in antennas ]

        fig, ax = plt.subplots()
        ax.set_title("Maximum Beacon/Noise SNR")
        ax.set_xlabel("Antenna no.")
        ax.set_ylabel("SNR")
        ax.plot([ int(ant.name) for ant in antennas], beacon_snrs, 'o', ls='none')
        
        if fig_dir:
            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".beacon_vs_noise_snr.pdf"))

    ##
    ## Beacon vs Total SNR
    ##
    if True:
        beacon_snrs = [ lib.signal_to_noise(myfilter(beacon_amp*ant.beacon), ant.E_AxB, samplerate=1/dt, signal_band=beacon_pb) for ant in antennas ]

        fig, ax = plt.subplots()
        ax.set_title("Maximum Beacon/Total SNR")
        ax.set_xlabel("Antenna no.")
        ax.set_ylabel("SNR")
        ax.plot([ int(ant.name) for ant in antennas], beacon_snrs, 'o', ls='none')

        if fig_dir:
            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".beacon_vs_total_snr.pdf"))
    ##
    ## Airshower signal vs Noise SNR
    ##
    if True:
        shower_snrs = [ lib.signal_to_noise(ant.E_AxB, myfilter(ant.noise), samplerate=1/dt, signal_band=pb) for ant in antennas ]

        fig, ax = plt.subplots()
        ax.set_title("Total (Signal+Beacon+Noise)/Noise SNR")
        ax.set_xlabel("Antenna no.")
        ax.set_ylabel("SNR")
        ax.plot([ int(ant.name) for ant in antennas], shower_snrs, 'o', ls='none')
        
        if fig_dir:
            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".total_snr.pdf"))
    
    if show_plots:
        plt.show()