def calcphierr(mom,momerr):
     phierrY=( 1. / (1+mom.X()/mom.Y()) ) * (1. / mom.X() )
     phierrY*=momerr.Y()
     phierrX=( 1. / (1+mom.X()/mom.Y()) ) * (1. / mom.Y() )
     phierrX*=momerr.X()
     phierr= math.sqrt(phierrY**2+phierrX**2)
     return phierr

def calcthetaerr(mom,momerr):
    beta=mom.Perp()/mom.Z()
    thetaerrX=( 1. / (1+beta) ) * 1./mom.Perp() * mom.X()/mom.Z()
    thetaerrX*=momerr.X()
    thetaerrY=( 1. / (1+beta) ) * 1./mom.Perp() * mom.Y()/mom.Z()
    thetaerrY*=momerr.Y()
    thetaerrZ=( 1. / (1+beta) ) * mom.Perp()/(mom.Z()**2) * (-1)
    thetaerrZ*=momerr.Z()
    thetaerr=math.sqrt(thetaerrX**2+thetaerrY**2+thetaerrZ**2)
    return thetaerr

def calcmomerr(mom,momerr):
    momerrX=mom.X()/mom.Mag() * momerr.X()
    momerrY=mom.Y()/mom.Mag() * momerr.Y()
    momerrZ=mom.Z()/mom.Mag() * momerr.Z()
    return math.sqrt( momerrX**2 + momerrY**2 + momerrZ**2 )

def fillHists(tree,hdict,hdict2=None,hdict3=None):
    event=0
    for e in tree:
        event+=1
        if args.nevts!=-1 and event>args.nevts:
            break
        if event%100==0:
            print event

        for tr in e.CutCosmics:
            theta=tr.GetTheta()
            phi=tr.GetPhi()
            if phi<args.phicut[0] or phi>args.phicut[1]:
                continue
            if tr.GetMeanZ()<args.zcut[0] or tr.GetMeanZ()>args.zcut[1]:
                continue

            mom=tr.GetMom()
            mcmom=tr.GetMCmom()
            mctheta=mcmom.Theta()*ROOT.TMath.RadToDeg()
            mcphi=mcmom.Phi()*ROOT.TMath.RadToDeg()
            mz=tr.GetMeanZ()
            momerr=tr.GetMomErr()

            thetasign=1
            if abs(phi)<90:
                thetasign=-1
            res_phi=phi-mcphi
            res_theta=((theta-mctheta) % 180 ) * thetasign
            res_mom=mom.Mag()-mcmom.Mag()

            hdict['theta'].Fill(res_theta)
            hdict['mom'].Fill(res_mom)
            hdict['phi'].Fill(res_phi)
            hdict['phi2'].Fill(phi)
            hdict['phi z'].Fill(mz,res_phi)
            hdict['pull_phi'].Fill( res_phi/calcphierr(mom,momerr) )
            hdict['pull_theta'].Fill( res_theta/calcthetaerr(mom,momerr) )
            hdict['pull_mom'].Fill( res_mom/calcmomerr(mom,momerr) )
            hdict['pull_phi_z'].Fill(mz , res_phi/calcphierr(mom,momerr) )
            hdict['pull_theta_z'].Fill(mz , res_theta/calcthetaerr(mom,momerr) )
            hdict['pull_mom_z'].Fill(mz , res_mom/calcmomerr(mom,momerr) )
            hdict['res_phi_z'].Fill(mz, res_phi)
            hdict['res_theta_z'].Fill(mz, res_theta)
            hdict['res_mom_z'].Fill(mz, res_mom)

            if hdict2!=None:
                theta=tr.GetzAxisDir().Theta()*ROOT.TMath.RadToDeg()
                phi=tr.GetzAxisDir().Phi()*ROOT.TMath.RadToDeg()
                hdict2['theta'].Fill(((theta-mctheta)%180) * thetasign )
                hdict2['mom'].Fill(mom.Mag()-mcmom.Mag())
                hdict2['phi'].Fill(phi-mcphi)
                hdict2['phi2'].Fill(phi)
                hdict2['phi z'].Fill(mz,phi-mcphi)
                #hdict['pull_phi'].Fill( (phi-mcphi)/calcphierr(mom,momerr) )
                #hdict['pull_theta'].Fill( (theta-mctheta)/calcthetaerr(mom,momerr) )

            if hdict3!=None:
                theta=tr.GetFirstMCDir().Theta()*ROOT.TMath.RadToDeg()
                phi=tr.GetFirstMCDir().Phi()*ROOT.TMath.RadToDeg()
                hdict3['theta'].Fill(((theta-mctheta) % 180) * thetasign)
                hdict3['mom'].Fill(mom.Mag()-mcmom.Mag())
                hdict3['phi'].Fill(phi-mcphi)
                #hdict3['phi2'].Fill(phi)
                #hdict3['phi z'].Fill(mz,phi-mcphi)

                
    return

