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

def getProfile(hist):
    profile=ROOT.TH1D(hist.GetName()+"_profile",
                      hist.GetTitle()+" Profile",
                      hist.GetNbinsX(),
                      hist.GetXaxis().GetXmin(),
                      hist.GetXaxis().GetXmax())
    #hist.GetYaxis().GetXmin(),
    #                   hist.GetYaxis().GetXmax())
    
    binwidth=1
    i=binwidth
    nbins=hist.GetNbinsX()
    #print hist.GetTitle()
    #ctmp=ROOT.TCanvas()
    #hist.Draw()
    #u=raw_input("start?")
    while i<nbins:
        proj=hist.ProjectionY("{0}_proj_{1}".format(hist.GetName(),i),i-binwidth,i)
        #proj.Draw()
        maxpos=proj.GetMaximumBin()*proj.GetXaxis().GetBinWidth(proj.GetMaximumBin()) + proj.GetXaxis().GetXmin()
        maxpos -= float(proj.GetXaxis().GetBinWidth(proj.GetMaximumBin())/2.)
        if(proj.GetEntries()!=0):
            profile.SetBinContent(i,maxpos)
            print i,maxpos
        i+=1
        #ctmp.Update()
        #u=raw_input('next')
    
    #profile.Draw()
    #u=raw_input('done')
    #ctmp.Update()
    return profile



parser=argparse.ArgumentParser(description='plot correlation of 2pad cluster')
parser.add_argument('recofile',help='recofile',type=str)
parser.add_argument('digifile',help='path to digi file')
parser.add_argument('--events',help='how many events to process',type=int,default=-1)
parser.add_argument('--padf',help='the padplane file to use',type=str,default='tpc/parfiles/padPlane_FOPI.dat')
parser.add_argument('--rfile',help='the rootfile to save the shit into',type=str,default='None')
parser.add_argument('--pfile',help='the pdf file to save the shit into',type=str,default='None')
parser.add_argument('--read',help='Read the histograms from file',action='store_true')
parser.add_argument('--data',help='read real data, not sim',action='store_true')
parser.add_argument('--moreData',help='additional data files',type=str,default=['None'],nargs='+')
parser.add_argument('--batch',help='run in batch mode (not on batch farm)',action='store_true')

args=parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") 
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

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

DigiMapper = []
file = open(args.padf, "r")
lines=file.readlines()
file.close()
lines.reverse()
for line in lines:
    koord = line.split(" ") 
    DigiMapper.append([koord[3], koord[4]])

tree=ROOT.TChain("cbmsim","cbmsim")
tree.Add(args.recofile)

if args.moreData[0]!='None':
    for addfile in args.moreData:
        tree.Add(addfile)

if not args.data:
    tree.AddFriend("cbmsim",args.digifile)

tree.SetBranchStatus("CutCosmics.*",1)
tree.SetBranchStatus("TpcSPHit.*",1)
tree.SetBranchStatus("TpcCluster.*",1)
if not args.data:
    tree.SetBranchStatus("TpcDigi.*",1)

ev=-1
c1=ROOT.TCanvas("c1","c1",1000,1400)
c1.Divide(1,2)
hresVsAmpfrac=ROOT.TH2D("hresVsAmpfrac","MC Residual vs Amp Fraction",1200,-0.1,1.1,1000,-0.3,0.3)
haxphi=ROOT.TH1D("haxphi","axis phi size2",720,-360,360)
c2=ROOT.TCanvas("c2","amp fracs",1000,1000)
c2.Divide(2,2)
c3=ROOT.TCanvas("c3","amp fracs profiles",1000,1000)
c3.Divide(2,2)
c4=ROOT.TCanvas("c4","amp frac res norm dir",1000,1000)
c4.Divide(2,2)

