import ROOT as r
import numpy as np
import funcsByBarbara as tools

r.gROOT.SetBatch(r.kTRUE)
r.gROOT.ProcessLine(".X lhcbStyle.C")
plots = {}

def addPlot(key, x_axis, tree, string, cuts=None):
	if not cuts:
		tree.Draw(string+'>>h1')
	else:
		tree.Draw(string+'>>h1', cuts)
	plots[key] = r.gDirectory.Get('h1')
	plots[key].GetXaxis().SetTitle(x_axis)
	plots[key].SetTitle(key)
	plots[key].SetName(key)
	c1 = r.TCanvas(); c1.cd()
	plots[key].Draw()
	c1.Print('accPlots/%s.pdf'%key)
	c1.Print('accPlots/%s.png'%key)
	key += "-WEIGHTED"
	if 'bg' in key:
		if not cuts:
			tree.Draw(string+'>>h1', "NuWeight")
		else:
			cuts = "NuWeight*("+cuts+")"
			tree.Draw(string+'>>h1', cuts)
		plots[key] = r.gDirectory.Get('h1')
		plots[key].GetXaxis().SetTitle(x_axis)
		plots[key].SetTitle(key)
		plots[key].SetName(key)
		c1 = r.TCanvas(); c1.cd()
		plots[key].Draw()
		c1.Print('accPlots/%s.pdf'%key)
		c1.Print('accPlots/%s.png'%key)
	else:
		if not cuts:
			tree.Draw(string+'>>h1',"HNLw")
		else:
			cuts = "HNLw*("+cuts+")"
			tree.Draw(string+'>>h1', cuts)
		plots[key] = r.gDirectory.Get('h1')
		plots[key].GetXaxis().SetTitle(x_axis)
		plots[key].SetTitle(key)
		plots[key].SetName(key)
		c1 = r.TCanvas(); c1.cd()
		plots[key].Draw()
		c1.Print('accPlots/%s.pdf'%key)
		c1.Print('accPlots/%s.png'%key)




fpimu = r.TFile('../Geraldine/MUPI/ShipAna.root','read')
pimu = fpimu.Get('ShipAna')
#pimu_geo = tools.searchForNodes2_xyz_dict('../Geraldine/old_samples/MUPI/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008.root')
pimu_geo = tools.searchForNodes2_xyz_dict('../Geraldine/MUPI/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008-2985556.root')
fmumunu = r.TFile('../Geraldine/MUMUNU/ShipAna.root','read')
mumunu = fmumunu.Get('ShipAna')
#mumunu_geo = tools.searchForNodes2_xyz_dict('../Geraldine/old_samples/MUMUNU/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008.root')
mumunu_geo = tools.searchForNodes2_xyz_dict('../Geraldine/MUMUNU/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008-TOTAL.root')
fbg = r.TFile('../Geraldine/KLONG/ShipAna.root','read')
bg = fbg.Get('ShipAna')
#bg_geo = tools.searchForNodes2_xyz_dict('../Geraldine/KLONG/geofile_full.10.0.Genie-TGeant4.root')
bg_geo = tools.searchForNodes2_xyz_dict('../Geraldine/KLONG/geofile_full.10.0.Genie-TGeant4_9files.root')

studies = {
	'pimu': {'data':pimu, 'geo':pimu_geo, 'ntot':1000},
	'mumunu': {'data':mumunu, 'geo':mumunu_geo, 'ntot':7000},
	'bg': {'data':bg, 'geo':bg_geo, 'ntot':99999}
}

# Cuts:
# - z_vertex in [vol[0]+5m, straw[0]]
# - has track converged
# - chi2/ndof < 10
# - DOCA < 10 cm
# - has no veto?

def fit_converged(event):
	flags = event.DaughtersFitConverged
	if np.product(flags):
		return True
	return False

def chi2_cut(event, chi2max):
	chi2 = event.DaughtersChi2
	n = event.DaughtersNPoints
	for i in xrange(len(n)):
		if chi2[i]/n[i] > chi2max:
			return False
	return True

def doca_cut(event, doca_max):
	if event.DOCA < doca_max:
		return True
	return False

def z_cut(event, study):
	z = event.vtxz
	minz = studies[study]['geo']['lidT1I_1']['z']['pos'] + studies[study]['geo']['lidT1I_1']['z']['dim'] + 500.
	maxz = studies[study]['geo']['lidT6I_1']['z']['pos'] - studies[study]['geo']['lidT6I_1']['z']['dim']
	if minz < z < maxz:
		return True
	return False

def ip_cut(event, maxip):
	if event.IP0 < maxip:
		return True
	return False

