#!/usr/bin/env python3
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.

import ROOT
import numpy as np

# on retina screens, the default plots are way too small
# by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
# Qt5 will scale everything using the dpi in ~/.Xresources
import matplotlib
matplotlib.use("Qt5Agg")

def tick_formatter(x, pos):
    if x < 1:
        return '%.1f' % x
    else:
        return '%.0f' % x

# FIXME: What is this for the salt phase?
NEUTRON_DETECTION_EFFICIENCY = 0.5

# FIXME: What is the index for D2O?
INDEX_HEAVY_WATER = 1.3

# fractional energy resolution
# from Richie's thesis page 134
ENERGY_RESOLUTION = 0.05

def pdg_code_to_string(pdg):
    A = int(("%010i" % event.tgt)[-4:-1])
    Z = int(("%010i" % event.tgt)[-7:-4])
    if Z == 1:
        atom = 'H'
    elif Z == 8:
        atom = 'O'
    elif Z == 6:
        atom = 'C'
    else:
        raise NotImplementedError("unknown atom %i" % Z)
    return '%i%s' % (A,atom)

def get_reaction(event):
    reactants = []
    products = []
    if event.neu == 12:
        reactants.append('ve')
    elif event.neu == -12:
        reactants.append('vebar')
    elif event.neu == 14:
        reactants.append('vu')
    elif event.neu == -14:
        reactants.append('vubar')
    elif event.neu == 16:
        reactants.append('vt')
    elif event.neu == -16:
        reactants.append('vtbar')

    if event.hitnuc == 2212:
        reactants.append('p')
    elif event.hitnuc == 2112:
        reactants.append('n')
    elif event.hitnuc == 0:
        reactants.append(pdg_code_to_string(event.tgt))
    else:
        print("unknown nucleon %i" % event.hitnuc)

    if event.cc:
        if event.neu == 12:
            products.append('e-')
        elif event.neu == -12:
            products.append('e+')
        elif event.neu == 14:
            products.append('u-')
        elif event.neu == -14:
            products.append('u+')
        elif event.neu == 16:
            products.append('t-')
        elif event.neu == -16:
            products.append('t+')
    elif event.nc:
        if event.neu == 12:
            products.append('ve')
        elif event.neu == -12:
            products.append('vebar')
        elif event.neu == 14:
            products.append('vu')
        elif event.neu == -14:
            products.append('vubar')
        elif event.neu == 16:
            products.append('vt')
        elif event.neu == -16:
            products.append('vtbar')
    else:
        products.append("???")

    for pdg in event.pdgf:
        if pdg == 2112:
            products.append('n')
        elif abs(pdg) == 11:
            # e- or e+
            if pdg == 11:
                products.append('e-')
            else:
                products.append('e+')
        elif pdg == 22:
            # gamma
            products.append('gamma')
        elif pdg == 111:
            # pi0
            products.append('pi0')
        elif abs(pdg) == 211:
            # pi+/-
            if pdg == 211:
                products.append('pi+')
            else:
                products.append('pi-')
        elif abs(pdg) == 311:
            if pdg == 311:
                products.append('K0')
            else:
                products.append('K0bar')
        elif abs(pdg) == 321:
            # K+/-
            if pdg == 321:
                products.append('K+')
            else:
                products.append('K-')
        elif abs(pdg) == 3222:
            products.append('Sigma+')
        elif abs(pdg) == 3112:
            products.append('Sigma-')
        elif abs(pdg) == 3122:
            products.append('Delta')
        elif pdg == 2212:
            products.append('p')
        elif int(("%010i" % abs(pdg))[0]) == 1:
            products.append(pdg_code_to_string(pdg))
        else:
            print("unknown pdg code %i" % pdg)

    return ' + '.join(reactants) + ' -> ' + ' + '.join(products)

