import ROOT,sys,math,os,time,glob
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 *
from  pion_histos import *

parser=argparse.ArgumentParser(description='analyze cosmics from ana file')
#parser.add_argument("--branch",help="which branch to use",default="TrackFitStat_0")
parser.add_argument("filename",help="the file to analyze")
parser.add_argument("--slicef",help='a root file with the slices histo',default="")
parser.add_argument("--nEvts",help='number of events to analyze',default=-1,type=int)
parser.add_argument("--outfname",help="name of the output file",default="",type=str)
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("--drawSingle",help="draw a single histo or dict",default="",type=str)
parser.add_argument("--covErr",help="use errors derived from shapematrix",action="store_true")
parser.add_argument("--axErr",help='use errors from cov matrix along axis',action="store_true")
parser.add_argument("--cuts",help='enable cosmics cuts (normally not used)',action='store_true')
parser.add_argument("--hlp",help="display help",action='store_true')
parser.add_argument('--batch',help="batchmode",action='store_true')
parser.add_argument('--rangefile',help='file with the ranges for the histograms',type=str,default='defaultRanges.txt')
parser.add_argument('--nopostprocess',help='do not post process histograms',action='store_true')
parser.add_argument('--parfile',help='the param file',type=str,default='')
parser.add_argument('--info',help='additional information. a string: name=val,name2=val2,...',type=str,default='')
parser.add_argument('--detonly',help='use oly detector with ID',type=int,default=-1)
parser.add_argument('--tpcID',help='the tpc detector if (default=%(default)s)',type=int,default=22)
parser.add_argument('--tpconly',help='also analyze tpconly branch',action='store_true')

args=parser.parse_args()

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

if args.batch:
    ROOT.gROOT.SetBatch(True)

print 'the pandapath=',pandapath

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gStyle->SetPalette(1)')
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')

if args.outfname=="":
    args.outfname=args.filename.replace('.reco.root','.hist.root')

print 'histos will be written into:', args.outfname
histsdict=pion_histos(args.rangefile)

zSlices = []
if args.slicef!="":
    print "slicename", args.slicef
    infiles=ROOT.TFile(args.slicef)
    zSlices = []
    histsdict['hslices'] = infiles.Get("hslices")
    for k in range(1,7):
        print histsdict['hslices'].GetBinContent(k)
        zSlices.append(histsdict['hslices']).GetBinContent(k)
else:
    zSlices = (  12.0, 24.0, 36.0,  48.0, 60.0, 73.0 ) # cosmic ne
    #zSlices =  (   20,  30,  40,  49,  57, 70 ) # beam
    for i in range(len(zSlices)):
        histsdict['hslices'].SetBinContent(i+1,zSlices[i])

phiSlices = (30,60,90,120,150,180)
thetaSlices=[i*20 for i in range(9)]
hslices=cosmic_hslices(zSlices,args.rangefile)

tree=openTree(args.filename,args.runs,args.pattern)
if tree == None :
    print "No cbmsim in file", file
    if type(Rfile)==ROOT.TFile:
        Rfile.Close()
        exit()


haspar=False
if args.parfile!='':
    parfilename=args.parfile
else:
    parfilename=args.filename.replace('reco','param')

if os.path.isfile(parfilename):
    parfile=ROOT.TFile.Open(parfilename)
    digiPar=parfile.Get('TpcDigiPar')
    haspar=True
if not haspar:
    paths=parfilename.split('/')
    path_idx=[i for i,p in enumerate(paths) if 'parts' in p]
    for idx in path_idx:
        paths.pop(idx)
    paths.pop(-1)
    paths.pop(-1)
    parfilepath="/".join(paths)
    parfilename=glob.glob(parfilepath+'/*.param.root')
    if len(parfilename)!=0:
        parfilename=parfilename[0]
    else:
        parfilename=''
    print parfilename
    if os.path.isfile(parfilename):
        parfile=ROOT.TFile.Open(parfilename)
        digiPar=parfile.Get('TpcDigiPar')
        haspar=True


tree.SetBranchStatus("*", 0)
#tree.SetBranchStatus(args.branch+".*",1)
tree.SetBranchStatus("Particles.*",1)
if args.tpconly:
    tree.SetBranchStatus("TpcOnlyParticles.*",1)

