#!/usr/bin/env python
"""
Script to calculate probabilities for atmospheric neutrino oscillations. To run
it:

    $ ./calculate-atmospheric-oscillations --atmosphere-mode 3 --num-prec 5e-4

which will then produce the following files:

nue_osc_prob.txt: probability for nues to oscillate
nbe_osc_prob.txt: probability for anti nues to oscillate
num_osc_prob.txt: probability for numus to oscillate
nbm_osc_prob.txt: probability for anti numus to oscillate

The format of each text file is the same. It has the following columns:

    1. Energy of the neutrino in GeV.
    2. Cosine of the zenith angle.
    3. Probability to oscillate into nue (or antinue)
    4. Probability to oscillate into numu (or antinumu)
    5. Probability to oscillate into nutau (or antinutau)
"""
from __future__ import print_function, division
import numpy as np
from time import time
from NuCraft import *

# depth of the SNO detector (km)
# currently just converted 6800 feet -> km
SNO_DEPTH = 2.072

# number of energy bins
eBins = 1000
# number of zenith angle bins
zBins = 100
# zenith angle bins
zList = np.arccos(np.linspace(-1,1,zBins))
# energy range in GeV
eList = np.logspace(np.log10(0.01), np.log10(10), eBins)

# parameters from PDG 2019
THETA23 = np.arcsin(np.sqrt(0.512))/np.pi*180.
THETA13 = np.arcsin(np.sqrt(0.0218))/np.pi*180.
THETA12 = np.arcsin(np.sqrt(0.307))/np.pi*180.
DM21   = 7.53e-5
DM31   = 2.444e-3 + DM21

# Akhemdov is implicitely assuming an electron-to-neutron ratio of 0.5; 
# he is also using the approximation DM31 = DM32;
# if you want to reproduce his numbers exactly, switch the lines below, and turn
# atmosphereMode to 0 (no handling of the atmosphere because of )
# AkhmedovOsci = NuCraft((1., DM21, DM31-DM21), [(1,2,theta12),(1,3,theta13,0),(2,3,theta23)],
#                        earthModel=EarthModel("prem", y=(0.5,0.5,0.5))
#                        detectorDepth=0., atmHeight=0.)
AkhmedovOsci = NuCraft((1., DM21, DM31), [(1,2,THETA12),(1,3,THETA13,0),(2,3,THETA23)],detectorDepth=SNO_DEPTH)

def pdg_code_to_string(pdg):
    if pdg == 12:
        return "nue"
    elif pdg == -12:
        return "nbe"
    elif pdg == 14:
        return "num"
    elif pdg == -14:
        return "nbm"
    elif pdg == 16:
        return "nut"
    elif pdg == -16:
        return "nbt"
    else:
        raise ValueError("Unknown pdg code %i" % pdg)

if __name__ == '__main__':
    import argparse

    atmosphereMode = 3   # default: efficiently calculate eight path lenghts per neutrino and take the average

    # This parameter governs the precision with which nuCraft computes the weights; it is the upper
    # limit for the deviation of the sum of the resulting probabilities from unitarity.
    # You can verify this by checking the output plot example-standAlone2.png.
    numPrec = 5e-4

    parser = argparse.ArgumentParser("script to calculate atmospheric neutrino oscillations")
    parser.add_argument("--atmosphere-mode", type=int, default=atmosphereMode, help="controls the handling of the atmospheric interaction altitude")
    parser.add_argument("--num-prec", type=float, default=numPrec, help="parameter which governs the precision with which nuCraft computs the weights")
    args = parser.parse_args()


    # 12, -12:  NuE, NuEBar
    # 14, -14:  NuMu, NuMuBar
    # 16, -16:  NuTau, NuTauBar
    for pType in [12,-12,14,-14]:
        print("Calculating oscillation probabilities for %s..." % pdg_code_to_string(pType))

        # saving the current time to measure the time needed for the execution of the following code
        t = time()

        ### using lists (arrays, actually)
        zz, ee = np.meshgrid(zList, eList)
        zListLong = zz.flatten()
        eListLong = ee.flatten()
        tListLong = np.repeat(pType,len(eListLong))

        # actual call to nuCraft for weight calculations:
        prob = AkhmedovOsci.CalcWeights((tListLong, eListLong, zListLong), numPrec=args.num_prec, atmMode=args.atmosphere_mode)

        prob = np.array(prob)

        save = np.empty((prob.shape[0],5))
        save[:,0] = eListLong
        save[:,1] = np.cos(zListLong)
        save[:,2:] = prob

        print("Calculating the probabilities took %.0f seconds." % (time()-t))

        np.savetxt("%s_osc_prob.txt" % pdg_code_to_string(pType),save,header="Energy (GeV), Cos(zenith), P(nue), P(num), P(nut)")