import os,sys
from math import sqrt
import ROOT
import hddm_s
import itertools
dir_path = os.path.dirname(os.path.realpath(__file__))
ifs = hddm_s.istream(sys.argv[1])
ifs4 = hddm_s.istream(sys.argv[2])
abbrev = "tof"
name = "Time of Flight Detector - "
particle_list = ["Gamma","Positron","Electron","Neutrino","Muon+","Muon-","Pion0","Pion+","Pion-","Kaon0 Long","Kaon+","Kaon-","Neutron","Proton","Antiproton"]

#This function sets axis titles for an array of histograms
def AxisTitles(hists, x, y="Frequency", z=""):
        for hist in hists:
                hist.GetXaxis().SetTitle(x)
                hist.GetYaxis().SetTitle(y)
                if z != "":
                        hist.GetZaxis().SetTitle(z)
#This function sets an array of histograms' line color to red
def SetRed(hists):
        for hist in hists:
                hist.SetLineColor(2)

#Initializes the histograms to be filled by the data
t_list=[]
for num in range(0,2):
	t_list.append(ROOT.TH1D(abbrev+"_t" + str(num),"TOF Truth Hit Time (HDGeant-HDGeant4): End " + str(num),500,-10,10))
pt_list = []
pp_list = []
pp2d_list = []
for num in range(0,2):
	s = "True" if num else "False"
	pt_list.append(ROOT.TH1D(abbrev+"_pt" + str(num),"TOF Truth Point Time (HDGeant-HDGeant4): Primary=" + s,500,-10,10))
	pp_list.append(ROOT.TH1D(abbrev+"_pp" + str(num),name + "TOF Truth Point Momentum (HDGeant-HDGeant4): Primary=" + s,200,-0.1,0.1))
	pp2d_list.append(ROOT.TH2D(abbrev+"_pp2d" + str(num),name + "TOF Truth Point Momentum (HDGeant-HDGeant4): Primary=" + s,60,0,6,100,-0.1,0.1))
	
#Sets the axes titles for each of the histograms
AxisTitles(t_list, "#Delta Hit Time (HDGeant-HDGeant4) (ns)")
AxisTitles(pt_list, "#Delta Point Time (HDGeant-HDGeant4) (ns)")
AxisTitles(pp_list, "#Delta Momentum (HDGeant-HDGeant4) (GeV)")
AxisTitles(pp2d_list, "Momentum (HDGeant) (GeV)", "Momentum (HDGeant-HDGeant4) (GeV)")

#t_list[1].SetLineColor(3)

count = 0
for rec,rec4 in itertools.izip(ifs,ifs4):
        physicsEvent = rec.getPhysicsEvents()[0]
        physicsEvent4 = rec4.getPhysicsEvents()[0]
        hitViews = physicsEvent.getHitViews()
        hitViews4 = physicsEvent4.getHitViews()

        reaction = physicsEvent.getReaction()
        vtx_t = reaction.getVertex().getOrigin().t
        reaction4 = physicsEvent4.getReaction()
        vtx_t4 = reaction4.getVertex().getOrigin().t

        for hitView,hitView4 in itertools.izip(hitViews,hitViews4):
                systems = hitView.getForwardTOFs()
                systems4 = hitView4.getForwardTOFs()
                if len(systems) != 0 and len(systems4) != 0:
                        for system,system4 in itertools.izip(systems, systems4): 
                                pts = system.getFtofTruthPoints()
                                pts4 = system4.getFtofTruthPoints()
                                for pt in pts:
                                        for pt4 in pts4:
                                                if pt.ptype == pt4.ptype and pt.getTrackID().itrack == pt4.getTrackID().itrack:
                                                        num = 1 if pt.primary else 0
                                                        flight_time = pt.t - vtx_t
                                                        flight_time4 = pt4.t - vtx_t4
                                                        mom = sqrt(pt.px**2 + pt.py**2 + pt.pz**2)
                                                        mom4 = sqrt(pt4.px**2 + pt4.py**2 + pt4.pz**2) 
                                                        #pt_list[num].Fill(pt.t - pt4.t)
                                                        pt_list[num].Fill(flight_time - flight_time4)
                                                        pp_list[num].Fill( mom - mom4 )
                                                        pp2d_list[num].Fill( mom, mom - mom4 )
                                detectors = system.getFtofCounters()
                                detectors4 = system4.getFtofCounters()
                                if len(detectors) != 0 and len(detectors4) != 0:
                                        for detector,detector4 in itertools.izip(detectors, detectors4):
                                                hits = detector.getFtofTruthHits()
                                                hits4 = detector4.getFtofTruthHits()
                                                for hit in hits:
                                                        extra = hit.getFtofTruthExtra()
                                                        for hit4 in hits4:
                                                                extra4 = hit4.getFtofTruthExtra()
                                                                if hit.end == hit4.end and extra.ptype == extra4.ptype and extra.itrack == extra4.itrack:
                                                                        flight_time = pt.t - vtx_t
                                                                        flight_time4 = pt4.t - vtx_t4
                                                                        #t_list[hit.end].Fill(hit.t-hit4.t)
                                                                        t_list[hit.end].Fill(flight_time - flight_time4)
								
		
                count+=1
                if (count % 1000 == 0):
                        print count
print("\n")

#Create the output file to keep the completed histograms for data
outf = ROOT.TFile("tof_diff.root","recreate")

for h in t_list + pt_list + pp_list + pp2d_list:
	h.Write()

outf.Close()

print("Complete")