import ROOT,sys,math,os
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
import numpy as np
from functions import sort_inner, sort_inner2, Drawh, Divideh, fillcutred, plotfullevent
from  cosmic_histos import *


def setRange(h,mini=0,maxi=0):
    for i in range(h.GetNbinsX()+1):
        for j in range(h.GetNbinsY()+1):
            if (h.GetBinContent(i,j)==0.) :
                h.SetBinContent(i,j,-9999999)
            if (maxi!=mini):
                cont=h.GetBinContent(i,j)
                if (cont>maxi):
                    h.SetBinContent(i,j,maxi)
                if (cont<mini and cont!=-9999999):
                    h.SetBinContent(i,j,mini)
    if(mini!=maxi):
        h.GetZaxis().SetRangeUser(mini,maxi)
    return

parser=argparse.ArgumentParser(description='analyze cosmics from ana file')
parser.add_argument("filename",help="the file to analyze")
parser.add_argument("--hlp",help="display help",action='store_true')
parser.add_argument("--nEvts",help='number of events to analyze',default=-1,type=int)
parser.add_argument('--pdir',help='the picture print directory',type=str,default="not")
args=parser.parse_args()

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

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

RFile=ROOT.TFile()
Rfile = ROOT.TFile.Open(args.filename,"read")
tree = Rfile.Get("cbmsim")

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

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

tree.SetBranchStatus("*", 0)
#tree.SetBranchStatus(args.branch+".*",1)
tree.SetBranchStatus("CutCosmics.*",1)



trackbranch=ROOT.TBranch()
canvases=[]
histosMean=[]
histosMeanErr=[]
res=[]
bins=64
binsZ=39
for bi in range(bins):
    res.append([])
    for i in range(bins):
        res[bi].append([])
        for k in range(bins):
            res[bi][i].append([])
       
dx=1
dy=1
dz=2

for canv in range(binsZ):
    canvases.append(ROOT.TCanvas())
    canvases[canv].SetCanvasSize(600,1000)
    histosMean.append(ROOT.TH2D("","",32,-16.,16.,32,-16.,16.))
    histosMeanErr.append(ROOT.TH2D("","",32,-16.,16.,32,-16.,16.))
    
print len(canvases)


poZ=[]

j=-4.

i=-16
while j < 74.5:
    poZ.append(j)
    j+=dz



count=-1
evt=-1
done=0
olddone=0

cou=0
if args.nEvts==-1:
    todoev=tree.GetEntriesFast()
else:
    todoev=args.nEvts
for e in tree:
    evt+=1
    if evt%1000==0:
        print 'At event:',evt
    trk=e.CutCosmics.GetEntries()
    trackbranch=e.CutCosmics
            
    for tr in trackbranch:
 
        resX=tr.GetResX()
        posX=tr.GetHitPositionsX()
        posY=tr.GetHitPositionsY()
        posZ=tr.GetHitPositionsZ()

        for icl in range(len(resX)):
            
            resX[icl]*=  10000
            '''
            if posX[icl]>-6 and posX[icl]<-5:
                pass
            else:
                continue

            if posY[icl]>-9 and posY[icl]<-8:
                pass
            else:
                continue
            if posZ[icl]>30 and posZ[icl]<32:
                pass
            else:
                continue
                '''
            h=int((16.+posX[icl])/dx)
            z=int((16.+posY[icl])/dy)
            e=int((4.+posZ[icl])/dz)
        
            
            res[h][z][e].append(resX[icl])

#    if evt == 5000:
#        break


#ErrorMean=ROOT.TH2D("ErrorMean","ErrorMean",64,-16.,16.,64,-16.,16.)
#Mean=ROOT.TH2D("Mean","Mean",64,-16.,16.,64,-16.,16.)
mw=0
std=0
for i in range(bins):
    for j in range(bins):
        for l in range(binsZ):
            if len(res[i][j][l])==0:
                continue
            histosMeanErr[l].SetTitle("ErrorMean Z-Pos "+str(poZ[l]) + "cm")
            histosMean[l].SetTitle("Mean  Z-Pos "+str(poZ[l]) + "cm")
            histosMeanErr[l].SetName("ErrorMean_Z-Pos_"+str(poZ[l]) + "cm")
            histosMean[l].SetName("Mean_Z-Pos_"+str(poZ[l]) + "cm")

            if len(res[i][j][l])==1:
                histosMean[l].SetBinContent(i+1,j+1,res[i][j][l][0])
                histosMeanErr[l].SetBinContent(i+1,j+1,-1)
                continue
            for k in res[i][j][l]:
                mw+=k
            MW=mw/(len(res[i][j][l]))
            histosMean[l].SetBinContent(i+1,j+1,MW)
            for k in res[i][j][l]:
                std+=((MW-k)**2)
            STD=math.sqrt(std/( len(res[i][j][l]) * (len(res[i][j][l])-1) ) )
            if STD/abs(MW)>1:
                print 'i=',i,'j=',j,',l=',l
                print 'err:',STD,'len:',len(res[i][j][l]),'mw:',MW,'relerr:',STD/abs(MW)
            #print res[i][j][l]
            
            histosMeanErr[l].SetBinContent(i+1,j+1,STD/abs(MW))

                        
            mw=0
            std=0

#c1=ROOT.TCanvas("MeanEr","MeanEr")
#c2=ROOT.TCanvas("Mean","Mean")

for l in range(binsZ):
    canvases[l].SetTitle("Mean - MeanError Z- Pos:" +str(poZ[l]))
    canvases[l].Divide(1,2)
    setRange(histosMean[l], -1000,1000)
    setRange(histosMeanErr[l], -1,1)
    canvases[l].cd(1)
    histosMean[l].Draw("colz")
    canvases[l].cd(2)
    #histosMeanErr[l].GetZaxis().SetRangeUser(-1,1)
    histosMeanErr[l].Draw("colz")
    
    canvases[l].Update()
    if args.pdir!="not":
        rfile=ROOT.TFile(args.pdir+'.root','recreate')
        histosMeanErr[l].Write()
        histosMean[l].Write()
        rfile.Close()

if args.pdir!="not":
    for l in range(binsZ):
        if l==0:
            canvases[l].SaveAs(args.pdir+".pdf"+'(')
        elif l!=0 and l!= 38:
            canvases[l].SaveAs(args.pdir+".pdf")
        elif l==38:
            canvases[l].SaveAs(args.pdir+".pdf"+')')

#c1.cd()
#ErrorMean.Draw("colz")
#c1.Update()

#c2.cd()
#Mean.Draw("colz")
#c2.Update()



u='w'
while u!='q':
    u=raw_input("Done?")