import ROOT, argparse, math

parser=argparse.ArgumentParser(description='plot mc stuff')
parser.add_argument('mcfile',help='the mc file',type=str)
parser.add_argument("--thetacut",help='only take values inside theta range',type=float,default=[-360,360],nargs=2)
parser.add_argument('--zcut',help='cut in z',type=float,default=[-90,20],nargs=2)
parser.add_argument('--threeD',help='plot 3d points',action='store_true')
parser.add_argument('--hlp',help='show help',action='store_true')
args=parser.parse_args()

if args.hlp:
    parser.show_help()
    exit(0)


ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") 
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

Infile=ROOT.TFile(args.mcfile,"read")
tree=Infile.Get("cbmsim")

tree.SetBranchStatus("MCTrack.*",1)
tree.SetBranchStatus("TpcPoint.*",1)

htheta=ROOT.TH1D("htheta","MCTrack theta",360*2,-360,360)
hthetavsZp=ROOT.TH2D("hthetavsZp","theta vs ZP",900,-70,20,360*2,-360,360)
hphivsZp=ROOT.TH2D("hphivsZp","phi vs ZP",900,-70,20,360*2,-360,360)
hzp=ROOT.TH1D('hzp','Point Z',900,-70,20)
hVyVsZp=ROOT.TH2D('hVyVsZp',' Start Y vs ZP',900,-70,20,2000,-100,100)
hVxVsZp=ROOT.TH2D('hVxVsZp',' Start X vs ZP',900,-70,20,2000,-100,100)
hVzVsZp=ROOT.TH2D('hVzVsZp',' Start Z vs ZP',900,-70,20,2000,-100,100)
hthetaVsVy=ROOT.TH2D('hthetaVsVy','Theta vs Vertex Y',720,-360,360,2000,-100,100)
hMzVsZp=ROOT.TH2D('hMzVsZp',' Mom Z vs ZP',900,-70,20,220,-11,11)
hMVsZp=ROOT.TH2D('hMVsZp',' Mom vs ZP',900,-70,20,220,-22,22)
hpoints3D=ROOT.TH3D('hpoints3D','3D tpc points',300,-15,15,300,-15,15,900,-70,20)
hlxy=ROOT.TH2D('hlxy','xy last point',300,-15,15,300,-15,15)
hlyz=ROOT.TH2D('hlyz','xy last point',900,-70,20,300,-15,15)

canvas1=ROOT.TCanvas()
canvas1.Divide(3,3)
canvas2=ROOT.TCanvas()
for e in tree:
    if e.MCTrack.GetEntries()>1:
        print "skipping: more than 1 track"
        continue

    vertex=e.MCTrack[0].GetStartVertex()
    mom=e.MCTrack[0].GetMomentum()

#    if mom.Y()!=0:
#        theta=math.atan(mom.Z()/mom.Y())
    theta=mom.Theta()
    phi=mom.Phi()
    htheta.Fill(math.degrees(theta))
    hthetaVsVy.Fill(math.degrees(theta),vertex.Y())
    if math.degrees(theta)<args.thetacut[0] or math.degrees(theta)>args.thetacut[1]:
        continue
    point=None
    for point in e.TpcPoint:

        if point.GetZ()<args.zcut[0] or point.GetZ()>args.zcut[1]:
            continue

        hthetavsZp.Fill(point.GetZ(),math.degrees(theta))
        hphivsZp.Fill(point.GetZ(),math.degrees(phi))
        hzp.Fill(point.GetZ())
        hVyVsZp.Fill(point.GetZ(),vertex.Y())
        hVxVsZp.Fill(point.GetZ(),vertex.X())
        hVzVsZp.Fill(point.GetZ(),vertex.Z())
        hMzVsZp.Fill(point.GetZ(),mom.Z())
        hMVsZp.Fill(point.GetZ(),mom.Mag())
        if args.threeD:
            hpoints3D.Fill(point.GetX(),point.GetY(),point.GetZ())
    if point!=None:
        if not(point.GetZ()<args.zcut[0] or point.GetZ()>args.zcut[1]):
            hlxy.Fill(point.GetX(),point.GetY())
            hlyz.Fill(point.GetZ(),point.GetY())

canvas1.cd(1)
htheta.Draw()
canvas1.cd(2)
hzp.Draw()
canvas1.cd(3)
hthetavsZp.Draw("colz")
canvas1.GetPad(3).SetLogz()
canvas1.cd(4)
hVyVsZp.Draw("colz")
canvas1.cd(5)
hVxVsZp.Draw("colz")
canvas1.cd(6)
hVzVsZp.Draw("colz")
#hthetaVsVy.Draw("colz")
canvas1.cd(7)
hMzVsZp.Draw("colz")
canvas1.cd(8)
#hMVsZp.Draw("colz")
hlxy.Draw("colz")
canvas1.cd(9)
#hphivsZp.Draw("colz")
hlyz.Draw("colz")
canvas1.Update()
if args.threeD:
    canvas2.cd()
    hpoints3D.Draw()
u=raw_input("done?")