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

r.gROOT.SetBatch(r.kTRUE)
r.gROOT.ProcessLine(".L 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)


fpimu = r.TFile('../Geraldine/MUPI/ShipAna.root','read')
pimu = fpimu.Get('ShipAna')
pimu_geo = tools.searchForNodes2_xyz_dict('../Geraldine/MUPI/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008.root')
fmumunu = r.TFile('../Geraldine/MUMUNU/ShipAna.root','read')
mumunu = fmumunu.Get('ShipAna')
mumunu_geo = tools.searchForNodes2_xyz_dict('../Geraldine/MUMUNU/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008.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')

studies = {
	'pimu': {'data':pimu, 'geo':pimu_geo},
	'mumunu': {'data':mumunu, 'geo':mumunu_geo},
	'bg': {'data':bg, 'geo':bg_geo}
}

# 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']
	print minz, maxz
	if minz < z < maxz:
		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

for s in studies:
	print 'Reading %s...'%s
	t = studies[s]['data']
	ntot = t.GetEntriesFast()
	nAfterCuts = [0]*6
	for event in t:
		if fit_converged(event): 
			nAfterCuts[0] += 1
			if chi2_cut(event,25):
				nAfterCuts[1] +=1
				if doca_cut(event, 50.):
					nAfterCuts[2] += 1
					if z_cut(event, s):
						nAfterCuts[3] += 1
						if nothing_in_veto5(event):
							nAfterCuts[4] += 1
							if nothing_in_liquidscint(event):
								nAfterCuts[5] += 1
	print '%s \t survived particles: '%s, nAfterCuts, ' of %s total'%ntot
	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 < 50"
	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)