import sys
sys.path.append('/home/mberger/fopiroot/fopiroot_dev/macro/tpc/FOPI/mberger')
import ROOT,glob,math,argparse
from anaFile import anaFile
from functions 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,-9999)
            if (maxi!=mini):
                cont=h.GetBinContent(i,j)
                if (cont>maxi):
                    h.SetBinContent(i,j,maxi)
                if (cont<mini and cont!=-9999):
                    h.SetBinContent(i,j,mini)
    if(mini!=maxi):
        h.GetZaxis().SetRangeUser(mini,maxi)
    return

def drawSlice(hd,c):
    counter=0
    for h in hd:
        c.cd(counter+1)
        setRange(h,args.borders[0],args.borders[1])
        h.SetStats(0)
        h.Draw("colz")
        counter+=1

parser=argparse.ArgumentParser(description="plot the residuals in xy plane")
parser.add_argument('anafiles',help='the ana files',type=str,nargs='+')
parser.add_argument('--pdir',help='the picture print directory',type=str,default="not")
parser.add_argument('--titles',help='the titles for',type=str,default=[""],nargs="+")
parser.add_argument('--diff',help='calculate the difference and other crap between 2 input files',
                    action='store_true')
parser.add_argument('--borders',help='boraders for z-axis',type=float,default=[-1000,1000],nargs=2)
parser.add_argument('--hlp',help='show help',action='store_true')
parser.add_argument('--mc',help='plot MC residuals instead of fit residuals',action='store_true')
parser.add_argument('--phi',help='plot the phi sliced x-residuals',action='store_true')
args=parser.parse_args()


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

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

rfiles=[]
dslices=[]
xyslice=[]
xyphislice=[]
xysliceR=[]
xysliceP=[]
xyres=[]
yzres=[]
xreshistZ=[]
xreshist=[]
canv=[[],[],[],[],[]]#*2
if args.phi:
    canvphi=[]
fcounter=-1
for ifile in range(len(args.anafiles)):
    fcounter+=1
    rfiles.append(anaFile(args.anafiles[ifile]))
    dslices.append(rfiles[-1].getDriftSlices())

    if (args.mc):
        xyres.append(rfiles[-1].getHist("hxyMCresnorm"))
        yzres.append(rfiles[-1].getHist("hyzMCresnorm"))
        xyslice.append(rfiles[-1].getSlicedHist("hxyMCresnorm"))
        xysliceR.append(rfiles[-1].getSlicedHist("hxyMCresRnorm"))
        xysliceP.append(rfiles[-1].getSlicedHist("hxyMCresPnorm"))
    else:
        xyres.append(rfiles[-1].getHist("hxyresnorm"))
        yzres.append(rfiles[-1].getHist("hyzresnorm"))
        xyslice.append(rfiles[-1].getXYsliceHists())
        xysliceR.append(rfiles[-1].getSlicedHist("hxyresRnorm"))
        xysliceP.append(rfiles[-1].getSlicedHist("hxyresPnorm"))

    if (args.phi):
        if args.mc:
            xyphislice.append(rfiles[-1].getPhiSlicedHist("hxyMCresnorm_phi"))
        else:
            xyphislice.append(rfiles[-1].getPhiSlicedHist("hxyresnorm_phi"))

    canv[0].append(ROOT.TCanvas())
    canv[2].append(ROOT.TCanvas())
    if args.titles[0]!="":
        canv[0][-1].SetTitle(args.titles[fcounter])
        canv[2][-1].SetTitle(args.titles[fcounter])
    canv[0][-1].SetRightMargin(0.19)
    canv[2][-1].SetRightMargin(0.19)
    factor=3.5

#    xyres[fcounter].GetZaxis().SetRangeUser(-2000,2000)
    xyres[fcounter].GetXaxis().SetTitle("X Pos (cm)")
    xyres[fcounter].GetYaxis().SetTitle("Y Pos (cm)")
    xyres[fcounter].GetZaxis().SetTitle("Residual (#mu m)")
    xyres[fcounter].GetZaxis().SetTitleOffset(1.6)
    xyres[fcounter].SetStats(0)

