#Description: draws BiCubSpline from a reconstruction root tree 
# (output of TpcLaserFitTask)
#Author: Felix Boehmer, E18

import ROOT, glob, math, sys, os
from ROOT import std
from array import array

options = ["-f", "-name","-bw"]
opts = set(options)

def getArgsToNextOpt(stringlist, index) :
    id = index+1
    args = []
    while(id < len(stringlist)):
        if stringlist[id] in options:
            return args
        args.append(stringlist[id])
        id+=1
    return args    

def getLast(pline,split) :
    chunks = pline.split(split)
    return chunks[len(chunks)-1:][0]

#default parameters
pandapath = os.environ['PANDAPATH']
print pandapath
pfilename = pandapath+"/tpc/parfiles/tpc.par"
splinebranch = "TpcLaserSplineFits"

#get geometry
#geometry:
pfile = open(pfilename, "r")

BW = 0   #plot in black and white

for line in pfile:
    if line.startswith("zGem") :
        zMin = (float)(getLast(line,' ').rstrip())
    if line.startswith("zMax") :
        zMax = (float)(getLast(line,' ').rstrip())
    if line.startswith("rMin") :
        rMin = (float)(getLast(line,' ').rstrip())
    if line.startswith("rMax") :
        rMax = (float)(getLast(line,' ').rstrip())    
        
pfile.close()

for iarg in range(len(sys.argv)) :
    arg = sys.argv[iarg]
    if arg == "-f":  #file containing E-field (ASCII)
        efieldfs = getArgsToNextOpt(sys.argv,iarg)
        if len(efieldfs) > 1 :
            print "Invalid argument to -ef!"
            exit
        filename = efieldfs[0]
    if arg == "-name":  #name of TF2 in file
        tf2name = getArgsToNextOpt(sys.argv,iarg)
        if len(tf2name) > 1 :
            print "Invalid argument to -name!"
            exit
        name = tf2name[0]    
    if arg == "-bw":
        BW = 1

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
if BW :
    rlogon = pandapath+"/rootlogon_Bernhard_New_BW.C"
else :
    rlogon = pandapath+"/rootlogon_Bernhard_New.C"

ROOT.gROOT.ProcessLine(".x "+rlogon)
if BW:
    ROOT.gROOT.ProcessLine('gROOT->SetStyle("bw")'); 
else :
    ROOT.gROOT.ProcessLine('gROOT->SetStyle("col")'); 
ROOT.gROOT.ProcessLine('gROOT->ForceStyle()'); 

nbinsr = 500
nbinsz = 2000


rfile = ROOT.TFile(filename, "READ")
tree = rfile.Get("cbmsim") 
tree.SetBranchStatus("*", 0)
tree.SetBranchStatus(splinebranch+".*", 1)

for e in tree:  #should be just one event normally
    #for spl in e.TpcLaserSplineFits:
    #    spl.getCoeffs().Print()
    splines = e.TpcLaserSplineFits
    recoRad = splines.At(0)
    recoPhi = splines.At(1)
    ifcRad = recoRad.getTF2(zMin,zMax,rMin,rMax)
    ifcPhi = recoPhi.getTF2(zMin,zMax,rMin,rMax)
    

histRad = ROOT.TH2D("histRad","dasd",nbinsz,zMin,zMax,nbinsr,rMin,rMax)
histRad.GetXaxis().SetTitle("z (cm)")
histRad.GetYaxis().SetTitle("r (cm)")
histPhi = ROOT.TH2D("histPhi","dasd",nbinsz,zMin,zMax,nbinsr,rMin,rMax)
histPhi.GetXaxis().SetTitle("z (cm)")
histPhi.GetYaxis().SetTitle("r (cm)")


rbinwidth = (rMax-rMin)/(nbinsr)
zbinwidth = (zMax-zMin)/(nbinsz)

for iz in range(nbinsz) :
    for ir in range(nbinsr) :
        z = (iz+0.5)*zbinwidth + zMin
        r = (ir+0.5)*rbinwidth + rMin
        val = ifcRad.Eval(z,r)
        histRad.SetBinContent(iz+1,ir+1,val)
        val = ifcPhi.Eval(z,r)
        histPhi.SetBinContent(iz+1,ir+1,val)


canv1 = ROOT.TCanvas("dsf1","DSD",800,400)
histRad.Draw("colz")
ROOT.gPad.Update()   #REQUIRED for the next line to return a valid object
pal = histRad.FindObject("palette")
pal.GetAxis().SetTitle('#xi_{rad, reco} (cm)')
pal.GetAxis().SetTitleOffset(0.65)
histRad.Draw("colz")

canv2 = ROOT.TCanvas("dsf2","DSD",800,400)
histPhi.Draw("colz")
ROOT.gPad.Update()   #REQUIRED for the next line to return a valid object
pal = histPhi.FindObject("palette")
pal.GetAxis().SetTitle('#xi_{#phi, reco} (cm)')
pal.GetAxis().SetTitleOffset(0.65)
histPhi.Draw("colz")


out = ROOT.TFile("tmp.root", "recreate")
histRad.Write()
histPhi.Write()

input()