import ROOT as r
import funcsByBarbara as tools
cSaver = []

def getName(code):
	try:
		mother = pdg.GetParticle( code ).GetName()
	except:
		mother = 'beam'
		if code == 9900015:
			mother = 9900015
	return mother

def numOfFakeCandidates(anaFile = '../DATA/NewPIMU/ShipAna.root'):
	f = r.TFile(anaFile, 'read')
	t = f.Get('ShipAna')
	numTot = 0
	numFake = 0
	for HNL in t:
		mums = HNL.DaughtersTruthMotherPDG
		for mum in mums:
			if mum != 9900015:
				numFake += 1
				break
		numTot +=1
	cSaver.append(r.TCanvas())
	t.Draw("vtxy:vtxx","DaughtersTruthMotherPDG != 9900015 && DOCA < 30. && IP0 < 250. && vtxz > -1958. && vtxz < 2578.")
	cSaver.append(r.TCanvas())
	t.Draw("vtxz","DaughtersTruthMotherPDG != 9900015 && DOCA < 30. && IP0 < 250. && vtxz > -1958. && vtxz < 2578.")
	cSaver.append(r.TCanvas())
	t.Draw("veto5_y:veto5_x","DaughtersTruthMotherPDG != 9900015 && DOCA < 30. && IP0 < 250. && vtxz > -1958. && vtxz < 2578.")
	#t.GetEntry(811)
	#print int(t.EventNumber)
	#print len(t.straw_x)
	#print len(t.veto5_x)
	#print
	#t.GetEntry(1298)
	#print int(t.EventNumber)
	#print len(t.straw_x)
	#print len(t.veto5_x)
	#print
	#t.GetEntry(2)
	#print int(t.EventNumber)
	#print len(t.straw_x)
	#print len(t.veto5_x)
	#print
	cutVeto = "!(Alt$(veto5_z,0))"
	cSaver.append(r.TCanvas())
	t.Draw("vtxz","DaughtersTruthMotherPDG != 9900015 && DOCA < 30. && IP0 < 250. && vtxz > -1958. && vtxz < 2578. && %s"%cutVeto)
	cutLS = "!(Alt$(liquidscint_z,0))"
	cSaver.append(r.TCanvas())
	t.Draw("vtxz","DaughtersTruthMotherPDG != 9900015 && DOCA < 30. && IP0 < 250. && vtxz > -1958. && vtxz < 2578. && %s && %s"%(cutVeto, cutLS))
	cutEllipse = "( (vtxx*vtxx)/(250.*250.) + (vtxy*vtxy)/(500.*500.) ) < 1."
	cutVtx = "!(BadTruthVtx)"
	badMum = "DaughtersTruthMotherPDG[][0] != 9900015 && DaughtersTruthMotherPDG[][1] != 9900015"
	goodMum = "DaughtersTruthMotherPDG[][0] == 9900015 && DaughtersTruthMotherPDG[][1] == 9900015"
	docaVtxIp = "DOCA < 30. && IP0 < 250. && vtxz > -1958. && vtxz < 2578."
	trackQ = "(DaughtersChi2/DaughtersNPoints)<5. && DaughtersAlwaysIn"
	#t.Scan("EventNumber","%s && %s && %s && %s && %s && %s && %s"%(badMum, docaVtxIp, cutVeto, cutLS, cutEllipse, cutVtx, trackQ))
	#t.Scan("EventNumber","%s && %s && %s && %s && %s && %s && %s"%(goodMum, docaVtxIp, cutVeto, cutLS, cutEllipse, cutVtx, trackQ))
	t.Draw(">>elistAll","%s && %s && %s && %s && %s && %s"%(docaVtxIp, cutVeto, cutLS, cutEllipse, cutVtx, trackQ))
	elistAll = r.gDirectory.Get("elistAll")
	print "Total candidates after the cuts: ", elistAll.GetN()
	t.Draw(">>elistBad","%s && %s && %s && %s && %s && %s && %s"%(badMum, docaVtxIp, cutVeto, cutLS, cutEllipse, cutVtx, trackQ))
	elistBad = r.gDirectory.Get("elistBad")
	print "Bad candidates after the cuts: ", elistBad.GetN()
	t.Draw(">>elistGood","%s && %s && %s && %s && %s && %s && %s"%(goodMum, docaVtxIp, cutVeto, cutLS, cutEllipse, cutVtx, trackQ))
	elistGood = r.gDirectory.Get("elistGood")
	print "Good candidates after the cuts: ", elistGood.GetN()
	return numFake, numTot


def lookAtSingleEvent(evNumber=5771, inputfile='../DATA/NewPIMU/ship.10.0.Pythia8-TGeant4-1.0GeV-008_rec.root'):
	#inputfile = '../DATA/NewPIMU/ship.10.0.Pythia8-TGeant4-1.0GeV-008_rec.root'
	f = r.TFile(inputfile,'read')
	t = f.Get('cbmsim')
	tGeom = r.TGeoManager("Geometry", "Geane geometry")
	gMan  = tGeom.Import(inputfile.replace('ship','geofile_full').replace('_rec',''))
	fGeo = r.gGeoManager
	pdg = r.TDatabasePDG().Instance()
	t.GetEntry(evNumber)
	nCandidates = 0
	for p in t.Particles:
		nCandidates += 1
	print
	print '\t %s candidates in this event.'%nCandidates
	print
	print 'Tracks in this event: name, fStartZ, node, mother'
	print
	for tr in t.MCTrack:
		name = getName(tr.GetPdgCode())
		try:
			mother = getName(t.MCTrack[tr.GetMotherId()].GetPdgCode())
		except:
			mother = 'beam'
		print '%s \t %.2f \t %s \t %s'%(name, tr.GetStartZ(), fGeo.FindNode(tr.GetStartX(),tr.GetStartY(),tr.GetStartZ()).GetName(), mother)
	print
	print 'Candidates in this event: daugther name, daughter node, daugther mother (x2)'
	print
	for p in t.Particles:
		t1,t2 = p.GetDaughter(0),p.GetDaughter(1)
		pdg1, mumPdg1, truthX1, truthY1, truthZ1 = tools.retrieveMCParticleInfo(t, t1)
		pdg2, mumPdg2, truthX2, truthY2, truthZ2 = tools.retrieveMCParticleInfo(t, t1)
		name1 = getName(pdg1)
		name2 = getName(pdg2)
		mother1 = getName(mumPdg1)
		mother2 = getName(mumPdg2)
		node1 = fGeo.FindNode(truthX1, truthY1, truthZ1).GetName()
		node2 = fGeo.FindNode(truthX2, truthY2, truthZ2).GetName()
		print '%s\t%s\t%s\t%s\t\t%s\t%s\t%s\t%s'%( truthZ1, name1, node1, mother1, truthZ2, name2, node2, mother2  )