def nothing_in_veto5(event):
	if len(event.veto5_x) > 0:
		return False
	return True

def nothing_in_liquidscint(event):
	if len(event.liquidscint_x) > 0:
		return False
	return True

def weight(event, study):
	w = event.HNLw / studies[s]['ntot']
	if 'bg' in study:
		w = event.NuWeight
	return w

results = {}
for s in studies:
	print 'Reading %s...'%s
	t = studies[s]['data']
	ntot = t.GetEntriesFast()
	nAfterCuts = [0]*10
	nAfterCuts_weighted = [0.]*10
	for event in t:
		w = weight(event,s)
		if fit_converged(event): 
			nAfterCuts[0] += 1
			nAfterCuts_weighted[0] += w
			if chi2_cut(event,25):
				nAfterCuts[1] +=1
				nAfterCuts_weighted[1] +=w
				if chi2_cut(event,15):
					nAfterCuts[2]+=1
					nAfterCuts_weighted[2]+=w
					if doca_cut(event, 50.):
						nAfterCuts[2+1] += 1
						nAfterCuts_weighted[2+1] += w
						if doca_cut(event, 30.):
							nAfterCuts[3+1]+=1
							nAfterCuts_weighted[3+1]+=w
							if z_cut(event, s):
								nAfterCuts[4+1] += 1
								nAfterCuts_weighted[4+1] += w
								if ip_cut(event,500.):
									nAfterCuts[5+1]+=1
									nAfterCuts_weighted[5+1]+=w
									if ip_cut(event,250.):
										nAfterCuts[6+1]+=1
										nAfterCuts_weighted[6+1]+=w
										if nothing_in_veto5(event):
											nAfterCuts[7+1] += 1
											nAfterCuts_weighted[7+1] += w
											if nothing_in_liquidscint(event):
												nAfterCuts[8+1] += 1
												nAfterCuts_weighted[8+1] += w
	print '%s \t survived particles: '%s, nAfterCuts, ' of %s total'%ntot
	print '%s \t weighted survived particles: '%s, nAfterCuts_weighted, ' of %s total'%ntot
	results[s] = {}
	results[s]['ntot']=studies[s]['ntot']
	results[s]['numbers']=nAfterCuts
	results[s]['numbersWithWeights']=nAfterCuts_weighted
	addPlot(s+'-0cuts-IP', 'Impact parameter to target [cm]', t, 'IP0')
	addPlot(s+'-0cuts-Mass', 'Reco inv. mass', t, 'Mass'              )
	addPlot(s+'-0cuts-dPt', 'Tracks Pt', t, 'DaughtersPt'             )
	addPlot(s+'-0cuts-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' )
	cut = "DaughtersChi2<25*DaughtersNPoints"
	addPlot(s+'-chi2cut-IP', 'Impact parameter to target [cm]', t, 'IP0', cut)
	addPlot(s+'-chi2cut-Mass', 'Reco inv. mass', t, 'Mass'              , cut)
	addPlot(s+'-chi2cut-dPt', 'Tracks Pt', t, 'DaughtersPt'             , cut)
	addPlot(s+'-chi2cut-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' , cut)
	cut += " && DOCA < 25"
	addPlot(s+'-chi2andDOCAcuts-IP', 'Impact parameter to target [cm]', t, 'IP0', cut)
	addPlot(s+'-chi2andDOCAcuts-Mass', 'Reco inv. mass', t, 'Mass'              , cut)
	addPlot(s+'-chi2andDOCAcuts-dPt', 'Tracks Pt', t, 'DaughtersPt'             , cut)
	addPlot(s+'-chi2andDOCAcuts-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' , cut)
	cut += " && vtxz > %s && vtxz < %s"%(studies[s]['geo']['lidT1I_1']['z']['pos'] + studies[s]['geo']['lidT1I_1']['z']['dim'] + 500.,
		studies[s]['geo']['lidT6I_1']['z']['pos'] - studies[s]['geo']['lidT6I_1']['z']['dim'])
	addPlot(s+'-chi2andDOCAandZcuts-IP', 'Impact parameter to target [cm]', t, 'IP0', cut)
	addPlot(s+'-chi2andDOCAandZcuts-Mass', 'Reco inv. mass', t, 'Mass'              , cut)
	addPlot(s+'-chi2andDOCAandZcuts-dPt', 'Tracks Pt', t, 'DaughtersPt'             , cut)
	addPlot(s+'-chi2andDOCAandZcuts-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' , cut)


for s in ['pimu', 'mumunu', 'bg']:
    for list in results[s]:
        print s,list,results[s][list]