import ROOT,sys,math,os,time,glob,copy,numpy
from array import array
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1')
sys.path.append(pandapath+'/macro/tpc/FOPI/mberger')
import argparse
from functions import *

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

parser=argparse.ArgumentParser(description="Plot the theta and phi a angles for the tracks with cuts and without cuts and inside cuts")
parser.add_argument("path",help='the path to the runfolders',type=str)
parser.add_argument("runs",help='the run numbers',type=str)
parser.add_argument("pattern",help="the pattern for file selection",type=str)
parser.add_argument("--events",help="number of events to process",type=int,default=-1)
parser.add_argument("--rout",help="name of the output file to save the histos into",type=str,default="")

args=parser.parse_args()

tree=openTree(args.path,args.runs,args.pattern)
tree.SetBranchStatus("*",0)
tree.SetBranchStatus("TrackFitStat_0.*",1)
tree.SetBranchStatus("CutCosmics.*",1)
tree.SetBranchStatus("DelTrack.*",1)

htheta={}
htheta['all']=ROOT.TH1D("hthetaAll",         "Theta All;#theta (#circ);#",180,0,180)
htheta['cut']=ROOT.TH1D("hthetaCut",         "Theta Cut;#theta (#circ);#",180,0,180)
htheta['cutted']=ROOT.TH1D("hthetaCutted","Theta Cutted;#theta (#circ);#",180,0,180)

hphi={}
hphi['all']=ROOT.TH1D("hphiAll",         "phi All;#phi (#circ);#",180,0,180)
hphi['cut']=ROOT.TH1D("hphiCut",         "phi Cut;#phi (#circ);#",180,0,180)
hphi['cutted']=ROOT.TH1D("hphiCutted","phi Cutted;#phi (#circ);#",180,0,180)

ctheta=ROOT.TCanvas("ctheta","theta",700,500)
cphi=ROOT.TCanvas("cphi","Phi",700,500)

ev=-1
for e in tree:
    ev+=1
    if args.events!=-1 and ev>args.events:
        break 
    if ev%1000==0:
        print 'at event',ev
    
    for hname in htheta:
        if hname=='all':
            branch=e.TrackFitStat_0
        if hname=='cut':
            branch=e.CutCosmics
        if hname=='cutted':
            branch=e.DelTrack
            
        for tr in branch:
            theta=tr.GetTheta()
            htheta[hname].Fill(theta)
            phi=tr.GetPhi()
            if phi<0:
                phi=180+phi
            hphi[hname].Fill(phi)
            
        
        
ctheta.cd()
htheta['all'].Draw()
htheta['cut'].SetLineColor(2)
htheta['cut'].Draw("same")
htheta['cutted'].SetLineColor(3)
htheta['cutted'].Draw('same')
ctheta.Update()

cphi.cd()
hphi['all'].Draw()
hphi['all'].SetMinimum(0)
hphi['cut'].SetLineColor(2)
hphi['cut'].Draw("same")
hphi['cutted'].SetLineColor(3)
hphi['cutted'].Draw('same')
cphi.Update()

thisIsTheEnd()

if args.rout!="":
    routfile=ROOT.TFile(args.rout,"recreate")
    for hname in htheta:
        htheta[hname].Write()
        hphi[hname].Write()
    routfile.Close()