import ROOT, glob, math, sys, os
from ROOT import std
from array import array

options = ["-f","-of","-clname","-if","-trname","-ml","-mh","-v","-gf","-resname"]
opts = set(options)
searchLabel = "reco"
mom_low = 0.
mom_high = 1000.

def getArgsToNextOpt(stringlist, index) :
    id = index+1
    args = []
    while(id < len(stringlist)):
        if stringlist[id] in options:
            return args
        args.append(stringlist[id])
        id+=1
    return args    

def getFileNumber(rf, label) :
    splits = rf.split(".")
    for i in splits:
        if i.count(label) :
            return int(i[i.find(label)+len(label):]) 
        
def compareFiles(f1, f2) :
    if1 = getFileNumber(f1,searchLabel)
    if2 = getFileNumber(f2,searchLabel)
    return if1-if2
    
  
bernhard = 0
verbose = 0
specOut = 0
batch = 0
genfit = 0

#argument parsing:

clname = "TpcCluster"
trname = "CdcTrack"
evtidname = "TpcEventIdentifier"
resname = "TpcRefTrackResidual" 
ofile = ""

for iarg in range(len(sys.argv)) :
    arg = sys.argv[iarg]
    if arg == "-f":   #distortion reco output file
        files = getArgsToNextOpt(sys.argv,iarg)
    if arg == "-if":   #input reco files containing the clusters and tracks    
        ifiles = getArgsToNextOpt(sys.argv,iarg)
    if arg == "-of":
        specOut = 1
        ofile = getArgsToNextOpt(sys.argv,iarg)
    if arg == "-clname":
        clname =  getArgsToNextOpt(sys.argv,iarg)[0]
    if arg == "-trname":
        trname =  getArgsToNextOpt(sys.argv,iarg)[0]
    if arg == "-resname":
        resname =  getArgsToNextOpt(sys.argv,iarg)[0]
    if arg == "-ml":
        mom_low = float(getArgsToNextOpt(sys.argv,iarg)[0])    
    if arg == "-mh":
        mom_high = float(getArgsToNextOpt(sys.argv,iarg)[0])
    if arg == "-v":
        verbose = 1 
    if arg == "-gf":
        genfit = 1     
        
if specOut == 0 :
    ofile = files[0]
    name = "_anaOut_ml"+str(mom_low)+"_mh"+str(mom_high)+".root"
    print name
    ofile = ofile.replace(".root",name)

print "Analysis input file: ",files 
print "Reconstruction Input files: "
ifiles = sorted(ifiles,cmp=compareFiles)
for f in ifiles:
    print "   ", f
print "Analysis output file:", ofile
print "Cluster branch name:", clname
print "CdcTrack branch name:", trname
print "Low CDC momentum cut:", mom_low
print ""
print ""

if len(files) < 1:
    print "No input files given!" 
    exit


ROOT.gROOT.ProcessLine(".x rootlogon.C") 

if bernhard :
    ROOT.gROOT.ProcessLine(".L rootlogon_Bernhard.C")
    ROOT.gROOT.ProcessLine("rootlogon_Bernhard()")
    ROOT.gROOT.ProcessLine('gROOT->SetStyle("col")') 
    ROOT.gROOT.ProcessLine("gROOT->ForceStyle()") 
else :
    ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")') 

ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") 

#construct input file chain:
ichain = ROOT.TChain("cbmsim")
for p in ifiles :
    ichain.Add(p)
itree = ichain
itree.SetBranchStatus("*", 0)
#itree.SetBranchStatus(clname+".*", 1)  #WHY IS THIS NOT WORKING RIGHT NOW ???
itree.SetBranchStatus(trname+".*", 1)
itree.SetBranchStatus("TpcEventIdentifier.*", 1)

print "Input-chain has", ichain.GetEntries(), " entries"
print
print "Building input chain index ... "
major = evtidname+ ".getSpill()"
minor = evtidname+ ".getEventInSpill()"
nE = ichain.BuildIndex(major,minor)
if nE < 1: #something went wrong
    print " ... could not build index - aborting"
    exit
print " ... Built chain index with ",nE," entries"
chainIndex = ichain.GetTreeIndex()

tc_cl = ROOT.TClonesArray("TpcCluster")
if genfit:
    tc_tr = ROOT.TClonesArray("GFTrack")
else:
    tc_tr = ROOT.TClonesArray("CdcTrack")

#itree.SetBranchAddress(clname, tc_cl)
itree.SetBranchAddress(trname, tc_tr)

histos = []

# ----------------------------- SET UP HISTOS -------------------------------------    
distVsR = ROOT.TH2D("distVsR","Cluster-CDCHit Residual vs Cluster R",
                    100,5,15,300,-5,5)
distVsR.GetXaxis().SetTitle("Radius (cm)")
histos.append(distVsR)
dPhiVsR = ROOT.TH2D("dPhiVsR","Cluster-CDCHit Residual (Phi) vs Cluster R",
                    100,5,15,300,-5,5)
dPhiVsR.GetXaxis().SetTitle("Radius (cm)")
histos.append(dPhiVsR)
distHist = ROOT.TH1D("distHist", "Residual Distribution",100,-5,5)
distHist.GetXaxis().SetTitle("Radius (cm)")
histos.append(distHist)
clPosR = ROOT.TH1D("clPosR", "Cluster Radial Distribution",100,5,15)
clPosR.GetXaxis().SetTitle("Radius (cm)")
histos.append(clPosR)
clPosZ = ROOT.TH1D("clPosZ", "Cluster Z Distribution",300,-50,100)
clPosZ.GetXaxis().SetTitle("Z (cm)")
histos.append(clPosZ)
cdcMom = ROOT.TH1D("cdcMom", "Momentum from CDC fit", 1000,0,5)
histos.append(cdcMom)
cdcNHits = ROOT.TH1D("cdcnHits", "Number of Hits in CDC tracks", 100,0,100)
histos.append(cdcNHits)
cdcVal = ROOT.TH1D("cdcVal", "Validity of CDC tracks", 10,-5,5)
histos.append(cdcVal)