def print_info(d1,d2,var):
    print args.titles[0]+' rms',var,':',d1[var].GetRMS()
    print args.titles[1]+' rms',var,':',d2[var].GetRMS()
    print 'rms diff',var,':',100*(d1[var].GetRMS()-d2[var].GetRMS())/d2[var].GetRMS(),'%'
    print args.titles[0]+' mean',var,':',d1[var].GetMean()
    print args.titles[1]+' mean',var,':',d2[var].GetMean()
    print 'mean diff',var,':',100*(d1[var].GetMean()-d2[var].GetMean())/d2[var].GetMean(),'%'
    print ''

def plotHists(hists,c1,c2,c3):
    maxi={}

    for f in hists:
        for hn in f:
            h=f[hn]
            if type(h)==ROOT.TH1D:
                scale=1
                if h.GetEntries()>0:
                    scale=h.GetEntries()
                h.Scale(1/scale)
                h.SetStats(0)
                maxi[hn]=max(maxi.get(hn,0),h.GetMaximum())
    opt=''
    counter=-1
    legends.append(ROOT.TLegend(0,0,1,1))
    for f in hists:
        counter+=1


        c1.cd(1)
        f['theta'].SetMaximum(maxi['theta']*1.1)
        f['theta'].Draw(opt)
        f['theta'].GetXaxis().SetTitle('Theta Residual (#circ)')
        f['theta'].SetLineColor(colors[counter])
        legends[-1].AddEntry(f['theta'],args.titles[counter],'l')
        c1.cd(2)
        f['phi'].SetMaximum(maxi['phi']*1.1)
        f['phi'].Draw(opt)
        f['phi'].GetXaxis().SetTitle('Phi Residual (#circ)')
        f['phi'].SetLineColor(colors[counter])
        c1.cd(3)
        f['mom'].SetMaximum(maxi['mom']*1.1)
        f['mom'].Draw(opt)
        f['mom'].GetXaxis().SetTitle('Momentum Residual (GeV/c)')
        f['mom'].SetLineColor(colors[counter])
        c1.Update()

        f['phi z'].GetXaxis().SetTitle('Mean Z pos of track (cm)')
        f['phi z'].GetYaxis().SetTitle('Phi Residual (#circ)')

        c3.cd(1)
        f['phi2'].SetLineColor(colors[counter])
        f['phi2'].Draw(opt)


        opt='same'

    c1.Update()
    c3.Update()


import ROOT,argparse,sys,math,os
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/mberger/')
sys.path.append(pandapath+'/macro/tpc/FOPI/')
from functions import *
from cosmic_histos import *
from parseFiles import PyFileParser

parser=argparse.ArgumentParser(description='plot differences of recl and not recl')
parser.add_argument('filenames',help='the reco files; if runs is given tchains with patterns are created (path also needed then)',default=['None'],type=str,nargs='+')
parser.add_argument('--runs',help='which runs to use',type=str,default="")
parser.add_argument('--pattern',help='pattern to identify the files which should be used',type=str,default='')
parser.add_argument('--path',help='the path to the runfolders',type=str,default='outfiles_e12/Data/Cosmics/')
parser.add_argument('--titles',help='the titles',default=['1','2'],type=str,nargs='+')
parser.add_argument('--nevts',help='number of events',type=int,default=-1)
parser.add_argument('--pfile',help='the file to save all',type=str,default='None')
parser.add_argument("--hlp",help="display help",action='store_true')
parser.add_argument('--phicut',help='only tracks inside phi window are taken',type=float,default=[-190,190],nargs=2)
parser.add_argument('--outfile',help='the root file to write the histos into',type=str,default='./reclcomp.root')
parser.add_argument('--zcut',help='only use tracks whose mean z position is in window',type=float,default=[0,99],nargs=2)
args=parser.parse_args()

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

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
set_palette()