resVsAmpfracHists=[]
resVsAmpfracHistst=[]
resAmpfrac=[]
Ampfrac=[]
zbincountHist=[]
axisPhi=[]
axisLen=[]
axisLenD={}
resVsAmpfracHistsLen={}
resVsAmpfracHistsLen_strict={}
resVsAmpfracHistsLen3={}
for i in range(9):
    resVsAmpfracHists.append(ROOT.TH2D("hresVsAmpfrac"+str(i),
                                       "MC Residual xy vs Amp Fraction (Size "+str(i+2)+")",
                                       100,0,1,160,-0.4,0.4))
    axisPhi.append(ROOT.TH1D("haxisPhi_{0}".format(i),
                             "Axis Phi (Size {0})".format(i+2),
                             150,0,30))
    axisLen.append(ROOT.TH1D("haxisLen_{0}".format(i),
                             "Axis Length (Size {0})".format(i+2),
                             1200,0,3))
    resVsAmpfracHistst.append(ROOT.TH2D("hresVsAmpfract"+str(i),
                                       "MC Residual t vs Amp Fraction (Size "+str(i+2)+")",
                                       100,0,1,160,-0.4,0.4))
    zbincountHist.append(ROOT.TH1D("hzbincount_{0}".format(i),
                                   "z bin occupancy (Size {0})".format(i),
                                   20,0,20))

    resAmpfrac.append(ROOT.TH2D("hresAmpfrac"+str(i),"Ampfrac dir res (Size "+str(i+2)+")",
                                22,-1.1,1.1,22,-1.1,1.1))
    Ampfrac.append(ROOT.TH2D("hAmpfrac"+str(i),"Ampfrac dir (Size "+str(i+2)+")",
                             22,-1.1,1.1,22,-1.1,1.1))

resVsAmpfracHistsPhi=[]
for i in range(18):
    resVsAmpfracHistsPhi.append(ROOT.TH2D("hresVsAmpfrac_2_{0}".format(i),
                                          "MC Residual xy vs Amp Fraction (Size 2, Phi {0}-{1} )".format(i*10,(i+1)*10),
                                          100,0,1,160,-0.4,0.4))

resVsAmpfracHistsPhiSteps=[]
resVsAmpfracHistsPhiSteps3=[]
phisteps=[0,5,12,17,22,26,30]
for i in range(len(phisteps)-1):
    resVsAmpfracHistsPhiSteps.append(ROOT.TH2D("hresVsAmpfracStep_2_{0}_{1}".format(phisteps[i],phisteps[i+1]),
                                               "MC Residual xy vs Amp Fraction (Size 2, Phi {0}-{1} )".format(phisteps[i],phisteps[i+1]),
                                               100,0,1,160,-0.4,0.4))
    resVsAmpfracHistsPhiSteps3.append(ROOT.TH2D("hresVsAmpfracStep_3_{0}_{1}".format(phisteps[i],phisteps[i+1]),
                                                "MC Residual xy vs Amp Fraction (Size 3, Phi {0}-{1} )".format(phisteps[i],phisteps[i+1]),
                                                100,0,1,160,-0.4,0.4))


tree.SetBranchStatus("*",0)
tree.SetBranchStatus("TpcSPHit.*",1)
tree.SetBranchStatus("TpcCluster.*",1)
tree.SetBranchStatus("CutCosmics.*",1)
tree.SetBranchStatus("TpcDigi.*",1)

