############################################################################################################

# importing necessary packages

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import RegularGridInterpolator

############################################################################################################

# loading data from NuTauSim (cite https://arxiv.org/abs/1901.08498) and building interpolators

# the probability interpolator takes as arguments N pairs of (log10(neutrino_energy(eV)), Earth emergence angle (degrees))
# and returns an array of length N representing the log of the Earth emergence probability for each pair
prob_data = np.load("Tau_Emergence_Probabilities.npy", allow_pickle = True)
energies, angles, probabilities = np.log10(prob_data[0]), prob_data[1], np.log10(prob_data[2])
prob_interpolator = RegularGridInterpolator((energies, angles), probabilities, bounds_error=False, fill_value = -10)

# the energy interpolator takes as arguments N pairs of (log10(neutrino_energy(eV)), Earth emergence angle (degrees), y)
# and returns an array of length N representing the log of a randomly sampled tau-lepton fractional energy (Etau/Eneutrino) for each triple
# the actual random sampling (and need for the value y, which is randomly sampled) is done via inverse tranform sampling: https://en.wikipedia.org/wiki/Inverse_transform_sampling
cdf_data = np.load("Tau_Emergence_Energy_CDFs.npy", allow_pickle = True)
energies, angles, cdfs, cdfxs = np.log10(cdf_data[0]), cdf_data[1], cdf_data[2], cdf_data[3]
energy_interpolator = RegularGridInterpolator((energies, angles, cdfs), cdfxs, bounds_error=False, fill_value = -10)

# one note of caution: my data was not generated below Earth emergence angles of 0.1 degrees, so results below that are undefined

############################################################################################################

# example use 1: generate an Earth emergence probability curve

N = 1000 # number of angular samples
angles = np.linspace(0.1,30,N) # spacing of Earth emergence angle (in degrees)
ens = 1e18*np.ones(N) # input neutrino energy in eV (in this case, the same for every angle, but doesn't need to be)
probs = 10**prob_interpolator(np.array([np.log10(ens), angles]).T) # probability of Earth emergence

fig = plt.figure()
plt.plot(angles, probs)
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Earth Emergence Angle (Degrees)')
plt.ylabel('Earth Emergence Probability (Pexit)')
plt.grid(which = 'both')

############################################################################################################

# example use 2: get emergent energies given input events (how you will likely use this in your MC)

N = 1000000 # number of taus you want to simulate (may need to do this in a for loop if RAM has too high of a usage)
angles = 30*np.random.rand(N) # Earth emergence angle of each event (in degrees). Currently, uniformly random. This will be replaced by your MC
ens = 1e17*np.ones(N) # input neutrino energy in eV (in this case, the same for every angle, but doesn't need to be)
sample_ys = np.random.rand(N) # a random value between 0 and 1 to do inverse transform sampling on the tau energy CDF tables

# my data only goes down to 0.1 degree Earth emergence angle, so below that, results are undefined and you need to be somewhat careful
valid_indices = angles >= 0.1
angles = angles[valid_indices]
ens = ens[valid_indices]
sample_ys = sample_ys[valid_indices]

energies = ens*10**energy_interpolator(np.array([np.log10(ens), angles, sample_ys]).T) # energies of Earth emergent tau leptons

fig = plt.figure()
plt.hist(np.log10(energies), bins = 50, density = True)
plt.xlabel(r'$E_{\tau}$'+' (eV)')
plt.show()

############################################################################################################