import os,sys
import ROOT
import hddm_s
import itertools
#ifs = hddm_s.istream("g3/hdgeant.hddm")
#ifs4 = hddm_s.istream("g4/hdgeant.hddm")
ifs = hddm_s.istream(sys.argv[1])
ifs4 = hddm_s.istream(sys.argv[2])
abbrev = "cdc"
name = "Central Drift Chamber - "

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)
def SetRed(hists):
        for hist in hists:
                hist.SetLineColor(2)

straw_offset = [0,0,42,84,138,192,258,324,404,484,577,670,776,882,1005,1128,1263,1398,1544,1690,1848,2006,2176,2346,2528,2710,2907,3104,3313]

#Initializes the histograms to be filled by the data
q_hist = ROOT.TH1D(abbrev+"_q",name + "Charge (HDGeant)",200,0,20000)
q_hist4= ROOT.TH1D(abbrev+"_q4",name + "Charge (HDGeant4)",200,0,20000)
qstraw_hist = ROOT.TH2D(abbrev+"_qnum",name + "Charge vs. Straw (HDGeant)",200,0,20000,3500,1,3500)
qstraw_hist4=ROOT.TH2D(abbrev+"_qnum4",name + "Charge vs. Straw (HDGeant4)",200,0,20000,3500,1,3500)

#Sets the axes titles for each of the histograms
AxisTitles([q_hist, q_hist4], "Charge (C)")
AxisTitles([qstraw_hist, qstraw_hist4], "Charge (C)", "Straw number")

#Sets the color of the HDGeant4 histograms to red to distinguish
# from the HDGeant histograms
SetRed([q_hist4])

count = 0
for rec,rec4 in itertools.izip(ifs,ifs4):
        physicsEvent = rec.getPhysicsEvents()[0]
        physicsEvent4 = rec4.getPhysicsEvents()[0]
        hitViews = physicsEvent.getHitViews()
        hitViews4 = physicsEvent4.getHitViews()
        for hitView,hitView4 in itertools.izip(hitViews,hitViews4):
                systems = hitView.getCentralDCs()
                systems4 = hitView4.getCentralDCs()
                if len(systems) != 0:
                        for system in systems:
                                detectors = system.getCdcStraws()
                                if len(detectors) != 0:
                                        for detector in detectors:
                                                hits = detector.getCdcStrawTruthHits()
                                                for hit in hits:
                                                        straw_num = straw_offset[detector.ring] + detector.straw
                                                        q_hist.Fill(hit.q)
                                                        qstraw_hist.Fill(hit.q,straw_num)
                if len(systems4) != 0:
                        for system4 in systems4:
                                detectors = system4.getCdcStraws()
                                if len(detectors) != 0:
                                        for detector in detectors:
                                                hits = detector.getCdcStrawTruthHits()
                                                for hit in hits:
                                                        straw_num = straw_offset[detector.ring] + detector.straw
                                                        q_hist4.Fill(hit.q)
                                                        qstraw_hist4.Fill(hit.q,straw_num)

        count+=1
        if (count % 1000 == 0):
                print count
        
print("\n")

#Create the output file to keep the completed histograms for CDC
outf = ROOT.TFile("cdc.root","recreate")
q_hist.Write()
q_hist4.Write()
qstraw_hist.Write()
qstraw_hist4.Write()
outf.Close()