print tree.GetEntries(),'events to process'
if not args.read:
    for e in tree:
        ev+=1
        if ev%100==0:
            print "Event:",ev
        if args.events!=-1 and ev>=args.events:
            break

        for tr in e.CutCosmics:
            if args.data and tr.GetLength()<3:
                continue
            size2D=tr.Get2DClSizes()
            size3D=tr.GetClSizes()
            if args.data:
                mcres=tr.GetRes()
            else:
                mcres=tr.GetResMC()
            amps=tr.GetAmps()
            xpos=tr.GetHitPositionsX()
            ids=tr.GetHitIDs()

            for icl in range(len(mcres)):
                sphit=e.TpcSPHit.At(ids[icl])
                cluster=e.TpcCluster.At(sphit.getClusterID())
                digimap=cluster.getDigiMap()

                eta=-1
                leftamp=-1
                rightamp=-1
                leftpos=[80,80]
                rightpos=[-80,-80]
                bigpos=[0,0]
                bigpost=0
                bigamp=0
                bigampt=0
                smallpos=[0,0]
                smallposct=0
                smallamp=9999999999999999999                
                smallampt=999999999999999999
                digicol={}
                digicolt={}
                
                for dig in digimap:
                    digi=e.TpcDigi.At(dig.first)
                    
                    padId=str(digi.padId())
                    dx=float(DigiMapper[int(padId)][0])
                    dy=float(DigiMapper[int(padId)][1])
                    dt=digi.t()
                    amp=digi.amp()

                    if digicol.get(padId,None)!=None:
                        digicol[padId]['amp']+=amp
                    else:
                        digicol[padId]={}
                        digicol[padId]['amp']=amp
                        digicol[padId]['dx']=dx
                        digicol[padId]['dy']=dy

                    if digicolt.get(dt,None)!=None:
                        digicolt[dt]['amp']+=amp
                        digicolt[dt]['counter']+=1
                    else:
                        digicolt[dt]={}
                        digicolt[dt]['amp']=amp
                        digicolt[dt]['dx']=dx
                        digicolt[dt]['dy']=dy 
                        digicolt[dt]['counter']=0
                    
                
                for padId in digicol:
                    dx=digicol[padId]['dx']
                    dy=digicol[padId]['dy']
                    amp=digicol[padId]['amp']

                    if dx<leftpos[0]:
                        leftpos=[dx,dy]
                        leftamp=amp
                    if dx>rightpos[0]:
                        rightpos=[dx,dy]
                        rightamp=amp

                    #if amp<smallamp:
                    #    smallamp=amp
                    #    smallpos=[dx,dy]
                    if amp>bigamp:
                        smallamp=bigamp
                        smallpos=bigpos
                        bigamp=amp
                        bigpos=[dx,dy]
    
                for t in digicolt:
                    dt=float(t)
                    amp=digicolt[t]['amp']
                    #if amp<smallampt:
                    #    smallampt=amp
                    #    smallpost=t
                    if amp>bigampt:
                        smallpost=bigpost
                        smallampt=bigampt
                        bigampt=amp
                        bigpost=t
                    if len(digicolt)<11 and len(digicolt)>1:
                        zbincountHist[len(digicolt)-2].Fill(digicolt[t]['counter'])
                        
                if (size2D[icl]==2 and size3D[icl]==2):
                    eta=leftamp/amps[icl]
                    axis=ROOT.TVector3(rightpos[0]-leftpos[0],rightpos[1]-leftpos[1],0)
                    if axis.X()<0:
                        axis*=-1
                    if axis.Mag()!=0:
                        axis.SetMag(1)
                        hresVsAmpfrac.Fill(eta,mcres[icl]*axis)

                axis=ROOT.TVector3(smallpos[0]-bigpos[0],smallpos[1]-bigpos[1],0)
                phi=abs(axis.Phi()*ROOT.TMath.RadToDeg())
                while phi>60:
                    phi-=60
                if phi>30:
                    phi=60-phi

                mag=round(axis.Mag(),3)
                if len(digicol)<11 and len(digicol)>1:
                    #eta=smallamp/(smallamp+bigamp)
                    eta=bigamp/amps[icl]
                    #print 'eta:',eta
                    #axis.Print()

                    axisPhi[len(digicol)-2].Fill(phi)
                    axisLen[len(digicol)-2].Fill(mag)
                    axisLenD[str(mag)]=1
                    if len(digicol)==2:
                        haxphi.Fill(axis.Phi()*ROOT.TMath.RadToDeg())
                    if axis.Mag()!=0:
                        axis.SetMag(1)
                        resVsAmpfracHists[len(digicol)-2].Fill(eta,mcres[icl]*axis)
                        resVsAmpfracHists[len(digicol)-2].Fill(1-eta,-1*(mcres[icl]*axis))
                       # print 'residual',mcres[icl]*axis
                       # print ''

                        weight=mcres[icl]*axis
                        resAmpfrac[size2D[icl]-2].Fill(axis.X()*(1-eta),axis.Y()*(1-eta),weight)
                        Ampfrac[size2D[icl]-2].Fill(axis.X()*(1-eta),axis.Y()*(1-eta))

                        if len(digicol)==2:
                            for i in range(18):
                                if abs(phi)<(i+1)*10:
                                    resVsAmpfracHistsPhi[i].Fill(eta,mcres[icl]*axis)
                                    resVsAmpfracHistsPhi[i].Fill(1-eta,-1*(mcres[icl]*axis))
                                break

                            if resVsAmpfracHistsLen.get(str(mag),None)==None:
                                #print mag
                                resVsAmpfracHistsLen[str(mag)]=ROOT.TH2D("hresVsAmpfrac_2_L{0}".format(mag),
                                                                         "MC Residual xy vs Amp Fraction (Size 2, Len {0})".format(mag),
                                                                         100,0,1,80,-0.4,0.4)
                           

                            resVsAmpfracHistsLen[str(mag)].Fill(eta,mcres[icl]*axis)
                            resVsAmpfracHistsLen[str(mag)].Fill(1-eta,-1*(mcres[icl]*axis))

                            if size3D[icl]==2:
                                if resVsAmpfracHistsLen_strict.get(str(mag),None)==None:
                                    resVsAmpfracHistsLen_strict[str(mag)]=ROOT.TH2D("hresVsAmpfrac_2_L{0}_strict".format(mag),
                                                                                    "MC Residual xy vs Amp Fraction (Size 2 (strict), Len {0})".format(mag),
                                                                                    100,0,1,80,-0.4,0.4)
                                resVsAmpfracHistsLen_strict[str(mag)].Fill(eta,mcres[icl]*axis)
                                resVsAmpfracHistsLen_strict[str(mag)].Fill(1-eta,-1*(mcres[icl]*axis))

                        if len(digicol)==3:
                            if resVsAmpfracHistsLen3.get(str(mag),None)==None:
                                resVsAmpfracHistsLen3[str(mag)]=ROOT.TH2D("hresVsAmpfrac_3_L{0}".format(mag),
                                                                         "MC Residual xy vs Amp Fraction (Size 3, Len {0})".format(mag),
                                                                         100,0,1,80,-0.4,0.4)

                            resVsAmpfracHistsLen3[str(mag)].Fill(eta,mcres[icl]*axis)
                            resVsAmpfracHistsLen3[str(mag)].Fill(1-eta,-1*(mcres[icl]*axis))


                        for i in range(len(phisteps)-1):
                            if phi<phisteps[i+1]:
                                if len(digicol)==2:
                                    resVsAmpfracHistsPhiSteps[i].Fill(eta,mcres[icl]*axis)
                                    resVsAmpfracHistsPhiSteps[i].Fill(1-eta,-1*(mcres[icl]*axis))
                                if len(digicol)==3:
                                    resVsAmpfracHistsPhiSteps3[i].Fill(eta,mcres[icl]*axis)
                                    resVsAmpfracHistsPhiSteps3[i].Fill(1-eta,-1*(mcres[icl]*axis))
                                break

                if len(digicolt)<11 and len(digicolt)>1:
                    eta=smallampt/(smallampt+bigampt)
                    axis=smallpost-bigpost
                    if axis==0:
                        axis=1
                    else:
                        axis/=abs(axis)

                    resVsAmpfracHistst[len(digicolt)-2].Fill(eta,mcres[icl].Z()*axis)
                    resVsAmpfracHistst[len(digicolt)-2].Fill(1-eta,-1* (mcres[icl].Z()*axis) )
                    

