#!/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 recreate MCPL files for events which never got simulated in SNOMAN.
"""
from __future__ import print_function, division
import numpy as np
import pandas as pd
from sddm import read_hdf
import gzip

def read_mcpl(filename):
    data = np.genfromtxt(filename,skip_header=1)
    rv = {}
    n = 1
    i = 1
    with gzip.open(filename,"r") as f:
        lines = f.readlines()
        nevents = int(lines[0].split()[0])
        while True:
            nparticles = int(data[i-1,0])
            rv[n] = lines[i:i+nparticles]
            n += 1
            i += nparticles

            if i - 1 >= data.shape[0]:
                break

    assert nevents == len(rv)

    return rv

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

    parser = argparse.ArgumentParser("recreate MCPL files")
    parser.add_argument("filenames", nargs='+', help="atmospheric MC files")
    parser.add_argument("--mcpl-directory", required=True, help="directory where MCPL files are stored")
    parser.add_argument("-o", "--output", required=True, help="output directory")
    args = parser.parse_args()

    for filename in args.filenames:
        mcgn = get_mcgn(filename)
        print("Analyzing run %i" % mcgn.run[0])
        mcpl_filename = "mc_atm_genie_%06i_0.dat.gz" % mcgn.run[0]
        mcpl_file = join(args.mcpl_directory,mcpl_filename)
        mcpl = read_mcpl(mcpl_file)
        for ev in mcpl.keys():
            if ev in mcgn.evn.values:
                del mcpl[ev]
        output = "mc_atm_genie_%06i_0.dat" % mcgn.run[0]
        with open(join(args.output,output),"w") as fout:
            fout.write("%i %i\n" % (len(mcpl), 19))
            for ev, lines in mcpl.iteritems():
                fout.write(''.join(lines))
        print("Wrote new MCPL file with %i events" % len(mcpl))