# -------------------------- DEFINE zCUTS ----------------------------------------
zCuts = [10.,20.,30.,40.,50.,60.,70.]
zOffset = -25.
dictdPhiVsR = {}

# ------------------------- ANALYSIS CODE ----------------------------------------

for f in files:
    print "opening ", f, " ...." 
    rfile = ROOT.TFile.Open(f, "read")
    if rfile.IsZombie():
        print "file not ok, skipping ..."
        continue

    atree = rfile.Get("cbmsim")
    atree.SetBranchStatus("*", 0)
    resbr = resname+".*"
    atree.SetBranchStatus(resbr, 1)
    atree.SetBranchStatus("TpcEventIdentifierOut.*", 1)
    
    counter = 0
    for e in atree :
        if counter % 1000 == 0 :
            sys.stdout.write(". ")
            sys.stdout.flush()
        evtid = e.TpcEventIdentifierOut.At(0)
        evt = evtid.getEventInSpill()
        spill = evtid.getSpill()
        if verbose:
            print "Spill: ",spill,"  Event: ",evt
        #try to find matching event in input chain
        matchId = chainIndex.GetEntryNumberWithIndex(spill,evt)
        if verbose:
            print "... matches entry ",matchId," of input chain"
        if matchId < 0:
            print "Could not associate with an event from the input chain - skipping event"
            continue
        itree.GetEntry(matchId)
        if verbose:
            print "Input tree has ", tc_cl.GetEntriesFast()," Clusters in event ", matchId
        for rc in e.TpcRefTrackResidualAligned:
            if verbose:
                print "   ...looping over ", rc.getSize(), " residuals"
            trackIndex = rc.getRefIndex()
            cdcTrack = tc_tr.At(trackIndex)
           
            if genfit :
                rep = cdcTrack.getCardinalRep()
                mom = rep.getMom(rep.getReferencePlane()).Mag()
                cdcNHits.Fill(cdcTrack.getNumHits())
                stat = cdcTrack.getCardinalRep().getStatusFlag()
                cdcVal.Fill(stat)
                if stat:
                    continue
                
            else :
                mom = cdcTrack.GetMom()
                cdcNHits.Fill(cdcTrack.GetNpoint())
                cdcVal.Fill(cdcTrack.GetValid())
                
            cdcMom.Fill(mom)    
            if mom < mom_low or mom > mom_high :
                continue
            for nR in range(rc.getSize()) :
                res = rc.getResidual(nR)
                dist = res.getResXYZ().Perp()
                icl = res.getHitIndex()
                #get cluster:
                #cl = tc_cl.At(icl)
                #print icl, "      ", cl.get2DSize()
                               
                refPos = res.getRefPos()
                clpos = res.getHitPos()
                clPosR.Fill(clpos.Perp())
                clPosZ.Fill(clpos.Z())
                rad=refPos.Perp()
                dPhi = (clpos.Phi()-refPos.Phi())
                dPhiVsR.Fill(rad,dPhi*rad)
                zCount=0
                for i in range(len(zCuts)):
                    zCount=i
                    if clpos.Z() < (zCuts[i]+zOffset):
                        break
                #print "Slice:",zCount,"zValue =",clpos.Z()    
                if zCount not in dictdPhiVsR.keys() :
                    dictdPhiVsR[zCount] = ROOT.TH2D(dPhiVsR)
                    histos.append(dictdPhiVsR[zCount])
                    dictdPhiVsR[zCount].Reset()
                    title = dictdPhiVsR[zCount].GetTitle()
                    title+=(" [z>"+str(zCuts[zCount])+"]")
                    name = dictdPhiVsR[zCount].GetName()
                    name+=(" [z>"+str(zCuts[zCount])+"]")
                    dictdPhiVsR[zCount].SetTitle(title)
                    dictdPhiVsR[zCount].SetName(name)
                    
                dictdPhiVsR[zCount].Fill(rad,dPhi*rad)
                distHist.Fill(dist)
                distVsR.Fill(rad,dist)
        counter+=1                        
print ""
print ""
canv = ROOT.TCanvas()
canv.Divide(2,2)
canv.cd(1)
distVsR.Draw("colz")
canv.cd(2)
dPhiVsR.Draw("colz")
canv.cd(3)
clPosR.Draw()
canv.cd(4)
clPosZ.Draw()
canv2 = ROOT.TCanvas()
canv2.Divide(3,1)
canv2.cd(1)
distHist.Draw()
canv2.cd(2)
cdcMom.Draw()
max = cdcMom.GetMaximum()
l1 = ROOT.TLine()
l1.SetLineColor(ROOT.kRed+1)
l1.DrawLine(mom_low,0,mom_low,max)
l2 = ROOT.TLine()
l2.SetLineColor(ROOT.kRed+1)
l2.DrawLine(mom_high,0,mom_high,max)
canv2.cd(3)
cdcNHits.Draw()
c3 = ROOT.TCanvas()
c3.Divide(len(zCuts)/2 + 1,2)
for i in dictdPhiVsR.keys() :
    c3.cd(i+1)
    dictdPhiVsR[i].Draw("colz")
c4 = ROOT.TCanvas()
c4.Divide(3,1)
c4.cd(1)
cdcVal.Draw()

print "Creating output file and storing histograms ..."
orfile = ROOT.TFile(ofile,"recreate")
orfile.cd()
for h in histos:
    h.Write()
if not batch:
    input()