mcangle=90
canv=[]
canv.append(ROOT.TCanvas('c1','First Point',500,1000))
canv[-1].Divide(1,3)
canv.append(ROOT.TCanvas('c2','First Point',500,1000))
canv[-1].Divide(1,3)
canv.append(ROOT.TCanvas('c3','First Point',500,1000))
canv[-1].Divide(1,3)

canv.append(ROOT.TCanvas('c4','zAxis',500,1000))
canv[-1].Divide(1,3)
canv.append(ROOT.TCanvas('c5','zAxis',500,1000))
canv[-1].Divide(1,3)
canv.append(ROOT.TCanvas('c6','zAxis',500,1000))
canv[-1].Divide(1,3)

canv.append(ROOT.TCanvas('c7','MC Point',500,1000))
canv[-1].Divide(1,3)
canv.append(ROOT.TCanvas('c8','MC Point',500,1000))
canv[-1].Divide(1,3)
canv.append(ROOT.TCanvas('c9','MC Point',500,1000))
canv[-1].Divide(1,3)

if len(args.titles)<len(args.filenames):
    args.titles=[]
    for i in range(len(filenames)):
        args.titles.append(str(i))

rfile=[]
trees=[]
hists=[]
hists2=[]
hists3=[]

fcounter=-1
for f in args.filenames:
    fcounter+=1
    print f
    '''
    if args.runs!="":
        fp=PyFileParser()
        trees.append(ROOT.TChain('cbmsim'))
        
        fp.setRuns(args.runs)
        fp.setPattern(f)
        fp.setPath(args.path)
        fp.setInRunFolder()
        fp.makeFileList()
        for i in range(fp.getNum()):
            trees[-1].AddFile(fp.getFile(i))
            print "adding",fp.getFile(i),"to chain"
        
    else:
        rfile.append(ROOT.TFile(f,'read'))
        trees.append(rfile[-1].Get('cbmsim'))
    '''
    trees.append(openTree(f, args.runs, args.pattern))
    trees[-1].SetBranchStatus('*',0)
    trees[-1].SetBranchStatus('CutCosmics.*',1)

    hists.append({})
    hists2.append({})
    hists3.append({})

    hists[-1]['theta']=ROOT.TH1D('dtheta '+args.titles[fcounter],'Theta-MCTheta '+args.titles[fcounter],200,-1,1)
    hists[-1]['theta'].Sumw2()
    hists[-1]['theta'].SetDefaultSumw2()
    hists[-1]['phi']=ROOT.TH1D('dphi '+args.titles[fcounter],'Phi-MCPhi '+args.titles[fcounter],200,-4,4)
    hists[-1]['phi2']=ROOT.TH1D('phi2 '+args.titles[fcounter],'Phi '+args.titles[fcounter],360*2,-180,180)
    hists[-1]['mom']=ROOT.TH1D('dmom '+args.titles[fcounter],'Mom-MCmom '+args.titles[fcounter],1000,-1,1)
    hists[-1]['phi z']=ROOT.TH2D('phiZ '+args.titles[fcounter],'Phi vs Z '+args.titles[fcounter],80,0,80,400,-4,4)
    hists[-1]['pull_phi']=ROOT.TH1D('pullPhi_{0}'.format(args.titles[fcounter]),
                                    'Pull Phi {0}'.format(args.titles[fcounter]),
                                    400,-4,4)
    hists[-1]['pull_theta']=ROOT.TH1D('pullTheta_{0}'.format(args.titles[fcounter]),
                                      'Pull Theta {0}'.format(args.titles[fcounter]),
                                      400,-4,4)
    hists[-1]['pull_mom']=ROOT.TH1D('pullMom_{0}'.format(args.titles[fcounter]),
                                    'Pull Mom {0}'.format(args.titles[fcounter]),
                                    400,-4,4)
    hists[-1]['pull_phi_z']=ROOT.TH2D('pullPhi_z_{0}'.format(args.titles[fcounter]),
                                      'Pull Phi Vs Z{0}'.format(args.titles[fcounter]),
                                      75,0,75,400,-4,4)
    hists[-1]['pull_theta_z']=ROOT.TH2D('pullTheta_z_{0}'.format(args.titles[fcounter]),
                                        'Pull Theta Vs Z{0}'.format(args.titles[fcounter]),
                                        75,0,75,400,-4,4)
    hists[-1]['pull_mom_z']=ROOT.TH2D('pullMom_z_{0}'.format(args.titles[fcounter]),
                                      'Pull Mom Vs Z{0}'.format(args.titles[fcounter]),
                                      75,0,75,400,-4,4)
    hists[-1]['res_phi_z']=ROOT.TH2D('resPhi_z_{0}'.format(args.titles[fcounter]),
                                     'Residual Phi Vs Z{0}'.format(args.titles[fcounter]),
                                     75,0,75,400,-4,4)
    hists[-1]['res_theta_z']=ROOT.TH2D('resTheta_z_{0}'.format(args.titles[fcounter]),
                                       'Residual Theta Vs Z{0}'.format(args.titles[fcounter]),
                                       75,0,75,400,-4,4)
    hists[-1]['res_mom_z']=ROOT.TH2D('resMom_z_{0}'.format(args.titles[fcounter]),
                                     'Residual Mom Vs Z{0}'.format(args.titles[fcounter]),
                                     75,0,75,400,-4,4)

    for h in hists[-1]:
        hists2[-1][h]=hists[-1][h].Clone(hists[-1][h].GetName()+'zAxis')
        hists2[-1][h].SetTitle(hists2[-1][h].GetTitle()+' Z-Axis')
        hists3[-1][h]=hists[-1][h].Clone(hists[-1][h].GetName()+'firstMC')
        hists3[-1][h].SetTitle(hists3[-1][h].GetTitle()+' first MC')
        hists[-1][h].SetTitle(hists[-1][h].GetTitle()+' First Point')
        hists[-1][h].SetName(hists[-1][h].GetName()+'FirstPoint')

    fillHists(trees[-1],hists[-1],hists2[-1],hists3[-1])

