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

parser=argparse.ArgumentParser(description='plot correlation of 2pad cluster')
parser.add_argument('recofile',help='recofile',type=str)
parser.add_argument('--events',help='how many events to process',type=int,default=-1)
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='+')

args=parser.parse_args()

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

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

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

tree.SetBranchStatus("CutCosmics.*",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(3,3)
c3=ROOT.TCanvas("c3","amp fracs profiles",1000,1000)
c3.Divide(3,3)
c4=ROOT.TCanvas("c4","amp frac res norm dir",1000,1000)
c4.Divide(3,3)

resVsAmpfracHists=[]
resAmpfrac=[]
Ampfrac=[]

for i in range(2):
    resVsAmpfracHists.append([])
    resAmpfrac.append([])
    Ampfrac.append([])
    for j in range(2):
        resVsAmpfracHists[i].append(ROOT.TH2D("hresVsAmpfrac_{0}_{1}".format(i,j),
                                              "MC Residual vs Amp Fraction (Size {0}; Axis {1})".format(i,j),
                                              1200,-0.1,1.1,1000,-0.3,0.3))
        '''
        resAmpfrac[i].append(ROOT.TH2D("hresAmpfrac_{0}_{1}".format(i,j),
                                       "Ampfrac dir res (Size {0}; Axis {1})".format(i,j),
                                       22,-1.1,1.1,22,-1.1,1.1))
        Ampfrac[i].append(ROOT.TH2D("hAmpfrac_{0}_{1}".format(i,j),
                                    "Ampfrac dir (Size {0}; Axis {1})".format(i,j),
                                    22,-1.1,1.1,22,-1.1,1.1))
                                    '''

tree.SetBranchStatus("*",0)
tree.SetBranchStatus("CutCosmics.*",1)


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()
                mcresA=tr.GetResA()
            else:
                mcres=tr.GetResMC()
                mcres=tr.GetResAMC()

            amps=tr.GetAmps()

            for icl in range(len(mcres)):
                dop=tr.GetDigisOnPlane(icl)
                dopa=tr.GetDigisOnPlaneAmp(icl)
                cop=tr.GetCovOnPlane(icl)

                eigenproblem=ROOT.TMatrixDEigen(cop)
                eigenwert=eigenproblem.GetEigenValues()
                eigenvector=eigenproblem.GetEigenVectors()
                ev1=ROOT.TVector3(eigenvector[0][0],eigenvector[1][0],0)
                ev2=ROOT.TVector3(eigenvector[0][1],eigenvector[1][1],0)
                
                ndigi=len(dop)
                maxd1=0
                maxd2=0
                maxa1=0
                maxa2=0
                for idi in range(ndigi):
                    dopa1=dop[idi]*ev1/ev1.Mag()
                    dopa2=dop[idi]*ev2/ev2.Mag()

                    if dopa1>maxd1:
                        maxd1=dopa1
                        maxa1=dopa[idi]
                    if dopa2>maxd2:
                        maxd2=dopa2
                        maxa2=dopa[idi]

                    
                resVsAmpfracHists

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=[]
fits=[]
fitfunc=ROOT.TF1("func","pol4",-0.1,0.6)
resAmpfracNorm=[]
for i in range(9):
    c2.cd(i+1)
    resVsAmpfracHists[i].Draw("colz")
    profiles.append(resVsAmpfracHists[i].ProfileX())
    c3.cd(i+1)
    profiles[-1].Draw()
    profiles[-1].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)+")"))
#    for idi in range(resAmpfracNorm[-1].GetBinsX()):
#        for ibi in range(resAmpfracNorm[-1].GetBinsY()):
#            if resAmpfracNorm[-1].GetBinsContent(idi,ibi)==0:
#                resAmpfracNorm[-1].SetBinsContent(idi,ibi-)
    resAmpfracNorm[-1].Draw("colz")
    #resAmpfrac[i].Draw("colz")

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()
    for i in range(9):
        resVsAmpfracHists[i].Write()
        profiles[i].Write()
        resAmpfrac[i].Write()
        Ampfrac[i].Write()
u='n'
while u!='q':
    u=raw_input('quit?')