#    yzres[fcounter].GetZaxis().SetRangeUser(-1000,1000)
    yzres[fcounter].GetXaxis().SetTitle("X Pos (cm)")
    yzres[fcounter].GetYaxis().SetTitle("Y Pos (cm)")
    yzres[fcounter].GetZaxis().SetTitle("Residual (#mu m)")
    yzres[fcounter].GetZaxis().SetTitleOffset(1.6)
    yzres[fcounter].SetStats(0)


    setRange(xyres[fcounter],args.borders[0],args.borders[1])
    setRange(yzres[fcounter],args.borders[0],args.borders[1])

    canv[0][-1].cd()
    xyres[fcounter].Draw("colz")
    canv[2][-1].cd()
    yzres[fcounter].Draw("colz")

    canv[1].append(ROOT.TCanvas())
    canv[3].append(ROOT.TCanvas())
    canv[4].append(ROOT.TCanvas())

    if args.titles[0]!="":
        canv[1][-1].SetTitle(args.titles[fcounter]+" X")
        canv[3][-1].SetTitle(args.titles[fcounter]+" R")
        canv[4][-1].SetTitle(args.titles[fcounter]+" P")

    canv[1][-1].Divide(3,2)
    canv[3][-1].Divide(3,2)
    canv[4][-1].Divide(3,2)
    counter=0
    for h in xyslice[fcounter]:
        canv[1][-1].cd(counter+1)
        setRange(h,args.borders[0],args.borders[1])
#        h.GetZaxis().SetRangeUser(args.borders[0],args.borders[1])
        h.SetStats(0)
#        setRange(h)
        h.Draw("colz")
        counter+=1

    drawSlice(xysliceR[fcounter],canv[3][-1])
    drawSlice(xysliceP[fcounter],canv[4][-1])
#    counter=0
#    for h in xysliceR[fcounter]:
#        canv[3][-1].cd(counter+1)
#        setRange(h,args.borders[0],args.borders[1])
#        h.SetStats(0)
#        h.Draw("colz")
#        counter+=1


    counter=0
    if args.phi:
        canvphi.append(ROOT.TCanvas())
        canvphi[-1].Divide(3,2)
        if args.titles[0]!="":
            canvphi[-1].SetTitle(args.titles[fcounter])
        for h in xyphislice[fcounter]:
            canvphi[-1].cd(counter+1)
            setRange(h,args.borders[0],args.borders[1])
            h.SetStats(0)
            h.Draw("colz")
            counter+=1

    for c in canv:
        c[-1].Update()
#    canv[0][-1].Update()
#    canv[1][-1].Update()
#    canv[2][-1].Update()

    if args.phi:
        canvphi[-1].Update()

if args.diff:
#    hdiff=Addh(xyres[0],xyres[1],-1.,"hdiff","Abs Difference "+str(args.titles[0])+"-"+str(args.titles[1]))
    hdiff=xyres[0].Clone("hdiff")
    hdiff.SetTitle("Abs Difference "+str(args.titles[0])+"-"+str(args.titles[1]))
    c3=ROOT.TCanvas("diff","diff")
    c3.SetRightMargin(0.19)
    hdiff.GetZaxis().SetRangeUser(-1000,1000)
    hdiff.GetXaxis().SetTitle("X Pos (cm)")
    hdiff.GetYaxis().SetTitle("Y Pos (cm)")
    hdiff.GetZaxis().SetTitle("Difference (#mu m)")
    hdiff.GetZaxis().SetTitleOffset(1.6)
#    setRange(hdiff)
#    hdiff.Draw("colz")
    c3.Update()
#    reldiff=Divideh(hdiff,xyres[0],"resdiff","Rel Difference")
    reldiff=hdiff.Clone("reldiff")
    reldiff.SetTitle("Rel Difference "+str(args.titles[0])+"-"+str(args.titles[1]))
    reldiff1d=ROOT.TH1D("reldiff1d","Rel Diffeerence 1D "+str(args.titles[0])+"-"+str(args.titles[1])+"/"+str(args.titles[0]),100,-.5,.5)
    for ibi in range(hdiff.GetNbinsX()+1):
        for jbi in range(hdiff.GetNbinsY()+1):
            #cont=hdiff.GetBinContent(ibi,jbi)
            cont=xyres[0].GetBinContent(ibi,jbi)/xyres[1].GetBinContent(ibi,jbi)
            hdiff.SetBinContent(ibi,jbi,cont)
            divider=xyres[0].GetBinContent(ibi,jbi)
            newcont=-9999999999
            if divider > -999 and cont > -999:
                newcont=cont/divider
                reldiff1d.Fill(newcont)
                newcont=abs(newcont)
            reldiff.SetBinContent(ibi,jbi,newcont)