if len(args.filenames)<10:
    colors=getColors2()
else:
    colors=getColors()


global legends
legends=[]
print 'First track point'
plotHists(hists,canv[0],canv[1],canv[2])
print 'Z-Axis'
plotHists(hists2,canv[3],canv[4],canv[5])
print 'first MC point'
plotHists(hists3,canv[6],canv[7],canv[8])
cleg=ROOT.TCanvas('leg','leg',500,1000)
cleg.SetFillColor(0)
legends[-1].SetFillColor(0)
legends[-1].Draw()


if args.pfile!='None':
    cleg.SaveAs(args.pfile+'(')
    canv[0].SaveAs(args.pfile)
    canv[1].SaveAs(args.pfile)
    canv[2].SaveAs(args.pfile)
    canv[3].SaveAs(args.pfile)
    canv[4].SaveAs(args.pfile)
    canv[5].SaveAs(args.pfile)
    canv[6].SaveAs(args.pfile)
    canv[7].SaveAs(args.pfile)
    canv[8].SaveAs(args.pfile+')')
    #canv[8].SaveAs(args.pfile+']')

rofile = ROOT.TFile(args.outfile,'RECREATE')

for hset in hists:
    for h in hset:
        hset[h].Write()

for hset in hists2:
    for h in hset:
        hset[h].Write()

for hset in hists3:
    for h in hset:
        hset[h].Write()
        
cleg.Write()
for c in canv:
    c.Write()

tits=''
for t in args.titles:
    tits+=';'+t

tits=ROOT.TNamed("Titles",tits)

tits.Write()
rofile.Close()
u=raw_input("done?")