from __future__ import print_function, division
import numpy as np
from NuRadioMC.utilities import units
import os
from scipy import interpolate as intp
import glob
import pickle
import sys
from radiotools import plthelpers as php
from matplotlib import pyplot as plt
from NuRadioMC.SignalGen.ARZ import ARZ
import logging
from scipy import signal
import matplotlib.gridspec as gridspec

logger = logging.getLogger("test")
logging.basicConfig(level=logging.WARNING)
rho = 0.924 * units.g / units.cm ** 3  # density g cm^-3

plot=False

if __name__ == "__main__":
    sampling_rate = 5 * units.GHz
    dt = 1. / sampling_rate
    T2 = 50 * units.ns
    N2 = np.int(np.round(T2 / dt))
    n_index = 1.78
    theta = np.arccos(1. / n_index)
    R = 5 * units.km
    T = 600 * units.ns
    N = np.int(np.round(T / dt))
    tt = np.arange(0, dt * N, dt)
    
    dCs = np.append(np.arange(-20, -5, 1), np.append(np.arange(-5,-1,0.2), np.arange(-1, 0, .05)))
    dCs = np.append(np.append(dCs, [0]), -1 * dCs)
    dCs = np.sort(dCs) * units.deg
    print('generating Askaryan pulses for the following viewing angles')
    print(dCs/units.deg)
    
    lib_path = "/Users/cglaser/work/ARIANNA/data/ARZ/v1.1/library_v1.1.pkl"
    a = ARZ.ARZ(library=lib_path, interp_factor=100)
    
    showers = {}
    showers['meta'] = {'dt': dt,
                       'n_index': n_index,
                       'R': R}
    
    with open(lib_path, 'rb') as fin:
        lib = pickle.load(fin)
        for iS, shower_type in enumerate(lib):  # loop through shower types
            if(shower_type not in showers):
                showers[shower_type] = {}
            print("shower type {}".format(shower_type))
            b = []
            nb = []
            for iE, E in enumerate(lib[shower_type]):  # loop through energies
                print('E = {:.2g}eV'.format(E))
                
                nE = len(lib[shower_type][E]['charge_excess'])
                if(E not in showers[shower_type]):
                    showers[shower_type][E] = {}
                for iN in range(nE):  # loop through shower realizations
                    if(iN not in showers[shower_type][E]):
                        showers[shower_type][E][iN] = {}
                    if plot:
                        fig, ax = plt.subplots(1, 1)
                        fig2, ax2 = plt.subplots(1, 1)
                    for dC in dCs:  # loop through different cherenkov angles
                        theta1 = theta + dC
                        trace, Lmax = a.get_time_trace(E, theta1, N, dt, shower_type, n_index, R, iN=iN, output_mode='Xmax')
                        iMax = np.argmax(np.abs(trace[1]))
                        i1 = iMax - N2 // 2
                        if((i1 + N2) >= N) or (i1 < 0):
                            print("tracelength to short, i1 = {}".format(i1))
#                             raise IndexError("tracelength to short, i1 = {}".format(i1))
                        t0 = tt[i1] - tt.mean()
                        showers[shower_type][E][iN][theta1] = {}
                        showers[shower_type][E][iN][theta1]['t0'] = t0
                        showers[shower_type][E][iN][theta1]['Lmax'] = Lmax
                        showers[shower_type][E][iN][theta1]['trace'] = trace[1][i1:(i1 + N2)]
                        if(plot):
                            ax.plot(tt[i1:(i1 + N2)], trace[1][i1:(i1 + N2)] / np.abs(trace[1]).max())
                            ax2.plot(tt[i1:(i1 + N2)], np.abs(signal.hilbert(trace[1][i1:(i1 + N2)])))
                    if(plot):
                        ax.set_title("{} E = {:.2g}eV i = {:d}".format(shower_type, E / units.eV, iN))
                        ax2.set_title("{} E = {:.2g}eV i = {:d}".format(shower_type, E / units.eV, iN))
                        ax.set_xlabel('time [ns]')
                        ax2.set_xlabel('time [ns]')
                        ax.set_ylabel('normalized amplitude')
                        ax2.semilogy(True)
                        fig.tight_layout()
                        fig2.tight_layout()
#                         plt.show()
                        fig.savefig("plots/{}_{:.2g}eV_{:03d}.png".format(shower_type, E / units.eV, iN))
                        fig2.savefig("plots/{}_{:.2g}eV_{:03d}_log.png".format(shower_type, E / units.eV, iN))
                        plt.close("all")
                            
        with open("ARZ_library_v1.1.pkl", 'wb') as fout:
            pickle.dump(showers, fout, protocol=2)