if __name__ == '__main__':
    import argparse
    import matplotlib.pyplot as plt
    from collections import Counter

    parser = argparse.ArgumentParser("script to analyze GENIE 'ntuple' ROOT files")
    parser.add_argument("filenames", nargs='+', help="GENIE ROOT files")
    args = parser.parse_args()

    bins = np.logspace(-1,2,100)

    El = []
    total_neutrons = []
    total_neutrons_detected = []
    E = []
    KE = []
    r = []
    total_nrings = []
    total_e_like_rings = []
    total_u_like_rings = []
    reactions = Counter()
    for filename in args.filenames:
        print("analyzing %s" % filename)
        f = ROOT.TFile(filename)
        T = f.Get("gst")

        for event in T:
            neutrons = 0
            nrings = 0
            e_like_rings = 0
            u_like_rings = 0
            ke = 0

            if event.cc:
                if abs(event.neu) == 12:
                    e_like_rings = 1
                else:
                    u_like_rings = 1
                nrings = 1
                ke += event.El
            elif event.nc:
                pass
            else:
                print("event is not cc or nc!")
                continue

            for i, pdg in enumerate(event.pdgf):
                if pdg == 2112:
                    neutrons += 1
                elif abs(pdg) == 11:
                    # e- or e+
                    if event.Ef[i] > 0.1:
                        # for now assume we only count rings from electrons
                        # with > 100 MeV
                        nrings += 1
                        e_like_rings += 1
                    ke += event.Ef[i]
                elif pdg == 22:
                    # gamma
                    if event.Ef[i] > 0.1:
                        # for now assume we only count rings from gammas with >
                        # 100 MeV
                        nrings += 1
                        e_like_rings += 1
                    ke += event.Ef[i]
                elif pdg == 111:
                    # pi0
                    nrings += 1
                    e_like_rings += 1
                    ke += event.Ef[i]
                elif abs(pdg) == 211:
                    # pi+/-
                    # momentum of ith particle in hadronic system
                    p = np.sqrt(event.pxf[i]**2 + event.pyf[i]**2 + event.pzf[i]**2)
                    # velocity of ith particle (in units of c)
                    # FIXME: is energy total energy or kinetic energy?
                    v = p/event.Ef[i]
                    if v > 1/INDEX_HEAVY_WATER:
                        # if the pion is above threshold, we assume that it
                        # produces 2 muon like rings
                        nrings += 2
                        u_like_rings += 2
                    else:
                        # if the pion is not above threshold, we assume that it
                        # produces 1 muon like ring
                        nrings += 1
                        u_like_rings += 1
                    # FIXME: should actually be a beta distribution
                    p = np.sqrt(event.pxf[i]**2 + event.pyf[i]**2 + event.pzf[i]**2)
                    m = np.sqrt(event.Ef[i]**2 - p**2)
                    ke += event.Ef[i] - m
                elif abs(pdg) in [2212,3222,311,321,3122,3112]:
                    # momentum of ith particle in hadronic system
                    p = np.sqrt(event.pxf[i]**2 + event.pyf[i]**2 + event.pzf[i]**2)
                    # velocity of ith particle (in units of c)
                    # FIXME: is energy total energy or kinetic energy?
                    v = p/event.Ef[i]
                    if v > 1/INDEX_HEAVY_WATER:
                        # above cerenkov threshold
                        nrings += 1
                        u_like_rings += 1
                        m = np.sqrt(event.Ef[i]**2 - p**2)
                        ke += event.Ef[i] - m
                elif int(("%010i" % abs(pdg))[0]) == 1:
                    # usually just excited 16O atom which won't produce a lot
                    # of light
                    pass
                else:
                    print("unknown pdg code %i" % pdg)
            total_neutrons.append(neutrons)
            total_neutrons_detected.append(np.random.binomial(neutrons,NEUTRON_DETECTION_EFFICIENCY))
            total_nrings.append(nrings)
            total_e_like_rings.append(e_like_rings)
            total_u_like_rings.append(u_like_rings)
            El.append(event.El)
            E.append(event.calresp0)
            KE.append(ke + np.random.randn()*ke*ENERGY_RESOLUTION)
            r.append(np.sqrt(event.vtxx**2 + event.vtxy**2 + event.vtxz**2))

            if total_neutrons_detected[-1] == 0 and nrings >= 2 and ((e_like_rings == 0) or (u_like_rings == 0)):
                reactions.update([get_reaction(event)])

    total = sum(reactions.values())
    for reaction, count in reactions.most_common(10):
        print("%.0f%% %s" % (count*100/total, reaction))

    El = np.array(El)
    total_neutrons = np.array(total_neutrons)
    total_neutrons_detected = np.array(total_neutrons_detected)
    E = np.array(E)
    KE = np.array(KE)
    r = np.array(r)
    total_nrings = np.array(total_nrings)
    total_e_like_rings = np.array(total_e_like_rings)
    total_u_like_rings = np.array(total_u_like_rings)

    cut1 = (total_neutrons_detected == 0)
    cut2 = (total_neutrons_detected == 0) & (total_nrings >= 2)
    cut3 = (total_neutrons_detected == 0) & (total_nrings >= 2) & ((total_e_like_rings == 0) | (total_u_like_rings == 0))

    El1 = El[cut1]
    El2 = El[cut2]
    El3 = El[cut3]
    E1 = E[cut1]
    E2 = E[cut2]
    E3 = E[cut3]
    KE1 = KE[cut1]
    KE2 = KE[cut2]
    KE3 = KE[cut3]

    plt.figure()
    bincenters = (bins[1:] + bins[:-1])/2
    x = np.repeat(bins,2)
    El_hist, _ = np.histogram(El, bins=bins)
    total_events = El_hist.sum()
    # FIXME: this is just a rough estimate of how many events we expect in 3
    # years based on Richie's thesis which says "Over the 306.4 live days of
    # the D2O phase we expect a total of 192.4 events within the acrylic vessel
    # and 504.5 events within the PSUP.
    El_hist = El_hist*230/total_events
    y = np.concatenate([[0],np.repeat(El_hist,2),[0]])
    El1_hist, _ = np.histogram(El1, bins=bins)
    El1_hist = El1_hist*230/total_events
    y1 = np.concatenate([[0],np.repeat(El1_hist,2),[0]])
    El2_hist, _ = np.histogram(El2, bins=bins)
    El2_hist = El2_hist*230/total_events
    y2 = np.concatenate([[0],np.repeat(El2_hist,2),[0]])
    El3_hist, _ = np.histogram(El3, bins=bins)
    El3_hist = El3_hist*230/total_events
    y3 = np.concatenate([[0],np.repeat(El3_hist,2),[0]])
    plt.plot(x, y, label="All events")
    plt.step(x, y1, where='mid', label="n")
    plt.step(x, y2, where='mid', label="n + nrings")
    plt.step(x, y3, where='mid', label="n + nrings + same flavor")
    plt.xlabel("Energy (GeV)")
    plt.ylabel("Events/year")
    plt.gca().set_xscale("log")
    plt.gca().xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(tick_formatter))
    plt.xlim((0.02,bins[-1]))
    plt.ylim((0,None))
    plt.legend()
    plt.title("Primary Lepton Energy")

    plt.figure()
    KE_hist, _ = np.histogram(KE, bins=bins)
    KE_signal, _ = np.histogram(np.random.randn(1000)*1.0*ENERGY_RESOLUTION + 1.0, bins=bins)
    total_events = KE_hist.sum()
    # FIXME: this is just a rough estimate of how many events we expect in 3
    # years based on Richie's thesis which says "Over the 306.4 live days of
    # the D2O phase we expect a total of 192.4 events within the acrylic vessel
    # and 504.5 events within the PSUP.
    KE_hist = KE_hist*230/total_events
    y = np.concatenate([[0],np.repeat(KE_hist,2),[0]])
    KE1_hist, _ = np.histogram(KE1, bins=bins)
    KE1_hist = KE1_hist*230/total_events
    y1 = np.concatenate([[0],np.repeat(KE1_hist,2),[0]])
    KE2_hist, _ = np.histogram(KE2, bins=bins)
    KE2_hist = KE2_hist*230/total_events
    y2 = np.concatenate([[0],np.repeat(KE2_hist,2),[0]])
    KE3_hist, _ = np.histogram(KE3, bins=bins)
    KE3_hist = KE3_hist*230/total_events
    y3 = np.concatenate([[0],np.repeat(KE3_hist,2),[0]])
    KE_signal = KE_signal*10/np.sum(KE_signal)
    y4 = np.concatenate([[0],np.repeat(KE_signal,2),[0]])
    plt.plot(x, y, label="All events")
    plt.plot(x, y1, label="n")
    plt.plot(x, y2, label="n + nrings")
    plt.plot(x, y3, label="n + nrings + same flavor")
    plt.plot(x, y4, label="1 GeV signal")
    plt.xlabel("Energy (GeV)")
    plt.ylabel(r"Expected Event Rate (year$^{-1}$)")
    plt.gca().set_xscale("log")
    plt.gca().xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(tick_formatter))
    plt.xlim((0.02,bins[-1]))
    plt.ylim((0,None))
    plt.legend()
    plt.title("Approximate Visible Energy")

    plt.figure()
    plt.hist(r, bins=np.linspace(0,8,100), histtype='step')
    plt.xlabel("R (m)")
    plt.title("Radius of Events")

    plt.figure()
    plt.hist(total_neutrons, bins=np.arange(11)-0.5, histtype='step')
    plt.xlabel("Number of Neutrons")
    plt.title("Number of Neutrons")

    plt.figure()
    plt.hist(total_nrings, bins=np.arange(11)-0.5, histtype='step')
    plt.xlabel("Number of Rings")
    plt.title("Number of Rings (approximate)")
    plt.show()