import ROOT,sys,math,os,time,glob
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
from functions import *
from anaFile import anaFile

#def calcRelDif(h1,h2):
    
def checkInside(rad,z,opt='xy'):
    if opt=='xy':
        if rad>5 and rad<15:
            return True
    elif opt=='yz':
        if rad>-15 and rad<15:# and z>0 and z<73:
            return True
    elif opt=='xyz':
        if rad>5 and rad<15 :#and z>0 and z<73:
            return True
    return False
    
def removeEmptyBins(h,opt='xy'):
    mini= 100000
    maxi=-100000
    
    for ibi in range(0,h.GetNbinsX()+2):
        for jbi in range(0,h.GetNbinsY()+2):
            mini=min(mini,float(h.GetBinContent(ibi,jbi)))
            maxi=max(maxi,float(h.GetBinContent(ibi,jbi)))
            binx=0
            biny=0
            binz=0
            if opt=='xy':
                binx=h.GetXaxis().GetBinCenter(ibi)
                biny=h.GetYaxis().GetBinCenter(jbi)
                rad=math.sqrt(binx**2+biny**2)
                binz=-10000
            if opt=='yz':
                binx=0
                rad=biny=h.GetYaxis().GetBinCenter(jbi)
                binz=h.GetXaxis().GetBinCenter(ibi)
                
            
            if(abs(h.GetBinContent(ibi,jbi))==0 and not checkInside(rad,binz,opt) ):
                h.SetBinContent(ibi,jbi,-100000)
                
            if args.range[0]!=-1:
                if(abs(h.GetBinContent(ibi,jbi))<args.range[0] and checkInside(rad,binz,opt)):
                    h.SetBinContent(ibi,jbi,args.range[0])
            if args.range[1]!=-1:
                if(abs(h.GetBinContent(ibi,jbi))>args.range[1] and checkInside(rad,binz,opt)):
                    h.SetBinContent(ibi,jbi,args.range[1])
                
    mini=-max(abs(mini),maxi)
    maxi=max(abs(mini),maxi)
    if args.range[0]!=-1:
        mini=args.range[0]
    if args.range[1]!=-1:
        maxi=args.range[1]
    return mini,maxi

def normalize(h1,h2):
    newh=ROOT.TH3D()
    newh=h1.Clone()
    newh.Reset()
    newh.SetName(h1.GetName()+"norm")
    newh.SetTitle(h1.GetTitle()+"norm")
    for ibin in range(0,h1.GetSize()):
        if h2.GetBinContent(ibin)!=0:
            newh.SetBinContent(ibin,h1.GetBinContent(ibin)/h2.GetBinContent(ibin))     
        else:
            newh.SetBinContent(ibin,0);         
    return newh

