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_hist = ROOT.TH1D(abbrev+"_t","TOF Hit Time (HDGeant)",250,-50,200)
t_hist4=ROOT.TH1D(abbrev+"_t4","TOF Hit Time (HDGeant4)",250,-50,200)
tnum_hist=ROOT.TH2D(abbrev+"_tnum","TOF Hit Time vs. Bar (HDGeant)",150,0,150,88,1,89)
tnum_hist4=ROOT.TH2D(abbrev+"_tnum4","TOF Hit Time vs. Bar (HDGeant4)",150,0,150,88,1,89)
t_tof_hist = ROOT.TH1D(abbrev+"_tof_t","TOF Hit Time - Primary Vertex time (HDGeant)",250,-50,200)
t_tof_hist4=ROOT.TH1D(abbrev+"_tof_t4","TOF Hit Time - Primary Vertex time (HDGeant4)",250,-50,200)
tnum_tof_hist=ROOT.TH2D(abbrev+"_tof_tnum","TOF Hit - Primary Vertex time Time vs. Bar (HDGeant)",150,0,150,88,1,89)
tnum_tof_hist4=ROOT.TH2D(abbrev+"_tof_tnum4","TOF Hit Time - Primary Vertex time vs. Bar (HDGeant4)",150,0,150,88,1,89)

#Sets the axes titles for each of the histograms
AxisTitles([t_hist, t_hist4], "Hit Time (ns)")
AxisTitles([tnum_hist, tnum_hist4], "Hit Time (ns)", "(Plane Number)*44 + Bar Number")
AxisTitles([t_tof_hist, t_tof_hist4], "Hit Time - Primary Vertex time (ns)")
AxisTitles([tnum_tof_hist, tnum_tof_hist4], "Hit Time - Primary Vertex time (ns)", "(Plane Number)*44 + Bar Number")

print("Creating the histograms...")

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:
                        for system in systems: 
                                detectors = system.getFtofCounters()
                                if len(detectors) != 0:
                                        for detector in detectors:
                                                hits = detector.getFtofTruthHits()
                                                hit_count = len(hits)
                                                for hit in hits:
                                                        t_hist.Fill(hit.t)
                                                        tnum_hist.Fill(hit.t,detector.plane*44 + detector.bar)
                                                        t_tof_hist.Fill(hit.t-vtx_t)
                                                        tnum_tof_hist.Fill(hit.t-vtx_t,detector.plane*44 + detector.bar)
                if len(systems4) != 0:
                        for system4 in systems4:
                                detectors = system4.getFtofCounters()
                                if len(detectors) != 0:
                                        for detector in detectors:
                                                hits = detector.getFtofTruthHits()
                                                hit_count4 = len(hits)
                                                for hit in hits:
                                                        t_hist4.Fill(hit.t)
                                                        tnum_hist4.Fill(hit.t,detector.plane*44 + detector.bar)
                                                        t_tof_hist4.Fill(hit.t-vtx_t4)
                                                        tnum_tof_hist4.Fill(hit.t-vtx_t4,detector.plane*44 + detector.bar)
		
        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.root","recreate")
t_hist.Write()
t_hist4.Write()
tnum_hist.Write()
tnum_hist4.Write()
t_tof_hist.Write()
t_tof_hist4.Write()
tnum_tof_hist.Write()
tnum_tof_hist4.Write()

outf.Close()