trackbranch=ROOT.TBranch()
if args.nEvts==-1 or tree.GetEntries()<args.nEvts:
    todoev=tree.GetEntries()
else:
    todoev=args.nEvts
print "starting:", args.filename, "("+str(todoev)+")"

main_start=time.time()
ev_start=time.time()
digiPerpVariance=[]
digiNum=[]
for iz in range(75):
    digiPerpVariance.append(0)
    digiNum.append(0)


evt=-1
step_ev=-1
done=0
olddone=0
for e in tree:
    evt+=1
    step_ev+=1
    if args.nEvts!=-1 and args.nEvts<=evt:
        break

    done=int(float(evt*100./todoev))
    if done!=olddone:
        if step_ev==0:
            continue
        t_per_ev=(time.time()-ev_start)/step_ev
        eta_tot=(todoev-evt)*t_per_ev
        eta_h=int(eta_tot/(60*60))
        eta_min=int((eta_tot-60*60*eta_h)/60)
        eta_s=eta_tot-60*60*eta_h-60*eta_min
        print "{0}% done, Good event: {1} (t_tot={2:.2}min; t_ev={3:.2}s; eta={4:02}:{5:02}:{6:02})".format(done,evt,(time.time()-main_start)/60,t_per_ev,eta_h,eta_min,int(eta_s))
        ev_start=time.time()
        step_ev=-1
        olddone=done
        
    if args.tpconly:
        for part in e.TpcOnlyParticles:
            histsdict['fitflag'].Fill(part.getFitFlag())
            if part.getFitFlag()!=0:
                continue
            nhits=part.getNumHits()
            for ihit in range(nhits):
                detID=part.getDetID(ihit)
                
                if args.detonly!=-1 and args.detonly!=detID:
                    continue
                
                hitPosXYZ  =part.getPosXYZ(ihit)
                trackResXYZ=part.getResXYZ(ihit)
                trackPosXYZ=part.getTrackPosXYZ(ihit)
                
                histsdict['hxyzTpcOnly'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z())
                histsdict['residuals Tpc Only']['hxyz res' ].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.X())
                histsdict['residuals Tpc Only']['hxyz resY'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.Y())
                histsdict['residuals Tpc Only']['hxyz resZ'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.Z())
            
        
    for part in e.Particles:
        histsdict['fitflag'].Fill(part.getFitFlag())
        
        dedx=part.getdEdx()
        ndedx=dedx.nEntries()
        histsdict['dedx']['fopi'].Fill(dedx.getMom().Mag(),dedx.truncMean(0.05,0.25))
        histsdict['dedx']['mom'].Fill(dedx.getMom().Mag())
        #if args.tpcID==detID:
        #    histsdict['dedx']['tpc'].Fill(dedx.truncMean(0.05,0.25),dedx.getMom().Mag())
        for i in range(ndedx):
            histsdict['dedx']['dedx'].Fill(dedx.getdE(i)/dedx.getdx(i))
            pass
        
        if part.getFitFlag()!=0:
            continue
        nhits=part.getNumHits()
        for ihit in range(nhits):
            detID=part.getDetID(ihit)
            
            if args.detonly!=-1 and args.detonly!=detID:
                continue
            
            hitPosXYZ   =part.getPosXYZ(ihit)
            trackResXYZ =part.getResXYZ(ihit)
            trackPosXYZ =part.getTrackPosXYZ(ihit)
            trackResA   =part.getResA(ihit)
            trackResUV  =part.getResUV(ihit)
            errXYZ      =part.getSigXYZ(ihit)
            errA        =part.getSigA(ihit)
            errUV       =part.getSigUV(ihit)
            clsizeXYZ   =part.getFullClusterSize(ihit)
            clsizeA     =part.getFullClusterSizeAxis(ihit)
            trackMom    =part.getTrackMom(ihit)
            theta=trackMom.Theta()*ROOT.TMath.RadToDeg()
            phi=trackMom.Phi()*ROOT.TMath.RadToDeg()

            histsdict['hxyz'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z())
            histsdict['residuals']['hxyz res' ].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.X())
            histsdict['residuals']['hxyz resY'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.Y())
            histsdict['residuals']['hxyz resZ'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.Z())
            
            
            histsdict['resSurf']['resX']    .Fill(hitPosXYZ.Z(),theta,trackResXYZ.X()*10000) #residual in mum
            histsdict['resSurf']['resY']    .Fill(hitPosXYZ.Z(),theta,trackResXYZ.Y()*10000)
            histsdict['resSurf']['resZ']    .Fill(hitPosXYZ.Z(),theta,trackResXYZ.Z()*10000) 
            histsdict['resSurf']['resA1']   .Fill(hitPosXYZ.Z(),theta,trackResA.Y()*10000)
            histsdict['resSurf']['resA2']   .Fill(hitPosXYZ.Z(),theta,trackResA.Z()*10000)
            histsdict['resSurf']['resU']    .Fill(hitPosXYZ.Z(),theta,trackResUV.X()*10000)
            histsdict['resSurf']['resV']    .Fill(hitPosXYZ.Z(),theta,trackResUV.Y()*10000)
            histsdict['resSurf']['errX']    .Fill(hitPosXYZ.Z(),theta,errXYZ.X()*10000)
            histsdict['resSurf']['errY']    .Fill(hitPosXYZ.Z(),theta,errXYZ.Y()*10000)
            histsdict['resSurf']['errZ']    .Fill(hitPosXYZ.Z(),theta,errXYZ.Z()*10000)
            histsdict['resSurf']['errA1']   .Fill(hitPosXYZ.Z(),theta,errA.Y()*10000)
            histsdict['resSurf']['errA2']   .Fill(hitPosXYZ.Z(),theta,errA.Z()*10000)
            histsdict['resSurf']['errU']    .Fill(hitPosXYZ.Z(),theta,errUV.X()*10000)
            histsdict['resSurf']['errV']    .Fill(hitPosXYZ.Z(),theta,errUV.Y()*10000)
            histsdict['resSurf']['pullX']   .Fill(hitPosXYZ.Z(),theta,trackResXYZ.X()/errXYZ.X())
            histsdict['resSurf']['pullY']   .Fill(hitPosXYZ.Z(),theta,trackResXYZ.Y()/errXYZ.Y())
            histsdict['resSurf']['pullZ']   .Fill(hitPosXYZ.Z(),theta,trackResXYZ.Z()/errXYZ.Z())
            histsdict['resSurf']['pullA1']  .Fill(hitPosXYZ.Z(),theta,trackResA.Y()/errA.Y())
            histsdict['resSurf']['pullA2']  .Fill(hitPosXYZ.Z(),theta,trackResA.Z()/errA.Z())
            histsdict['resSurf']['pullU']   .Fill(hitPosXYZ.Z(),theta,trackResUV.X()/errUV.X())
            histsdict['resSurf']['pullV']   .Fill(hitPosXYZ.Z(),theta,trackResUV.Y()/errUV.Y())
            histsdict['resSurf']['clsizeX'] .Fill(hitPosXYZ.Z(),theta,clsizeXYZ.X())
            histsdict['resSurf']['clsizeY'] .Fill(hitPosXYZ.Z(),theta,clsizeXYZ.Y())
            histsdict['resSurf']['clsizeZ'] .Fill(hitPosXYZ.Z(),theta,clsizeXYZ.Z())
            histsdict['resSurf']['clsizeA1'].Fill(hitPosXYZ.Z(),theta,clsizeA.Y())
            histsdict['resSurf']['clsizeA2'].Fill(hitPosXYZ.Z(),theta,clsizeA.Z())
            
            
              
            histsdict['res vs Z']['cluster track X'] .Fill(trackPosXYZ.Z(),trackResXYZ.X()*10000)
            histsdict['res vs Z']['cluster track Y'] .Fill(trackPosXYZ.Z(),trackResXYZ.Y()*10000)
            histsdict['res vs Z']['cluster track Z'] .Fill(trackPosXYZ.Z(),trackResXYZ.Z()*10000)
            histsdict['res vs Z']['cluster track U'] .Fill(trackPosXYZ.Z(),trackResUV.X()*10000)
            histsdict['res vs Z']['cluster track V'] .Fill(trackPosXYZ.Z(),trackResUV.Y()*10000)
            histsdict['res vs Z']['cluster track A1'].Fill(trackPosXYZ.Z(),trackResA.Y()*10000)
            histsdict['res vs Z']['cluster track A2'].Fill(trackPosXYZ.Z(),trackResA.Z()*10000)

            if(trackPosXYZ.Perp()<7 and trackPosXYZ.Perp()>13 and detID==args.tpcID):
                pass
            else:
                if (abs(abs(phi)-90)<20):
                    histsdict['res vs Z']['cluster track X Cut20'].Fill(trackPosXYZ.Z(),trackResXYZ.X()*10000)
                if (abs(phi)<20):
                    histsdict['res vs Z']['cluster track Y Cut20'].Fill(trackPosXYZ.Z(),trackResXYZ.Y()*10000)
                if (abs(abs(theta)-90)<20):
                    histsdict['res vs Z']['cluster track Z Cut20'].Fill(trackPosXYZ.Z(),trackResXYZ.Z()*10000)

                if (abs(abs(phi)-90)<10):
                    histsdict['res vs Z']['cluster track X Cut10'].Fill(trackPosXYZ.Z(),trackResXYZ.X()*10000)
                if (abs(phi)<10):
                    histsdict['res vs Z']['cluster track Y Cut10'].Fill(trackPosXYZ.Z(),trackResXYZ.Y()*10000)
                if (abs(abs(theta)-90)<10):
                    histsdict['res vs Z']['cluster track Z Cut10'].Fill(trackPosXYZ.Z(),trackResXYZ.Z()*10000)
                    
                histsdict['res vs Z']['cluster track U Cut'].Fill(trackPosXYZ.Z(),trackResUV.X()*10000)
                histsdict['res vs Z']['cluster track V Cut'].Fill(trackPosXYZ.Z(),trackResUV.Y()*10000)
        
        
            
            if args.tpcID==detID:
                histsdict['hxyzTpc'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z())
                histsdict['residuals Tpc']['hxyz res' ].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.X())
                histsdict['residuals Tpc']['hxyz resY'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.Y())
                histsdict['residuals Tpc']['hxyz resZ'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.Z())
            
            
            #histsdict['residuals']['hxyz resU'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.X())
            #histsdict['residuals']['hxyz resV'].Fill(trackPosXYZ.X(),trackPosXYZ.Y(),trackPosXYZ.Z(),trackResXYZ.Y())

                


