#!/usr/bin/env python3
# vim: fdm=indent 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()

    # Bandpass
    parser.add_argument('-p', '--use-passband', type=bool, default=True, help='(Default: %(default)d)')
    parser.add_argument('-l', '--passband-low',  type=float, default=30e-3, help='Lower frequency [GHz] of the passband filter. (set -1 for np.inf) (Default: %(default)d)')
    parser.add_argument('-u', '--passband-high', type=float, default=80e-3, help='Upper frequency [GHz] of the passband filter. (set -1 for np.inf) (Default: %(default)d)')

    args = parser.parse_args()

    figsize = (12,8)

    fig_dir = args.fig_dir
    show_plots = args.show_plots

    ####
    fname_dir = args.data_dir
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)
    tx_fname = path.join(fname_dir, beacon.tx_fname)
    snr_fname = path.join(fname_dir, beacon.snr_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, traces_key='filtered_traces')
    _, __, txdata = beacon.read_tx_file(tx_fname)

    # general properties
    dt = antennas[0].t[1] - antennas[0].t[0] # ns
    beacon_pb = lib.passband(f_beacon, None) # GHz
    beacon_amp = np.max(txdata['amplitudes'])# mu V/m

    # General Bandpass
    low_bp  = args.passband_low  if args.passband_low  >= 0 else np.inf # GHz
    high_bp = args.passband_high if args.passband_high >= 0 else np.inf # GHz
    pb = lib.passband(low_bp, high_bp) # GHz

    noise_pb = pb

    if args.use_passband: # 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

    ##
    ## Debug plot of Beacon vs Noise SNR
    ##
    if True:
        ant = antennas[0]

        fig, ax = plt.subplots(figsize=figsize)
        _debug_snrs = lib.signal_to_noise(myfilter(beacon_amp*ant.beacon), myfilter(ant.noise), samplerate=1/dt, signal_band=beacon_pb, noise_band=noise_pb, debug_ax=ax)

        ax.legend(title="$\\langle SNR \\rangle$ = {: .1e}".format(np.mean(_debug_snrs)))

        ax.set_title("Spectra and passband")
        ax.set_xlabel("Frequency [GHz]")
        ax.set_ylabel("Amplitude")
        low_x, high_x = min(beacon_pb[0], noise_pb[0]), max(beacon_pb[1] or 0, noise_pb[1])
        ax.set_xlim(low_x, high_x)

        if fig_dir:
            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".debug_plot.pdf"))
    ##
    ## Beacon vs Noise SNR
    ##
    if True:
        N_samples = len(antennas[0].beacon)
        beacon_snrs = [ lib.signal_to_noise(myfilter(beacon_amp*ant.beacon), myfilter(ant.noise), samplerate=1/dt, signal_band=beacon_pb, noise_band=noise_pb) for i, ant in enumerate(antennas) ]

        # write mean and std to file
        beacon.write_snr_file(snr_fname, beacon_snrs)

        fig, ax = plt.subplots(figsize=figsize)
        ax.set_title(f"Maximum Beacon/Noise SNR (N_samples:{N_samples:.1e})")
        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, noise_band=pb) for ant in antennas ]

        fig, ax = plt.subplots(figsize=figsize)
        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"))

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

        fig, ax = plt.subplots(figsize=figsize)
        ax.set_title("Maximum Airshower/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".airshower_vs_noise_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, noise_band=pb) for ant in antennas ]

        fig, ax = plt.subplots(figsize=figsize)
        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()