#!/usr/bin/env python

"""
 Usage deweight.py -f input_file -o output_file

 Simple example to show how to 'deweight' the events in a genhen or gseagen
 file, producing a realistic sample of events (i.e. events which all 
 have the same weight).

 The desired livetime of the output sample can be set with the -l option. If this
 is set to 0 (default), the maximum livetime possible for the input file will be used.

 options:
 -f : input file                                default=../data/mc5.1.numuCC.nohits.aa.root
 -o : output file                               default=deweighted.root
 -F : neutrino flux (Gev-1 m-2 s-1 sr-1)        default=1e-4*neutrino().E**-2
 -l : target livetime (yr) (0=maximum possible) default=0
 -h : this help message and exit                default=False       
 -i : Python interative mode (prompt when done) default=False  
"""

import aa, ROOT,sys

options = aa.Options( __doc__, sys.argv[1:] ) 


f = ROOT.EventFile(options.f)
E = f.roottree()

flux   = options.F
weight = "w[1]/" +str(f.header.ngen()) +"*" + flux

# get the total rate and the max. weight.

E.Draw("0>>hh",weight,"goff")
nevts     = ROOT.hh.Integral();
maxweight = max( E.GetW()[i] for i in range( E.GetEntries() ) )
maxlive   = 1.0 / maxweight


livetime = options.l

if options.l == 0 :
	print ("setting livetime of output sample to maximum: ", maxlive )
	livetime = maxlive

if options.l > maxlive :
	print ("input file does not have enough statistics for requested livetime; max=",maxlive)
	sys.exit()

print ("flux                                      : ", flux, "GeV-1 m-2 sr-2 s-1")
print ("mean number of events per year            : ", nevts)
print ("max. weight                               : ", maxweight )
print ("deweighted output will have livetime  (yr): ", livetime )

# randomly select the events

E.Draw(">>elist","("+weight + "/ "+str(maxweight * maxlive/livetime)+" ) > rndm" )
elist = ROOT.elist


# write the output file

fout = ROOT.TFile(options.o, "RECREATE")
E.SetEventList(elist)
Enew = E.CopyTree("")

print ("selected", Enew.GetEntries(),"de-weighted entries")

Enew.Write()
f.header.set_line("livetime", str(livetime*3600*24*365) ) 
# todo? : remove the original ngen, set w[1]=1,0,1/livetime ?

f.header.Write()
fout.Close()