outfile = ROOT.TFile(args.outfname,"recreate")
outfile.cd()

outfile.mkdir('info')
outfile.cd('info')

infos=args.info.split(',')
rinfos={}
rinfos['generator'] =ROOT.TNamed('generator','undefined')
rinfos['angle']     =ROOT.TNamed('angle','undefined')
rinfos['error']     =ROOT.TNamed('error','undefined')
rinfos['correction']=ROOT.TNamed('correction','undefined')
rinfos['field']     =ROOT.TNamed('field','undefined')
rinfos['gas']       =ROOT.TNamed('gas','undefined')
rinfos['gain']      =ROOT.TNamed('gain','undefined')
rinfos['name']      =ROOT.TNamed('name','undefined')
if haspar:
    rinfos['Dl']        =ROOT.TNamed('Dl',str(digiPar.getGas().Dl()))
    rinfos['Dt']        =ROOT.TNamed('Dt',str(digiPar.getGas().Dt()))
    rinfos['vd']        =ROOT.TNamed('vd',str(digiPar.getGas().VDrift()))
    rinfos['gain']      =ROOT.TNamed('gain',str(digiPar.getGain()))
    rinfos['field']     =ROOT.TNamed('field',str(digiPar.getEField()))

if len(infos)>2:
    for inf in infos:
        inf=inf.split('=')
        rinfos[inf[0]]=ROOT.TNamed(inf[0],inf[1])
        
for rinf in rinfos:
    rinfos[rinf].Write()
outfile.cd()

histcounter = 0
outfile.mkdir('single')
for hists in histsdict:
    histcounter+=1
    hist=histsdict[hists]
    outfile.cd()
    if type(hist)==dict:
        outfile.mkdir(hists)
        outfile.cd(hists)
        for h in hist:
            hist[h].Write()
    else:
        outfile.cd('single')
        hist.Write()


outfile.cd()
outfile.mkdir('Z-slices')
for d in hslices:
    sdict=hslices[d]
    outfile.cd()
    outfile.mkdir('Z-slices/'+d)
    outfile.cd('Z-slices/'+d)
    for i in range(len(zSlices)):
        h=sdict[i]
        h.Write()


outfile.cd()
outfile.Close()     

print 'Histos are written into file:', args.outfname
print 'total runtime: {0} '.format(float((time.time()-main_start)))