import glob, sys, ROOT


noDistortName = "Momres_PDG211_mom0.5_deg23_mult1_noDistort.reco.root"
distortName = "Momres_PDG211_mom0.5_deg23_mult1_distort.reco.root"
corrName = "Momres_PDG211_mom0.5_deg23_mult1_corrected.reco.root"

outname = "corrAna.root"


ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine(".x rootlogon_Bernhard_New.C")         
ROOT.gROOT.ProcessLine('gROOT->SetStyle("col")')
ROOT.gROOT.ProcessLine('gROOT->ForceStyle()')

NameList = [noDistortName,distortName,corrName]

colors = [ROOT.kBlack,ROOT.kRed-3,ROOT.kGreen-3]
histMap = {}
nBinsPhi = 50
nBinsTheta = 20
momresHistDiff = {}
momresHistMap = {}
momresHistNormalMap = {}
colorMap = {}
count = 0
for i in NameList:
    histMap[i] = ROOT.TH1D(i,"meh",300,0.3,0.7)
    momresHistMap[i] = ROOT.TH2D("momresHist"+i, "GDGFDG", nBinsPhi,-180,180,nBinsTheta,0,180)
    momresHistDiff[i] = ROOT.TH1D("momresDiff"+i, "GDGFDG", 300,-0.2,0.2)
    momresHistNormalMap[i] = ROOT.TMatrixD(nBinsPhi, nBinsTheta)
    histMap[i].GetXaxis().SetTitle("Momentum (GeV/c)")
    colorMap[i] = colors[count]
    count+=1

phiBinWidth = 360./((float)(nBinsPhi))
thetaBinWidth = 180./((float)(nBinsTheta))
    



nFile=0
for name in NameList:
    rfile = ROOT.TFile(name,"READ")
    tree = rfile.Get("cbmsim")
    tree.SetBranchStatus("*", 0)
    tree.SetBranchStatus("TrackPostFit.*", 1)
    #mcname = name.replace(".reco.root", ".mc.root")
    #mcfile = ROOT.TFile(mcname, "READ")
    #mctree = mcfile.Get("cbmsim")
    #mctree.SetBranchStatus("*", 0)
    #mctree.SetBranchStatus("MCTrack.*", 1)
    count = 0
    for e in tree:
        #mctree.GetEvent(count)
        count += 1
        ntracks = e.TrackPostFit.GetEntriesFast()
        if ntracks > 1 :
            continue
        for track in e.TrackPostFit:
            
            status = track.getCardinalRep().getStatusFlag()
            if status > 0:
                continue
            refplane = track.getCardinalRep().getReferencePlane()  #workaround
            mom = track.getCardinalRep().getMom(refplane)
            #mom.Print()
            histMap[name].Fill(mom.Mag())
            theta = mom.Theta()*180/ROOT.TMath.Pi()
            phi = mom.Phi()*180/ROOT.TMath.Pi()
            binPhi = (int)((phi+180)/phiBinWidth)
            binTheta = (int)((theta)/thetaBinWidth)
            #print binPhi,binTheta
            #momresHistMap[name].Fill(phi, theta)
            momResidual = mom.Mag()-0.5
            momresHistDiff[name].Fill(momResidual)
            if momResidual > 0.1 or  momResidual < -0.1 :
                continue
            momresHistMap[name].SetBinContent(binPhi+1, binTheta+1, momresHistMap[name].GetBinContent(binPhi+1, binTheta+1)+momResidual)
            (momresHistNormalMap[name])[binPhi][binTheta] += 1.
            
                            
    rfile.Close()    
    nFile+=1

#normalize momres-maps
for i in NameList :
    hist =  momresHistMap[i]
    norm = momresHistNormalMap[i]
    for phi in range(nBinsPhi):
        for theta in range (nBinsTheta):
            value = hist.GetBinContent(phi+1,theta+1)
            n = norm[phi][theta]
            if n > 0.5 :
                hist.SetBinContent(phi+1,theta+1,value/n)
            


canv = ROOT.TCanvas("dsadsad","dsdsd",800,500)

outfile = ROOT.TFile(outname, "RECREATE")
count = 0
for name in NameList:
    arg = ""
    if count > 0 :
        arg = "same"
    histo = histMap[name]    
    histo.SetLineColor(colorMap[name])
    histo.SetFillColor(colorMap[name])
    histo.Draw(arg)
    histo.Write()
    count += 1

histMap[NameList[0]].Draw("same")

canv2 = ROOT.TCanvas()
canv2.Divide(3,2)
count = 1
for i in NameList:
    canv2.cd(count)
    momresHistMap[i].Draw("colz")
    canv2.cd(count+3)
    momresHistDiff[i].Draw()
    
    count+=1
    

outfile.Close()

input()