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()
	if "IP" in key or "Mass" in key or "DOCA" in key or "dChi2" in key:
		c1.SetLogy()
	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()
		if "IP" in key or "Mass" in key or "DOCA" in key or "dChi2" in key:
			c1.SetLogy()
		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()
		if "IP" in key or "Mass" in key or "DOCA" in key or "dChi2" in key:
			c1.SetLogy()
		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')
#fbg = r.TFile('../Geraldine/KLONG_smallfile/ShipAna.root','read')
bg = fbg.Get('ShipAna')
#bg_geo = tools.searchForNodes2_xyz_dict('../Geraldine/KLONG_smallfile/geofile_full.10.0.Genie-TGeant4_9files.root')
#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_70cc30nc.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':142856}#99999}#199998
}

# 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[study]['ntot']
	if 'bg' in study:
		w = event.NuWeight
	return w

def computeResidualWeight(filename):
	if not filename:
		filename = "forElena_Ethr.txt"
	fScan = file(filename)
	lines = fScan.readlines()
	
	lines = lines[3:-1]
	events = []
	for line in lines:
	    events.append(int(line.rsplit("*")[1].replace(" ","")))
	
	count = 0
	weighted_count = 0.
	print "unvetoed: ", len(events)
	#for entry in events:
	#    bg.GetEntry(entry)
	#    particles = bg.Particles
	#    for p in particles:
	#        count+=1
	#        weighted_count += bg.NuWeight	
	for event in bg:
		if event.EventNumber in events:
			if fit_converged(event) and chi2_cut(event,25) and doca_cut(event, 50.) and ip_cut(event, 500.):
				#print event.EventNumber, 'mothers: ', event.DaughtersTruthMotherPDG[0], event.DaughtersTruthMotherPDG[1]
				print event.EventNumber
				print 'mothers: ', event.DaughtersTruthMotherPDG[0], event.DaughtersTruthMotherPDG[1]
				#print event.DOCA, event.IP0, event.DaughtersChi2[0]/event.DaughtersNPoints[0], event.DaughtersChi2[1]/event.DaughtersNPoints[1], event.DaughtersTruthPDG[0], event.DaughtersTruthPDG[1] 
				print "tracks: ", event.DaughtersTruthPDG[0], event.DaughtersTruthPDG[1] 
				print tools.interactionElement(event.DaughtersTruthProdX[0], event.DaughtersTruthProdY[0], event.DaughtersTruthProdZ[0], bg_geo), tools.interactionElement(event.DaughtersTruthProdX[1], event.DaughtersTruthProdY[1], event.DaughtersTruthProdZ[1], bg_geo)
				count +=1
				weighted_count += event.NuWeight
	return count, weighted_count


def do_my_studies():
	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-P','Reconstructed Momentum', t, 'P')
                addPlot(s+'-0cuts-Pt', 'Reconstructed Pt', t, 'Pt')
		addPlot(s+'-0cuts-dMinP', 'Min Tracks P', t, 'DaughterMinP'             )
		addPlot(s+'-0cuts-dMinPt', 'Min Tracks Pt', t, 'DaughterMinPt'             )
		addPlot(s+'-0cuts-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' )
		addPlot(s+'-0cuts-DOCA', 'Daughter tracks DOCA [cm]', t, 'DOCA' )
		addPlot(s+'-0cuts-dChi2', 'Daughter tracks #chi^{2}/n', t, 'DaughtersChi2/DaughtersNPoints' )
		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-P', 'Reconstructed Momentum', t, 'P', cut)
                addPlot(s+'-chi2cut-Pt','Reconstructed Pt', t, 'Pt', cut)
		addPlot(s+'-chi2cut-dMinP', 'Min Tracks P', t, 'DaughterMinP'        , cut     )
		addPlot(s+'-chi2cut-dMinPt', 'Min Tracks Pt', t, 'DaughterMinPt'           , cut   )
		addPlot(s+'-chi2cut-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' , cut)
		addPlot(s+'-chi2cut-DOCA', 'Daughter tracks DOCA [cm]', t, 'DOCA'  , 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-P','Reconstructed Momentum', t, 'P', cut)
                addPlot(s+'-chi2andDOCAcuts-Pt','Reconstructed Pt', t, 'Pt', cut)
		addPlot(s+'-chi2andDOCAcuts-dPt', 'Tracks Pt', t, 'DaughtersPt'             , cut)
		addPlot(s+'-chi2andDOCAcuts-dMinP', 'Min Tracks P', t, 'DaughterMinP'        , cut      )
		addPlot(s+'-chi2andDOCAcuts-dMinPt', 'Min Tracks Pt', t, 'DaughterMinPt'     , 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-P','Reconstructed Momentum', t, 'P', cut)
                addPlot(s+'-chi2andDOCAandZcuts-Pt','Reconstructed Pt', t, 'Pt', cut)
		addPlot(s+'-chi2andDOCAandZcuts-dPt', 'Tracks Pt', t, 'DaughtersPt'             , cut)
		addPlot(s+'-chi2andDOCAandZcuts-dMinP', 'Min Tracks P', t, 'DaughterMinP'         , cut     )
		addPlot(s+'-chi2andDOCAandZcuts-dMinPt', 'Min Tracks Pt', t, 'DaughterMinPt'      , 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]