else:
    infile=ROOT.TFile(args.rfile,'read')
    hresVsAmpfrac=copy.deepcopy(infile.Get("hresVsAmpfrac"))
    for i in range(9):
        resVsAmpfracHists[i]=copy.deepcopy(infile.Get("hresVsAmpfrac"+str(i)))
    infile.Close()

c1.cd(1)
hresVsAmpfrac.Draw("colz")
c1.cd(2)
haxphi.Draw()
profiles=[]
profilest=[]
profiles_phi=[]
profiles_phi3=[]
profiles_mag=[]
profiles_mag_strict=[]
profiles_mag3=[]
fits=[]
fitfunc=ROOT.TF1("func","pol5",-0.1,0.6)
resAmpfracNorm=[]
for i in range(9):
    profiles.append(resVsAmpfracHists[i].ProfileX())
    profiles[-1].GetYaxis().SetRangeUser(-0.06,0.06)
    profilest.append(resVsAmpfracHistst[i].ProfileX())
    profilest[-1].GetYaxis().SetRangeUser(-0.06,0.06)

for i in range(len(phisteps)-1):
    profiles_phi.append(resVsAmpfracHistsPhiSteps[i].ProfileX())
    profiles_phi3.append(resVsAmpfracHistsPhiSteps3[i].ProfileX())

