import sys, os, copy, math
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, ROOT
from functions import set_palette, openTree, thisIsTheEnd, getSlices, setGraph


def drawProj(h,col,plist):
    h.FitSlicesY(0,0,-1,0,"QNRSG3")
    h2=ROOT.gDirectory.Get(h.GetName()+"_2")
    if col>=10:
        col+=1
    h2.SetLineColor(col)
    if col==0:
        h2.Draw()
    else:
        h2.Draw("same")
    plist.append(h2)
    
def plotHist2(c,c_sig,hlist,plist,flist,fgraph,name,i):
    hlist.append(infiles[i].Get(name))
    hlist[-1].SetName('{0}_{1}'.format(hlist[-1].GetName(),i))
    c.cd(i+1)
    hlist[-1].Draw('colz')
    c_sig.cd()
    drawProj(hlist[-1],i+1,plist)
    fitprojection(plist,flist,fgraph)
    
def plotHist(c,c_sig,hlist,plist,name,i):
    hlist.append(infiles[i].Get(name))
    hlist[-1].SetName('{0}_{1}'.format(hlist[-1].GetName(),i))
    c.cd(i+1)
    hlist[-1].Draw('colz')
    c_sig.cd()
    drawProj(hlist[-1],i+1,plist)

       
def fitprojection(plist,flist,graphP):

    flist.append(ROOT.TF1("fitfunc{0}".format(plist[-1].GetName()),"pol1",25,60))
    flist[-1].SetLineColor(plist[-1].GetLineColor())
    plist[-1].Fit(flist[-1],"QRN+")
    theta=float(fileList[len(plist)-1][1].split('_')[-1])
    graphP[0].SetPoint(graphP[0].GetN(),theta,flist[-1].GetParameter(0))
    graphP[0].SetPointError(graphP[0].GetN()-1,0,flist[-1].GetParError(0))
    graphP[1].SetPoint(graphP[1].GetN(),theta,flist[-1].GetParameter(1))
    graphP[1].SetPointError(graphP[1].GetN()-1,0,flist[-1].GetParError(1))
               
 

parser=argparse.ArgumentParser(description='plot residuals as function of various variablers')
parser.add_argument('filelist',help='folder or file',type=str)
args=parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")
set_palette()

fileList=[]
for line in open(args.filelist):
    words=line.split(';')
    fileList.append(words)


htheta=[[],[]]
htheta_sig=[[],[]]
ctheta=[ROOT.TCanvas('ctheta1','Theta A1',1000,1000),ROOT.TCanvas('ctheta2','Theta A2',1000,1000)]
ctheta[0].DivideSquare(len(fileList))
ctheta[1].DivideSquare(len(fileList))
ctheta_sig=[ROOT.TCanvas('ctheta_sig1','Sigma Theta A1',1000,1000),ROOT.TCanvas('ctheta_sig2','Sigma Theta A2',1000,1000)]
ftheta=[[],[]]
cthetaFit=[ROOT.TCanvas('cthetaFitA1','Theta Fit A1',1000,1000),ROOT.TCanvas('cthetaFitA2','Theta Fit A2',1000,1000)]
for c in cthetaFit:
    c.Divide(2)
fitGraphTheta=[[],[]]
for fl in fitGraphTheta:
    fl.append(ROOT.TGraphErrors())
    fl.append(ROOT.TGraphErrors())


hthetaXYZ=[[],[],[]]
hthetaXYZ_sig=[[],[],[]]
cthetaXYZ=[ROOT.TCanvas('cthetaX','Theta X',1000,1000),ROOT.TCanvas('cthetaY','Theta Y',1000,1000),ROOT.TCanvas('cthetaZ','Theta Z',1000,1000)]
cthetaXYZ[0].DivideSquare(len(fileList))
cthetaXYZ[1].DivideSquare(len(fileList))
cthetaXYZ[2].DivideSquare(len(fileList))
cthetaXYZ_sig=[ROOT.TCanvas('ctheta_sigX','Sigma Theta X',1000,1000),ROOT.TCanvas('ctheta_sigY','Sigma Theta Y',1000,1000),ROOT.TCanvas('ctheta_sigZ','Sigma Theta Z',1000,1000)]
fthetaXYZ=[[],[],[]]
cthetaFitXYZ=[ROOT.TCanvas('cthetaFitX','Theta Fit X',1000,1000),ROOT.TCanvas('cthetaFitY','Theta Fit Y',1000,1000),ROOT.TCanvas('cthetaFitZ','Theta Fit Z',1000,1000)]
for c in cthetaFitXYZ:
    c.Divide(2)
fitGraphThetaXYZ=[[],[],[]]
for fl in fitGraphThetaXYZ:
    fl.append(ROOT.TGraphErrors())
    fl.append(ROOT.TGraphErrors())
    
    
ctheta_leg=ROOT.TCanvas('tleg','Theta Legend',500,1000)
tleg=ROOT.TLegend(0,0,1,1)
tleg.SetFillColor(0)

infiles=[]
for i in range(len(fileList)):
    print fileList[i][0]
    infiles.append(ROOT.TFile(fileList[i][0],'read'))
    
    names=['A1','A2']
    for j in range(2):
        #plotHist(ctheta[j],ctheta_sig[j],htheta[j],htheta_sig[j],'res vs Z/hClusterMCRes{0}vsZ'.format(names[j]),i)
        plotHist(ctheta[j],ctheta_sig[j],htheta[j],htheta_sig[j],'pulls/hpull_cl_a{0}VSz'.format(j+1),i)
        fitprojection(htheta_sig[j],ftheta[j],fitGraphTheta[j])
        ctheta_sig[j].cd()
        ftheta[j][-1].Draw("same")
    #plotHist(ctheta[1],ctheta_sig[1],htheta[1],htheta_sig[1],'res vs Z/hClusterMCResA2vsZ',i)
    
    #names=['X','Y','Z']
    names=['x','y','z']
    for j in range(3):
        #plotHist(cthetaXYZ[j],cthetaXYZ_sig[j],hthetaXYZ[j],hthetaXYZ_sig[j],'res vs Z/hClusterMCRes{0}vsZ'.format(names[j]),i)
        plotHist(cthetaXYZ[j],cthetaXYZ_sig[j],hthetaXYZ[j],hthetaXYZ_sig[j],'pulls/hpull_cl_{0}VSz'.format(names[j]),i)
        fitprojection(hthetaXYZ_sig[j],fthetaXYZ[j],fitGraphThetaXYZ[j])
        cthetaXYZ_sig[j].cd()
        fthetaXYZ[j][-1].Draw("same")

    tleg.AddEntry(htheta_sig[0][-1],fileList[i][1],'lp')
    
for j in range(3):
    cthetaFitXYZ[j].cd(1)
    fitGraphThetaXYZ[j][0].Draw('AP')       
    cthetaFitXYZ[j].cd(2)
    fitGraphThetaXYZ[j][1].Draw('AP')       

for j in range(2):
    cthetaFit[j].cd(1)
    fitGraphTheta[j][0].Draw('AP')       
    cthetaFit[j].cd(2)
    fitGraphTheta[j][1].Draw('AP')       
    
ctheta_leg.cd()
tleg.Draw()
ctheta[0].Update()
ctheta[1].Update()
ctheta_sig[0].Update()
ctheta_sig[1].Update()
for c in cthetaFitXYZ:
    c.Update()
for c in cthetaFit:
    c.Update()
thisIsTheEnd()