parser=argparse.ArgumentParser(description='plot residuals for comparison of devmaps')
parser.add_argument('file1',help='first file',type=str)
parser.add_argument('file2',help='second file',type=str)
parser.add_argument('--range',help='set bins outside range to maxval (default=%(default)s)',type=float,nargs=2,default=[-1,-1])
parser.add_argument('--data',help='no mc available (default=%(default)s)',action='store_true')
parser.add_argument('--rout',help='root output file (default=%(default)s)',type=str,default='')
parser.add_argument('--tpc',help='use data from tpc detector only, but from combined fit',action="store_true")
parser.add_argument('--tpcOnly',help='use only TPC data from TPC only fit(default=%(default)s)',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")')

    
anaFiles=[]
anaFiles.append(anaFile(args.file1))
anaFiles.append(anaFile(args.file2))

if args.data:
    reslist=['hxyzres','hxyzresY','hxyzresZ']#,'hxyzresU','hxyzresV','hxyzresA1','hxyzresA2','hxyzresA3']
else:
    reslist=['hxyzres','hxyzresY','hxyzresZ','hxyzresU','hxyzresV','hxyzresA1','hxyzresA2','hxyzresA3','hxyzMCres','hxyzMCresY','hxyzMCresZ','hxyzMCresU','hxyzMCresV','hxyzMCresA1','hxyzMCresA2','hxyzMCresA3']

tpcOnly=""
if args.tpcOnly and not args.tpc:
    reslist=[hname+"TpcOnly" for hname in reslist]
    tpcOnly="TpcOnly"

if args.tpc and not args.tpcOnly:
    reslist=[hname+"Tpc" for hname in reslist]
    tpcOnly="Tpc"
    
if args.tpc and args.tpcOnly:
    print "idiot"
    exit()
        
hxyz=[]
hxyz.append(anaFiles[0].getHist("hxyz"+tpcOnly))
hxyz[-1].SetName( hxyz[-1].GetName()+ "_1")
hxyz[-1].SetTitle(hxyz[-1].GetTitle()+"_1")
hxyz.append(anaFiles[1].getHist("hxyz"+tpcOnly))
hxyz[-1].SetName( hxyz[-1].GetName()+ "_2")
hxyz[-1].SetTitle(hxyz[-1].GetTitle()+"_2")

hxyzRebin=[]
hxyzRebin.append(hxyz[0].Rebin3D(2,2,2,hxyz[0].GetName()+"Rebin"))
hxyzRebin.append(hxyz[1].Rebin3D(2,2,2,hxyz[1].GetName()+"Rebin"))

hxy=[]
hxy.append(hxyz[0].Project3D("yx"))
hxy[-1].SetName( hxy[-1].GetName()+ "_1")
hxy[-1].SetTitle(hxy[-1].GetTitle()+"_1")
hxy[-1].SetStats(0)
hxy.append(hxyz[1].Project3D("yx"))
hxy[-1].SetName( hxy[-1].GetName()+ "_2")
hxy[-1].SetTitle(hxy[-1].GetTitle()+"_2")
hxy[-1].SetStats(0)

hyz=[]
hyz.append(hxyz[0].Project3D("yz"))
hyz[-1].SetName( hyz[-1].GetName()+ "_1")
hyz[-1].SetTitle(hyz[-1].GetTitle()+"_1")
hyz[-1].SetStats(0)
hyz.append(hxyz[1].Project3D("yz"))
hyz[-1].SetName( hyz[-1].GetName()+ "_2")
hyz[-1].SetTitle(hyz[-1].GetTitle()+"_2")
hyz[-1].SetStats(0)

hxyzRes=[]
hxyzRes.append({})
hxyzRes.append({})

hxyzResNorm=[]
hxyzResNorm.append({})
hxyzResNorm.append({})

hxyProj=[]
hxyProj.append({})
hxyProj.append({})

hyzProj=[]
hyzProj.append({})
hyzProj.append({})

habsDiff={}
hrelDiff={}


fname=['one','two']
for hname in reslist:
    habsDiff[hname]=ROOT.TH1D("habsDiff"+hname,"Absolute Difference "+hname,200,-1000,1000)
    hrelDiff[hname]=ROOT.TH1D("hrelDiff"+hname,"Relative Difference "+hname,200,-10,10)
    for ifile in range(len(anaFiles)):
        print 'at file',ifile,'histo',hname
        hxyzRes[ifile][hname]=anaFiles[ifile].getHist(hname)
        hxyzRes[ifile][hname].GetXaxis().SetTitle("X(cm)")
        hxyzRes[ifile][hname].GetYaxis().SetTitle("Y(cm)")
        hxyzRes[ifile][hname].GetZaxis().SetTitle("Z(cm)")
        
        hxyzResNorm[ifile][hname]=hxyzRes[ifile][hname].Clone(hname+'norm')
        hxyzResNorm[ifile][hname].Divide(hxyz[ifile])
        hxyzResNorm[ifile][hname].SetTitle(hxyzResNorm[ifile][hname].GetTitle()+" normalized")
        
        hxyProj[ifile][hname]=hxyzRes[ifile][hname].Project3D("yx")
        hxyProj[ifile][hname].Divide(hxy[ifile])
        hxyProj[ifile][hname].SetName(hname+str(ifile)+"XY")
        mini,maxi=removeEmptyBins(hxyProj[ifile][hname])
        hxyProj[ifile][hname].GetZaxis().SetRangeUser(mini,maxi)
        hxyProj[ifile][hname].SetStats(0)
        
        hyzProj[ifile][hname]=hxyzRes[ifile][hname].Project3D("yz") 
        hyzProj[ifile][hname].Divide(hyz[ifile])                    
        hyzProj[ifile][hname].SetName(hname+str(ifile)+"YZ")        
        mini,maxi=removeEmptyBins(hyzProj[ifile][hname],'yz')
        hyzProj[ifile][hname].GetZaxis().SetRangeUser(mini,maxi)
        hyzProj[ifile][hname].SetStats(0)
        
        hxyzRes[ifile][hname].Rebin3D(2,2,2)
        hxyzResNorm[ifile][hname]=hxyzRes[ifile][hname].Clone(hname+'norm')
        hxyzResNorm[ifile][hname].Divide(hxyzRebin[ifile])
        hxyzResNorm[ifile][hname].SetTitle(hxyzResNorm[ifile][hname].GetTitle()+" normalized")
        
        
    print 'getting differences'
    for xbin in range(hxyzResNorm[0][hname].GetNbinsX()):
        #print 'at xbin',xbin
        for ybin in range(hxyzResNorm[0][hname].GetNbinsY()):
            for zbin in range(hxyzResNorm[0][hname].GetNbinsZ()):
                binx=hxyzResNorm[0][hname].GetXaxis().GetBinCenter(xbin)
                biny=hxyzResNorm[0][hname].GetYaxis().GetBinCenter(ybin)
                binz=hxyzResNorm[0][hname].GetZaxis().GetBinCenter(zbin)
                rad=math.sqrt(binx**2+biny**2)
                if checkInside(rad,binz,'xyz'):
                    if hxyzResNorm[0][hname].GetBinContent(xbin,ybin,zbin)!=0:
                        habsDiff[hname].Fill(hxyzResNorm[0][hname].GetBinContent(xbin,ybin,zbin)-hxyzResNorm[1][hname].GetBinContent(xbin,ybin,zbin))
                        hrelDiff[hname].Fill((hxyzResNorm[0][hname].GetBinContent(xbin,ybin,zbin)-hxyzResNorm[1][hname].GetBinContent(xbin,ybin,zbin))/abs(hxyzResNorm[0][hname].GetBinContent(xbin,ybin,zbin)))
                    
        
u=raw_input("start plotting?")
cxyProj={}
cyzProj={}
cDiff={}
if args.rout!="":
    outfile=ROOT.TFile(args.rout,"recreate")
    
for hname in reslist:
    cxyProj[hname]=(ROOT.TCanvas("cProj_XY_"+hname,"XY "+hname,1000,500))
    cxyProj[hname].Divide(len(anaFiles),1)
    
    cyzProj[hname]=ROOT.TCanvas("cProj_YZ_"+hname,"YZ "+hname,1000,500)
    cyzProj[hname].Divide(len(anaFiles),1)
    
    cDiff[hname]=ROOT.TCanvas("cDiff"+hname,"Difference "+hname,1000,500)
    cDiff[hname].Divide(2,1)
    
    for ifile in range(len(anaFiles)):
        cxyProj[hname].cd(ifile+1)
        hxyProj[ifile][hname].Draw('colz')
        
        cyzProj[hname].cd(ifile+1)
        hyzProj[ifile][hname].Draw('colz')
        
        if args.rout:
            hxyProj[ifile][hname].Write()
            hyzProj[ifile][hname].Write()
            
    
    cDiff[hname].cd(1)
    habsDiff[hname].Draw()
    cDiff[hname].cd(2)
    hrelDiff[hname].Draw()
    
    if args.rout:
        habsDiff[hname].Write()
        hrelDiff[hname].Write()
        

thisIsTheEnd()
if args.rout!="":
    outfile.Close()