sortedmag=resVsAmpfracHistsLen.keys()
sortedmag.sort()
for mag in sortedmag:
    profiles_mag.append(resVsAmpfracHistsLen[mag].ProfileX())
    profiles_mag.append(copy.copy(getProfile(resVsAmpfracHistsLen[mag])))

sortedmag=resVsAmpfracHistsLen_strict.keys()
sortedmag.sort()
for mag in sortedmag:
    profiles_mag_strict.append(resVsAmpfracHistsLen_strict[mag].ProfileX())
    profiles_mag_strict.append(copy.copy(getProfile(resVsAmpfracHistsLen_strict[mag])))

sortedmag=resVsAmpfracHistsLen3.keys()
sortedmag.sort()
for mag in sortedmag:
    profiles_mag3.append(resVsAmpfracHistsLen3[mag].ProfileX())
    profiles_mag3.append(copy.copy(getProfile(resVsAmpfracHistsLen3[mag])))
    

for i in range(2):
    c2.cd(i+1)    
    resVsAmpfracHists[i].Draw("colz")
    #profiles.append(resVsAmpfracHists[i].ProfileX())
    c3.cd(i+1)
    profiles[i].Draw()
    profiles[i].Fit(fitfunc)
    c4.cd(i+1)
    print "drawing",i
    resAmpfracNorm.append(Divideh(resAmpfrac[i],Ampfrac[i],"hresAmpfracNorm"+str(i),"res amp frac norm (Size "+str(i+2)+")"))
    resAmpfracNorm[-1].Draw("colz")
for i in range(2):
    c2.cd(i+3)
    resVsAmpfracHistst[i].Draw("colz")
    #profilest.append(resVsAmpfracHists[i].ProfileX())
    c3.cd(i+3)
    profilest[i].Draw()
    profilest[i].Fit(fitfunc)
    c4.cd(i+3)
    zbincountHist[i].Draw()

print "updating"
c1.Update()
print "updating 2"
c2.Update()
print "updating 3"
c3.Update()
print "updating 4"
c4.Update()

if args.pfile!='None':
    c1.SaveAs(args.pfile+'(')
    c2.SaveAs(args.pfile)
    c3.SaveAs(args.pfile)
    c4.SaveAs(args.pfile+')')

