import sys
sys.path.append('/home/vwalbr/fopiroot/fopiroot_dev/macro/tpc/FOPI/vwalbre')
import ROOT,glob,math,argparse
from anaFile import anaFile
from functions import *
from histo_plot import *
#from histo_plot_TEST   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

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("--outfname",help='name of the output file',type=str,default="")
parser.add_argument('--resY',help='plot Y res',action='store_true')
parser.add_argument('--res3D',help='plot 3D res',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)



  

if args.diff:
    c3=ROOT.TCanvas("diff","diff")
    c4=ROOT.TCanvas("reldiff","reldiff",1000,500)
    c5=ROOT.TCanvas("diffbins","diffbins")
    c6=ROOT.TCanvas("reldiffbins","reldiffbins")
    c7=ROOT.TCanvas("reldiffbins1d","reldiffbins1d")
if args.resY:
    c18=ROOT.TCanvas("diff Y","diff Y")
    c19=ROOT.TCanvas("reldiff Y","reldiff Y",1000,500)
if args.res3D:
    c21=ROOT.TCanvas("reldiff 3D","reldiff 3D",1000,500)
    c22=ROOT.TCanvas("Proj reldiff 3D ZY","Proj reldiff 3D ZY",800,400)
    c23=ROOT.TCanvas("Fit","Fit",700,900)
    c24=ROOT.TCanvas("Proj reldiff 3D ZY - CUT ","Proj reldiff 3D ZY - CUT ",800,400)
    c25=ROOT.TCanvas("Proj reldiff 3D XY","Proj reldiff 3D XY",600,500)
    c26=ROOT.TCanvas("Proj reldiff 3D XY - CUT","Proj reldiff 3D XY - CUT",600,500)
    
 
    




rfiles=[]
dslices=[]

xyslice=[]
xyphislice=[]


xyres=[]
xyzres=[]
yzres=[]

xy=[]
xyz=[]
xySlice=[]
xyzSlice=[]

xreshistZ=[]
xreshist=[]

