import sys, os, copy
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, ROOT
from functions import set_palette, openTree, thisIsTheEnd, getSlices
import random



parser = argparse.ArgumentParser(description='plot the residual width as function of z and theta')
parser.add_argument('recofile', help='the reco file', type=str,nargs='+')
parser.add_argument('--list',help='the recofile is actually a list of files',action='store_true')
parser.add_argument('--runs', help='if data, the runs to use (%(default)s)', type=str, default='3888')
parser.add_argument('--pattern', help='if data, the file name patters ('')', type=str, default='')
parser.add_argument('--events',  type=int, default=-1,help="number of events to process ( all )")
parser.add_argument('--rout',help='root output file name',type=str,default="shapeCov.root")

args = parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")
set_palette("palette",99)

trees=[]
files=[]
if args.list:
    for line in open(args.recofile[0]):
        line=line.replace('\n','')
        words=line.split(';')
        files.append(words[0])
else:
    files=args.recofile

for itree in files:
    trees.append(openTree(itree,args.runs,args.pattern))
    trees[-1].SetBranchStatus("*", 0)
    #trees[-1].SetBranchStatus("CutCosmics.*", 1)
    trees[-1].SetBranchStatus("TrackFitStat_0.*", 1)
    
    
c1=ROOT.TCanvas("c1","shape3D",1500,1000)
c1.Divide(3,1)
hshape3DA1vsZ=ROOT.TH2D("hshape3DA1vsZ","shape A1 vs Z",75,0,75,100,0,0.2)
hshape3DA2vsZ=ROOT.TH2D("hshape3DA2vsZ","shape A2 vs Z",75,0,75,100,0,0.1)
hshape3DA3vsZ=ROOT.TH2D("hshape3DA3vsZ","shape A3 vs Z",75,0,75,100,0,0.1)

c2=ROOT.TCanvas("c2","shape",1000,1000)
c2.Divide(2,1)
hshapePlaneA1vsZ=ROOT.TH2D("hshapePlaneA1vsZ","shape plane A1 vs Z",75,0,75,100,0,0.1)
hshapePlaneA2vsZ=ROOT.TH2D("hshapePlaneA2vsZ","shape plane A2 vs Z",75,0,75,100,0,0.1)

iev=-1
for e in trees[-1]:
    iev+=1
    
    if iev>args.events and args.events!=-1:
        break
    if iev%1000==0:
        print 'at event:',iev
        
    for tr in e.TrackFitStat_0:
        #print 'next track ***********************************************************************'
        nclust=tr.nhits()
        theta=tr.GetTheta()
        statflag=tr.GetStatFlag()
        for icl in range(nclust):
            shapeCov=tr.GetShapeCov(icl)
            pos=tr.GetHitPosistionXYZ(icl)
            
            shapeEigen=ROOT.TMatrixDEigen(shapeCov)
            eigenval=shapeEigen.GetEigenValues()
            hshape3DA1vsZ.Fill(pos.Z(),eigenval[0][0])
            hshape3DA2vsZ.Fill(pos.Z(),eigenval[1][1])
            hshape3DA3vsZ.Fill(pos.Z(),eigenval[2][2])
            
            shapeCovPlane=tr.GetShapeCovOnPlane(icl)
            
            shapePlaneEigen=ROOT.TMatrixDEigen(shapeCovPlane)
            eigenvalPlane=shapePlaneEigen.GetEigenValues()
            hshapePlaneA1vsZ.Fill(pos.Z(),eigenvalPlane[0][0])
            hshapePlaneA2vsZ.Fill(pos.Z(),eigenvalPlane[1][1])
            

            ''' 
            if( abs(theta-90)<25 and pos.Z()<20 and statflag==0):
                cov=tr.GetCovOnPlane(icl)          
                cov.Print()
               '''     
            
c1.cd(1)        
hshape3DA1vsZ.Draw("colz")
c1.cd(2)        
hshape3DA2vsZ.Draw("colz")
c1.cd(3)        
hshape3DA3vsZ.Draw("colz")
c1.Update()
          
c2.cd(1)
hshapePlaneA1vsZ.Draw("colz")        
c2.cd(2)
hshapePlaneA2vsZ.Draw("colz")      
c2.Update()

thisIsTheEnd()