#!/usr/bin/env python
# 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/>.
"""
Script to print the event rate per year for different neutrino interactions. It
works by reading the MCPL files produced by GENIE to determine the number of
neutrino interactions and dividing by the livetime in the autosno run_info.log
file. To run it, just run:

    $ ./print-event-rate /path/to/mcpl/files/*.dat.gz --run-info /path/to/autosno/log/run_info.log

which will print out a latex style table of the event rates.
"""
from __future__ import print_function, division
import numpy as np
import pandas as pd
from sddm import read_hdf, AV_RADIUS, PSUP_RADIUS
import gzip
from sddm.renormalize import *

def get_mcgn(filename):
    ev = read_hdf(filename, "ev")
    mcgn = read_hdf(filename, "mcgn").assign(run=ev.run[0])
    mcgn = mcgn.groupby(['run','evn']).first().reset_index()
    return mcgn

if __name__ == '__main__':
    import argparse
    import sys
    from os.path import join, split
    from sddm import *

    parser = argparse.ArgumentParser("recreate MCPL files")
    parser.add_argument("filenames", nargs='+', help="mcpl filenames")
    parser.add_argument("--run-info", required=True, help="run_info.log autosno file")
    args = parser.parse_args()

    cc_av = {nu: 0 for nu in range(30,36)}
    cc_psup = {nu: 0 for nu in range(30,36)}
    nc_av = {nu: 0 for nu in range(30,36)}
    nc_psup = {nu: 0 for nu in range(30,36)}
    livetime = 0.0
    run_info = np.genfromtxt(args.run_info,usecols=range(4),dtype=(np.int,np.int,np.double,np.double))
    for filename in args.filenames:
        head, tail = split(filename)
        try:
            mcpl = read_mcpl_df(filename)
        except Exception as e:
            print_warning("failed to read mcpl file '%s'" % filename)
            continue
        run = int(tail.split('_')[3])
        for _, row in mcpl.iterrows():
            nu = row['nu_id']
            r = row['nu_r']
            if r < AV_RADIUS:
                if not row['cc']:
                    # nc event
                    nc_av[nu] += row['flux_weight']
                else:
                    # cc event
                    cc_av[nu] += row['flux_weight']
            if r < PSUP_RADIUS:
                if not row['cc']:
                    # nc event
                    nc_psup[nu] += row['flux_weight']
                else:
                    # cc event
                    cc_psup[nu] += row['flux_weight']
        for i in range(run_info.shape[0]):
            if run_info[i][0] == run:
                livetime += run_info[i][3]

    for nu in range(30,36):
        cc_av[nu] *= 3600*24*365/livetime/100
        cc_psup[nu] *= 3600*24*365/livetime/100
        nc_av[nu] *= 3600*24*365/livetime/100
        nc_psup[nu] *= 3600*24*365/livetime/100

    for nu in range(30,36):
        line = [cc_av[nu],nc_av[nu],
                cc_av[nu]+nc_av[nu],
                cc_psup[nu],nc_psup[nu],
                cc_psup[nu]+nc_psup[nu]]
        print(' & '.join(["%.1f" % rate for rate in line]))