#!/usr/bin/env python
"""

This script can be used to produce an event-weight-sum response matrix, using AanetParser reader.
It creates a response for pure tracks.

"""
from docopt import docopt
#args = docopt(__doc__)

import ROOT
import datetime
import sys
import numpy as np
from input_files_orca6 import *

dataSets = {**neutrinos_v9, **muons_v9} 


#===========================================================================================
# initialise the response for pure tracks
#===========================================================================================

# Set binning parameters:

# true space 
ebins_t  ,emin_t  ,emax_t  = [40,1,10000]
ctbins_t ,ctmin_t ,ctmax_t = [80,-1,1]
bybins_t ,bymin_t ,bymax_t = [1,0,1]

ebins_t=[1.0, 1.2589254117941673, 1.5848931924611136, 1.9952623149688797, 2.51188643150958, 3.1622776601683795, 3.981071705534973, 5.011872336272724, 6.309573444801933, 7.943282347242816, 10.0, 12.589254117941675, 15.848931924611142, 19.952623149688797, 25.11886431509581, 31.622776601683793, 39.810717055349734, 50.11872336272725, 63.09573444801933, 79.43282347242817, 100.0, 125.89254117941675, 158.48931924611142, 199.52623149688807, 251.18864315095823, 316.22776601683796, 398.1071705534973, 501.18723362727246, 630.9573444801937, 794.3282347242822, 1000.0, 1258.9254117941675, 1584.893192461114, 1995.262314968881, 2511.886431509582, 3162.2776601683795, 3981.0717055349733, 5011.872336272725, 6309.573444801937, 7943.282347242822, 10000]
e_bins_t=np.array(ebins_t, dtype='d')
ctbins_t=[-1.0, -0.975, -0.95, -0.925, -0.9, -0.875, -0.85, -0.825, -0.8, -0.775, -0.75, -0.725, -0.7, -0.675, -0.65, -0.625, -0.6, -0.575, -0.55, -0.525, -0.5, -0.475, -0.44999999999999996, -0.42500000000000004, -0.4, -0.375, -0.35, -0.32499999999999996, -0.30000000000000004, -0.275, -0.25, -0.22499999999999998, -0.19999999999999996, -0.17500000000000004, -0.15000000000000002, -0.125, -0.09999999999999998, -0.07499999999999996, -0.050000000000000044, -0.025000000000000022, 0.0, 0.02499999999999991, 0.050000000000000044, 0.07499999999999996, 0.10000000000000009, 0.125, 0.1499999999999999, 0.17500000000000004, 0.19999999999999996, 0.2250000000000001, 0.25, 0.2749999999999999, 0.30000000000000004, 0.32499999999999996, 0.3500000000000001, 0.375, 0.3999999999999999, 0.42500000000000004, 0.44999999999999996, 0.4750000000000001, 0.5, 0.5249999999999999, 0.55, 0.575, 0.6000000000000001, 0.625, 0.6499999999999999, 0.675, 0.7, 0.7250000000000001, 0.75, 0.7749999999999999, 0.8, 0.825, 0.8500000000000001, 0.875, 0.8999999999999999, 0.925, 0.95, 0.9750000000000001, 1.0]
ct_bins_t=np.array(ctbins_t, dtype='d')
bybins_t=[0,1]
by_bins_t=np.array(bybins_t, dtype='d')

h_bins_true=ROOT.TH3D("h_bins_true", "hist true bins" , len(ebins_t)-1, e_bins_t, len(ctbins_t)-1, ct_bins_t, len(bybins_t)-1, by_bins_t)
# reco space 
ebins_r  ,emin_r  ,emax_r  = [19,2,1000]
ctbins_r ,ctmin_r ,ctmax_r = [10,-1,0]
bybins_r ,bymin_r ,bymax_r = [1,0,1]

ebins_r=[1.0, 2.0, 4.0, 4.61340151, 5.32086836, 6.13682553, 7.07791004, 8.1633102, 9.4151569, 10.85897475, 12.52420262, 14.4447938, 16.65990837, 19.21471159, 22.16129485, 25.55973776, 29.47933316, 34.0, 50, 100, 1000]
e_bins_r=np.array(ebins_r, dtype='d')
ctbins_r=[-1.0, -0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0.0]
ct_bins_r=np.array(ctbins_r, dtype='d')
bybins_r=[0,1]
by_bins_r=np.array(bybins_r, dtype='d')


h_bins_reco=ROOT.TH3D("h_bins_reco", "hist reco bins" , len(ebins_r)-1, e_bins_r, len(ctbins_r)-1, ct_bins_r, len(bybins_r)-1, by_bins_r)


# response for pure tracks
ewsResp = ROOT.EWSResponse(ROOT.EWSResponse.track, "ewsResponse_pureTracks", h_bins_true, h_bins_reco)

# The boolean containing the information if the event pass the cut(s) is contained in this single boolean branch
ewsResp.AddCut( 'SummaryEvent.Get_track_ql0',   '==', 1   , True ) 

print('NOTE: Building response for: pure tracks')

for key, value in dataSets.items():

    print ('Processing \"{}\" ...'.format(key))
    mainFile = ROOT.TFile(value['filename'],'READ')
    cutTFile  = ROOT.TFile(value['cutFile'],'READ')
    cutTTree  = cutTFile.Get('tree')

    aap = ROOT.AanetParser(mainFile, value['type'], True)
    mainTTree = aap.GetTree()
    mainTTree.AddFriend(cutTTree)

    mainTTree.Draw(">>totlist", "", "entrylist")
    eventList = ROOT.gDirectory.Get('totlist')

    nEvents = 0

    cutLeaf = mainTTree.FindLeaf('pure_tracks_cuts')
    start = datetime.datetime.now()

    for i in range(eventList.GetN()):
        if ( i % 250000 == 0): 
            rate = float(i) / (datetime.datetime.now() - start).total_seconds()
            if rate == 0 : rate = 1e10
            print("Event: {}/{} - {} remaining".format(i, eventList.GetN(), datetime.timedelta(seconds=(eventList.GetN()-i) / rate)))
        evt = aap.GetEvt(eventList.GetEntry(i))

        evt.Set_track_ql0(cutLeaf.GetValue()) 	
        if evt.Get_track_ql0() == 1 : nEvents += 1
        ewsResp.Fill(evt, eventList.GetEntry(i))

    print (nEvents)
ewsResp.PrintRunData()
print ('Done, save to disk ...')
ewsResp.WriteToFile('./responses/response_pure_tracks_ews_orca6v0bkg.root')