#!/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 collections import namedtuple

from earsim import REvent
import aa_generate_beacon as beacon
import lib


passband = namedtuple("passband", ['low', 'high'], defaults=[0, np.inf])

def get_freq_spec(val,dt):
    """From earsim/tools.py"""
    fval = np.fft.fft(val)[:len(val)//2]
    freq =  np.fft.fftfreq(len(val),dt)[:len(val)//2]
    return fval, freq


def bandpass_samples(samples, samplerate, band=passband()):
        """
        Bandpass the samples with this passband.
        This is a hard filter.
        """
        fft, freqs = get_freq_spec(samples, samplerate)

        fft[ ~ self.freq_mask(freqs) ] = 0

        return np.fft.irfft(fft)

def bandpass_mask(freqs, band=passband()):
    low_pass = abs(freqs) <= band[1]
    high_pass = abs(freqs) >= band[0]

    return low_pass & high_pass

def bandpower(samples, samplerate=1, band=passband(), normalise_bandsize=True):
    fft, freqs = get_freq_spec(samples, samplerate)

    bandmask = bandpass_mask(freqs, band=band)

    if normalise_bandsize:
        bins = np.count_nonzero(bandmask, axis=-1)
    else:
        bins = 1

    power = np.sum(np.abs(fft[bandmask])**2)

    return power/bins

def signal_to_noise(samples, noise, samplerate=1, signal_band=passband(), noise_band=None):
    if noise_band is None:
        noise_band = signal_band

    if noise is None:
        noise = samples

    noise_power = bandpower(noise, samplerate, noise_band)

    signal_power = bandpower(samples, samplerate, signal_band)

    return (signal_power/noise_power)**0.5


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')

    f_beacon_band = (49e-3,55e-3) #GHz

    fname = "ZH_airshower/mysim.sry"

    fig_dir = "./figures/"
    show_plots = not False

    ####
    fname_dir = path.dirname(fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)
    time_diffs_fname = 'time_diffs.hdf5' if not True else antennas_fname

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

    # Read in antennas from file
    _, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)
    # Read original REvent
    ev = REvent(fname)

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

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

        fig, ax = plt.subplots()
        ax.set_title("Beacon SNR")
        ax.set_xlabel("Antenna")
        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_snr.pdf"))

    ##
    ## Airshower signal vs Noise SNR
    ##
    if True:
        shower_snrs = [ signal_to_noise(ant.E_AxB, ant.noise, samplerate=1/dt, signal_band=pb) for ant in antennas ]

        fig, ax = plt.subplots()
        ax.set_title("Shower SNR")
        ax.set_xlabel("Antenna")
        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".shower_snr.pdf"))
    
    if show_plots:
        plt.show()