import ROOT,argparse

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gSystem->Load("libPhysics")')
ROOT.gROOT.ProcessLine('gStyle->SetPalette(1)')
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
ROOT.gROOT.ProcessLine('gSystem->Load("libGeom")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

parser=argparse.ArgumentParser(description='Plot stuff along a track')
parser.add_argument('filename',help='filename of the reco file',type=str)
parser.add_argument('--ev',help='which event to go to',type=int)
parser.add_argument('--zcut',help='zcut',type=int,default=-80)
parser.add_argument('--data',help='analyse real data',action='store_true')
args=parser.parse_args()

rfile=ROOT.TFile.Open(args.filename,"read")
if rfile.IsOpen()==ROOT.kTRUE:
    tree = rfile.Get("cbmsim")
if tree==None:
    print "No cbmsim found"
    if type(Rfile)==ROOT.TFile:
        Rfile.Close()
    exit()

hamp=ROOT.TH1D("hamp","Clusteramp along track",600,0,60)
hresx=ROOT.TH1D("hresx","X-Residual along track",600,0,60)
hresu=ROOT.TH1D("hresu","u-Residual along track",600,0,60)
hresv=ROOT.TH1D("hresv","v-Residual along track",600,0,60)
hresp=ROOT.TH1D("hresp","Residual along track",600,0,60)
hres2d=ROOT.TH2D("hres2d","res vs amp along track",600,0,60,200,-1,1)
hxy=ROOT.TH2D("hxy","xypos, amp weighted",300,-15,15,300,-15,15)
hclsize=ROOT.TH1D("hclsize","Clustersize along track",600,0,60)
c1=ROOT.TCanvas("c1","bla",700,1000)
c2=ROOT.TCanvas("c2","bla",700,1000)
c1.Divide(1,2)
c2.Divide(1,2)
oldpos=ROOT.TVector3(0,0,0)
dist=0
evnum=-1
print "events:",tree.GetEntries()
for e in tree:
    evnum+=1
    trackindex=-1
    print "tracks:",e.TrackPostFit.GetEntries()
    if evnum<args.ev:
        continue
    for tr in e.TrackPostFit:
        trackindex+=1
        cand = tr.getCand()  
        tfs=e.TrackFitStat_0.At(trackindex)
        resX=tfs.GetResX()
        resY=tfs.GetResY()
        resZ=tfs.GetResZ()
        resU=tfs.GetResU()
        resV=tfs.GetResV()
        if args.data:
            hitIds = cand.getHitIDs(8)
        else:
            hitIds = cand.getHitIDs(14)
        hitindex=-1
        if resX.size()==0:
            continue
        for ihit in hitIds:
            hitindex+=1
            if hitindex>=resX.size():
                continue            
            hit=e.TpcSPHit.At(ihit)
            clid=hit.getClusterID()
            cl=e.TpcCluster.At(clid)
            pos=cl.pos()
            if pos.Z()<args.zcut:
                continue
            if oldpos.Mag()!=0:
                dist+=(pos-oldpos).Mag()
            hamp.Fill(dist,cl.amp())
            print hitindex,resX.size()
            res = ROOT.TVector3(resX[hitindex],resY[hitindex],resZ[hitindex])
            hres2d.Fill(dist,res.X(),cl.amp())
            hxy.Fill(cl.pos().X(),cl.pos().Y(),cl.amp())
            hresx.Fill(dist,res.Mag())
            hresu.Fill(dist,resU[hitindex])
            hresv.Fill(dist,resV[hitindex])
            hresp.Fill(dist,resV[hitindex]+resU[hitindex])
            hclsize.Fill(dist,cl.size())
            oldpos=pos
            print dist,cl.amp(),pos.Z(),evnum
            
        if dist!=0:
            c1.cd(1)
            hamp.Draw()
            c1.cd(2)
#            hres2d.Draw("colz")
#            hresx.Draw()
            hxy.Draw("colz")
            c1.Update()
            c2.cd(1)
#            hresu.Draw()
            hresp.Draw()
            c2.cd(2)
#            hresv.Draw()
            hclsize.Draw()
            c2.Update()
            dist=0
            u=raw_input("next track?")
            hamp.Reset()
            hresx.Reset()
            hres2d.Reset()
            hxy.Reset()
            hresu.Reset()
            hresv.Reset()
            hresp.Reset()
            hclsize.Reset()
            if u=='q':
                exit()