canv=[[],[],[]]#*2

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("hxyMCresYnorm"))
        yzres.append(rfiles[-1].getHist("hyzMCresYnorm"))
        xyslice.append(rfiles[-1].getSlicedHist("hxyMCresYnorm"))
      
    else:
        xyres.append(rfiles[-1].getHist("hxyresYnorm"))
        yzres.append(rfiles[-1].getHist("hyzresYnorm"))
        xy.append(rfiles[-1].getHist("hxy"))
        xyz.append(rfiles[-1].getHist("hxyz"))
        if args.res3D:
            xyzres.append(rfiles[-1].getHist("hxyzresnormY"))
        xyslice.append(rfiles[-1].getXYsliceHistsY())

    canv[0].append(ROOT.TCanvas())
    canv[1].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


    set_titles(xyres[fcounter])
    set_titles_hdiffyz(yzres[fcounter])
    #if args.res3D:
    #    set_titles3D(xyzres[fcounter])


    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")
    
    if args.titles[0]!="":
        canv[1][-1].SetTitle(args.titles[fcounter]+" X")

    canv[1][-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.SetStats(0)
        h.Draw("colz")
        counter+=1



    counter=0
    
    for c in canv:
        c[-1].Update()

if args.diff:
    # X-RESIDUALS
    hdiff=xyres[0].Clone("hdiff")
    hdiff.SetTitle("Abs Difference "+str(args.titles[0])+"-"+str(args.titles[1]))
    c3.SetRightMargin(0.19)
    set_titles_hdiff(hdiff)
    c3.Update()
    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[1])+"|",100,-10,10)

    histos = prod_reldiff(hdiff, reldiff, reldiff1d, xyres, xy)
 
    divROOT=xyres[0].Clone("divROOT")
    divROOT.Divide(xyres[1])

    c3.cd()
    setRange(hdiff,-5,5)
    hdiff.Draw("colz")
    c4.Divide(2,1)
    c4.GetPad(1).SetRightMargin(0.19)
    c4.cd(1)
    #reldiff.GetZaxis().SetRangeUser(0,0.2)
    setRange(histos[0])
    histos[0].Draw("colz")
    c4.cd(2)
    resdataX1d=[]
    resdataX1d.append(fitres(histos[1],0))
    histos[1].GetXaxis().SetTitle("Relative Difference")
    histos[1].Draw()
    drawresdata(resdataX1d[-1],"mean")
    
    c4.Update()
    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[1]))
        reldiffbin1d.append(ROOT.TH1D("reldiffbin1d"+str(counter),"Rel Diffeerence 1D "+str(counter)+" "+str(args.titles[0])+"-"+str(args.titles[1]),100,-10,10))
        
        histos_slice=prod_reldiff_slice(counter, hdiffbin, xyslice, reldiffbin, reldiffbin1d)
        c6.cd(counter+1)
        c6.GetPad(counter+1).SetRightMargin(0.19)
        resdata.append(fitres(histos_slice[1][-1],0))
        histos_slice[0][-1].GetZaxis().SetRangeUser(0,5)
        histos_slice[0][-1].Draw("colz")
        c7.cd(counter+1)
        histos_slice[1][-1].Draw()
        drawresdata(resdata[-1],"mean")
        c5.Update()
        c6.Update()
        c7.Update()
        counter+=1          

    if args.res3D:
        hdiff3D=xyzres[0].Clone("hdiff3D")
        hdiff3D.SetTitle("Abs Difference "+str(args.titles[0])+"-"+str(args.titles[1]))
        set_titles_hdiff3D(hdiff3D)
        reldiff3D=hdiff3D.Clone("reldiff3D")
        reldiff3D.SetTitle("Rel Difference "+str(args.titles[0])+"-"+str(args.titles[1]))
        
        reldiff1d3D=ROOT.TH1D("reldiff1d 3D","Rel Diffeerence 3D "+str(args.titles[0])+"-"+str(args.titles[1])+"/"+"|"+str(args.titles[1])+"|",400,-5,5)

        histos,RelDiffyx, RelDiffCUTyx,RelDiffyz, RelDiffCUTyz= prod_reldiff3D(hdiff3D, reldiff3D, reldiff1d3D, xyzres,xyz)
        FitPar,c23=RooFit(histos[1],args.outfname,c23)
        print "*****************", args.outfname
        divROOT=xyzres[0].Clone("divROOT")
        divROOT.Divide(xyzres[1])

        c21.Divide(2,1)
        c21.GetPad(1).SetRightMargin(0.19)
        c21.cd(1)
        setRange(histos[0])
        histos[0].Draw("colz")
        c21.cd(2)
        histos[1].GetXaxis().SetTitle("Relative Difference")
        histos[1].Draw()
        c21.Update()
        
        Nran=-1000
        Pran=1000

        
        c25.cd()
        setRange(RelDiffyx,Nran,Pran)
        set_titles_hdiff(RelDiffyx)
        c25.SetRightMargin(0.19)
        RelDiffyx.Draw("colz")
        c25.Update()

        c26.cd()
        setRange(RelDiffCUTyx,Nran,Pran)
        set_titles_hdiff(RelDiffCUTyx)
        RelDiffCUTyx.Draw("colz")
        c26.SetRightMargin(0.19)
        c26.Update()

        c22.cd()
        setRange(RelDiffyz,Nran,Pran)
        set_titles_hdiffyz(RelDiffyz)
        c22.SetRightMargin(0.19)
        RelDiffyz.Draw("colz")
        c22.Update()

        c24.cd()
        setRange(RelDiffCUTyz,Nran,Pran)
        set_titles_hdiffyz(RelDiffCUTyz)
        c24.SetRightMargin(0.19)
        RelDiffCUTyz.Draw("colz")
        c24.Update()


        

        
       
 
 

#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")

if args.pdir!="not" and args.outfname!="":
    c23.SaveAs(args.pdir+args.outfname[:-10]+".pdf"+'(')
    c3.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c4.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c5.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c6.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c7.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c21.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c22.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    for c in range(len(canv[0])):
        canv[0][c].SaveAs(args.pdir+args.outfname[:-10]+".pdf")
        canv[1][c].SaveAs(args.pdir+args.outfname[:-10]+".pdf")
        canv[2][c].SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c24.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c25.SaveAs(args.pdir+args.outfname[:-10]+".pdf")
    c26.SaveAs(args.pdir+args.outfname[:-10]+".pdf"+')')
    #c23.SaveAs(args.pdir+args.outfname[:-10]+".pdf"+')')

if args.res3D:
    if args.outfname!="":
        FILE=open("vw_pics/FieldCorr/Files/"+args.outfname, "w")
        FILE.write(str(FitPar[0])+ " " + str(FitPar[1]) + " " +str(FitPar[2]) + " " + str(FitPar[3]) + " " + str(FitPar[4]))
        FILE.close()
        print "Fit Parameters were written into:", args.outfname


if args.outfname == "":
    u=raw_input("Done?")