if args.rfile!='None':
    outfile=ROOT.TFile(args.rfile,'recreate')
    c1.Write()
    c2.Write()
    c3.Write()
    c4.Write()
    hresVsAmpfrac.Write()
    outfile.mkdir("AmpFracHists")
    outfile.mkdir("AmpFracHistsT")
    outfile.mkdir("Profiles")
    outfile.mkdir("ProfilesT")
    outfile.mkdir("Axis_Phi")
    outfile.mkdir("Axis_Len")
    outfile.mkdir("asorted")
    for i in range(9):
        outfile.cd("AmpFracHists")
        resVsAmpfracHists[i].Write()
        outfile.cd("AmpFracHistsT")
        resVsAmpfracHistst[i].Write()
        outfile.cd("Profiles")
        profiles[i].Write()
        outfile.cd("ProfilesT")
        profilest[i].Write()

        outfile.cd("asorted")
        resAmpfrac[i].Write()
        Ampfrac[i].Write()
        zbincountHist[i].Write()

        outfile.cd("Axis_Phi")
        axisPhi[i].Write()
        outfile.cd("Axis_Len")
        axisLen[i].Write()


    outfile.mkdir("AmpFracHists_Phi")
    for i in range(18):
        outfile.cd("AmpFracHists_Phi")
        resVsAmpfracHistsPhi[i].Write()

    outfile.mkdir("AmpFracHists_PhiStep")
    outfile.mkdir("AmpFracHists_PhiStep_Profile")
    outfile.mkdir("AmpFracHists_PhiStep3")
    outfile.mkdir("AmpFracHists_PhiStep_Profile3")
    for i in range(len(phisteps)-1):
        outfile.cd("AmpFracHists_PhiStep")
        resVsAmpfracHistsPhiSteps[i].Write()
        outfile.cd("AmpFracHists_PhiStep_Profile")
        profiles_phi[i].Write()
        outfile.cd("AmpFracHists_PhiStep3")
        resVsAmpfracHistsPhiSteps3[i].Write()
        outfile.cd("AmpFracHists_PhiStep_Profile3")
        profiles_phi3[i].Write()

    outfile.mkdir("AmpFracHists_Len")
    outfile.cd("AmpFracHists_Len")
    sortedmag=resVsAmpfracHistsLen.keys()
    sortedmag.sort()
    for  mag in sortedmag:
        resVsAmpfracHistsLen[mag].Write()

    outfile.mkdir("AmpFracHists_Len_profiles")
    outfile.cd("AmpFracHists_Len_profiles")
    for prof in profiles_mag:
        prof.Write()

    outfile.mkdir("AmpFracHists_Len_strict")
    outfile.cd("AmpFracHists_Len_strict")
    sortedmag=resVsAmpfracHistsLen_strict.keys()
    sortedmag.sort()
    for  mag in sortedmag:
        resVsAmpfracHistsLen_strict[mag].Write()
    outfile.mkdir("AmpFracHists_Len_profiles_strict")
    outfile.cd("AmpFracHists_Len_profiles_strict")
    for prof in profiles_mag_strict:
        prof.Write()

    outfile.mkdir("AmpFracHists_Len3")
    outfile.cd("AmpFracHists_Len3")

    sortedmag=resVsAmpfracHistsLen3.keys()
    sortedmag.sort()    
    for  mag in sortedmag:
        resVsAmpfracHistsLen3[mag].Write()
    outfile.mkdir("AmpFracHists_Len3_profiles")
    outfile.cd("AmpFracHists_Len3_profiles")
    for prof in profiles_mag3:
        prof.Write()

axisLenD2=[]
print "axis lengths:"
for l in axisLenD:
    axisLenD2.append(float(l))

axisLenD2.sort()
for i in axisLenD2:
    print i
print "in between"

for i in range(len(axisLenD2)-1):
    print axisLenD2[i] + 0.5*(axisLenD2[i+1]-axisLenD2[i])
if not args.batch:
    u='n'
    while u!='q':
        u=raw_input('quit?')