import ROOT,glob,math,argparse,sys
sys.path.append('$PANDAPATH/macro/tpc/FOPI/mberger')
from anaFile import anaFile
from functions import *

def setGraph(g,title,col=1):
    g.SetTitle(title)
    g.SetMarkerStyle(20)
    g.SetMarkerSize(1)
    g.SetMarkerColor(col)
    return

def fillGraph(hist,graph,fnum=0):
    nbins=hist.GetNbinsX()
    projs=[]
    graph.append({})
    graph[-1]['RMS']=ROOT.TGraph()
    setGraph(graph[-1]['RMS'],'RMS',fnum+1)
    graph[-1]['Mean']=ROOT.TGraph()
    setGraph(graph[-1]['Mean'],'Mean',fnum+1)

    pcounter=0
    for i in range(args.binwidth,nbins,args.binwidth):
        projs.append(h.ProjectionY(h.GetName()+str(i),i-args.binwidth,i))
        graph[-1]['RMS'].SetPoint(pcounter,i,projs[-1].GetRMS())
        graph[-1]['Mean'].SetPoint(pcounter,i,projs[-1].GetMean())
        pcounter+=1
    return projs

parser=argparse.ArgumentParser(description="calculate a parametrization for the cluster error")
parser.add_argument('anafiles',help='the file with the histograms',type=str,nargs='+')
parser.add_argument('--binwidth',help='width of bins',type=int,default=5)
parser.add_argument('--drawslices',help='draw slices or not?',action='store_true')
parser.add_argument('--titles',help='the titles',type=str,nargs='+',default=["nothing special"])
parser.add_argument('--fC',help='the fC parameter',type=float,default=10)
parser.add_argument('--fG',help='the fG parameter',type=float,default=600)
parser.add_argument('--gain',help='the gain',type=float,default=3771)
parser.add_argument('--hlp',help='print help',action='store_true')
args=parser.parse_args()

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

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

rfile=[]
cluster_MC_res=[]
cluster_MC_res_graph=[]
cluster_MC_resTh=[]
cluster_Err=[]
cluster_Err_graph=[]
errx=ROOT.TF1('errx',)
for ifile in range(len(args.anafiles)):
    rfile.append(anaFile(args.anafiles[ifile]))
    cluster_MC_res.append(rfile[-1].getHistsByPattern("hClusterMCRes","Theta"))
    cluster_MC_res[-1].sort()
    cluster_MC_res_graph.append([])
    cluster_MC_resTh.append(rfile[-1].getHistsByPattern("hClusterMCResvsTheta"))
    cluster_MC_resTh[-1].sort()
    cluster_Err.append(rfile[-1].getHistsByPattern("hpull_cl_err"))
    cluster_Err[-1].sort()
    cluster_Err_graph.append([])

cMC_res       = ROOT.TCanvas("cMC_res","MC-Res Z",700,1000)
cMC_res_graph = ROOT.TCanvas("cMC_res_graph","MC-Res Z graph",1400,1000)
cMC_res_slices= []
cMC_resTh     = ROOT.TCanvas("cMC_resTh","MC-Res Theta",700,1000)

ctmp=ROOT.TCanvas()

cMC_res.Divide(1,4)
cMC_res_graph.Divide(2,4)
cMC_resTh.Divide(1,4)
cc=0
cluster_MC_res_proj=[]
fcounter=-1
leg=ROOT.TLegend(0.7,0.75,0.9,0.9)
leg.SetFillColor(0)
leg.SetTextSize(0.05)

for MC_res in cluster_MC_res:
    fcounter+=1
    cc=0
    cluster_MC_res_proj.append([])
    for h in MC_res:
        cc+=1
        nbins=h.GetNbinsX()
        cluster_MC_res_proj[-1].append(fillGraph(h,cluster_MC_res_graph[fcounter],fcounter))
        cMC_res.cd(cc)
        h.Draw("colz")
        if args.drawslices:
            cMC_res_slices.append(ROOT.TCanvas("cMC_res_slice"+str(cc),h.GetName()+"slice"+str(cc),1000,1000))
            sqdiv=int(math.sqrt(len(cluster_MC_res_proj[cc-1]))+1)
            cMC_res_slices[-1].Divide(sqdiv,sqdiv)
            scc=0
            for hp in cluster_MC_res_proj[cc-1]:
                scc+=1
                cMC_res_slices[-1].cd(scc)
                hp.Draw()
                cMC_res_slices[-1].Update()

        opt="P"
        if fcounter==0:
            opt="AP"
            
        if cc==1:
            leg.AddEntry(cluster_MC_res_graph[fcounter][cc-1]["RMS"],args.titles[fcounter],"P")

        cMC_res_graph.cd(cc*2-1)
        cluster_MC_res_graph[fcounter][cc-1]["RMS"].Draw(opt)
        if fcounter==len(cluster_MC_res)-1:
            leg.Draw()

        cMC_res_graph.cd(cc*2)
        if fcounter==len(cluster_MC_res)-1:
            leg.Draw()
        cluster_MC_res_graph[fcounter][cc-1]["Mean"].Draw(opt)
    

cErr=[]
lErr=[]
for i in range(len(args.anafiles)):
    cErr.append(ROOT.TCanvas("cErr"+args.titles[i],"Cluster Error and MC res "+str(args.titles[i]),1400,1000))
    cErr[-1].Divide(2,4)
    lErr.append(ROOT.TLegend(0.7,0.75,0.9,0.9))

fcounter=-1
cluster_Err_proj=[]
for Err in cluster_Err:
    fcounter+=1
    cc=0
    cluster_Err_proj.append([])
    for h in Err:
        cc+=1
        cluster_Err_proj[-1].append(fillGraph(h,cluster_Err_graph[fcounter],fcounter))

        cErr[fcounter].cd(cc*2-1)
        cluster_Err_graph[fcounter][cc-1]["RMS"].SetMarkerStyle(22)
        cluster_MC_res_graph[fcounter][cc-1]["RMS"].Draw("AP")
        cluster_Err_graph[fcounter][cc-1]["RMS"].Draw("P")

        cErr[fcounter].cd(cc*2)
        cluster_Err_graph[fcounter][cc-1]["Mean"].SetMarkerStyle(22)
        cluster_MC_res_graph[fcounter][cc-1]["Mean"].Draw("AP")
        cluster_Err_graph[fcounter][cc-1]["Mean"].Draw("P")

        cErr[fcounter].Update()

cc=0
cMC_resTh.cd()
for MC_resTh in cluster_MC_resTh:
    for h in MC_resTh:
        cc+=1
        cMC_resTh.cd(cc)
        h.Draw("colz")
    

cMC_res.Update()
cMC_resTh.Update()
u='Null'
while u!='q':
    u=raw_input("q for quit")