#            print ibi,jbi,cont,divider,newcont
    c3.cd()
    setRange(hdiff,-5,5)
    hdiff.Draw("colz")
    c4=ROOT.TCanvas("reldiff","reldiff",1000,500)
    c4.Divide(2,1)
    c4.GetPad(1).SetRightMargin(0.19)
    c4.cd(1)
    reldiff.GetZaxis().SetRangeUser(0,0.2)
    setRange(reldiff)
    reldiff.Draw("colz")
    c4.cd(2)
    reldiff1d.GetXaxis().SetTitle("Relative Difference")
    reldiff1d.Draw()
    c4.Update()
    c5=ROOT.TCanvas("diffbins","diffbins")
    c6=ROOT.TCanvas("reldiffbins","reldiffbins")
    c7=ROOT.TCanvas("reldiffbins1d","reldiffbins1d")
    c5.Divide(3,2)
    c6.Divide(3,2)
    c7.Divide(3,2)
    hdiffbin=[]
    reldiffbin=[]
    reldiffbin1d=[]
    resdata=[]
    counter=0
    for h in xyslice[0]:
        c5.cd(counter+1)
        hdiffbin.append(Addh(h,xyslice[1][counter],-1,"hdiffbin"+str(counter),"Abs Diff Bin"+str(counter)+" "+str(args.titles[0])+"-"+str(args.titles[1])))
        hdiffbin[-1].GetZaxis().SetRangeUser(-1000,1000)
        hdiffbin[-1].SetStats(0)
        setRange(hdiffbin[-1])
        hdiffbin[-1].Draw("colz")
        reldiffbin.append(hdiff.Clone("reldiffbin"+str(counter)))
        reldiffbin[-1].SetTitle("Rel Diff Bin "+str(counter)+" ("+str(args.titles[0])+"-"+str(args.titles[1])+")/"+str(args.titles[0]))
        reldiffbin1d.append(ROOT.TH1D("reldiffbin1d"+str(counter),"Rel Diffeerence 1D "+str(counter)+" "+str(args.titles[0])+"-"+str(args.titles[1]),100,-5,5))
        for ibi in range(hdiffbin[-1].GetNbinsX()+1):
            for jbi in range(hdiffbin[-1].GetNbinsY()+1):
                cont=hdiffbin[-1].GetBinContent(ibi,jbi)
                divider=h.GetBinContent(ibi,jbi)
                newcont=-9999999999
                if divider > -999 and cont > -999:
                    newcont=cont/divider
                    reldiffbin1d[-1].Fill(newcont)
                    if newcont==1:
                        print ibi,jbi,cont,divider
                    #newcont=abs(newcont)
                reldiffbin[-1].SetBinContent(ibi,jbi,abs(newcont))



        c6.cd(counter+1)
        c6.GetPad(counter+1).SetRightMargin(0.19)
        resdata.append(fitres(reldiffbin1d[-1],0))
        reldiffbin[-1].GetZaxis().SetRangeUser(0,5)
        reldiffbin[-1].Draw("colz")
        c7.cd(counter+1)
        reldiffbin1d[-1].Draw()
        drawresdata(resdata[-1],"mean")
        c5.Update()
        c6.Update()
        c7.Update()
        counter+=1

if args.pdir!="not":
    for c in range(len(canv[0])):
        canv[0][c].SaveAs(args.pdir+"xyresnorm_"+canv[1][c].GetTitle()+".eps")
        canv[1][c].SaveAs(args.pdir+"xyresnormslice_"+canv[2][c].GetTitle()+".eps")
        canv[2][c].SaveAs(args.pdir+"yzresnorm_"+canv[2][c].GetTitle()+".eps")
u=raw_input("Done?")