from ROOT import gROOT, TFile, gStyle
from ROOT import kFALSE, kTRUE
from ROOT import RooPlot, RooHist
from ROOT import TPaveText, TGraph, TCanvas, TLatex, TF1, TH1D, gPad
from ROOT import kBlack, kBlue, kRed

from math import exp,log

import collections #To select subsets of dictionaries


verbose = True

def isMatched(name):
    return ("matched" in name)

def isFull(name):
    return ("full" in name)

def rangeMass(method):
    if (method == "Long"): return [2625.0,3575.0]
    if (method == "Velo"): return [2925.0,3275.0]
    if (method == "T"):    return [2625.0,3575.0]

def binEdges(name):
    nameTmp = name.replace(" ","") #Remove spaces, they are useless 
    #split makes a list, first part we throw away
    #Then we split the string of binLow<VAR<binHigh
    nameList =  (nameTmp.split(',')[1]).split('<') 
    return (nameList[0],nameList[2]) #We return the first and the third one

def methodTex(method):
    method_tex = TLatex()
    method_tex.SetNDC(True)
    method_tex.SetTextFont(132)
    method_tex.SetTextSize(0.06)
    method_tex.SetTextAlign(13)
    method_tex.DrawLatex(0.205, 0.905, method+" method")
    return 

def mainTag():
    tex = TLatex()
    tex.SetNDC(True)
    tex.SetTextFont(132)
    tex.SetTextSize(0.06)
    tex.SetTextAlign(33)  
    tex.DrawLatex(0.875, 0.905, "This thesis")

def ConvertRooHistToTH1D(hist):
    nBins = hist.GetN()
    newHist = TH1D(hist.GetName(),hist.GetTitle(),nBins, hist.GetXaxis().GetXmin(),hist.GetXaxis().GetXmax())
    for bin in range(nBins):
        #print(bin,hist.GetPointY(bin))
        newHist.SetBinContent(bin+1,hist.GetPointY(bin))
    print(newHist.GetEntries())        
    

    return newHist    


def parsePaveText(text, param):
    parStr = text.GetLineWith(param).GetTitle()
    parStr.replace("  "," ") #Replace double spaces by one to avoid trouble
    parList = parStr.split(" ")
    return (float(parList[-3]),float(parList[-1])) #the list is then [mess, mess, something, val,#pm,err]


    

def Plot(simVer = "Sim09h", year = "2018_25ns", data = True,
         method = "Long", var = "ETA", 
         outputPath = "./prettyPlots"
         ):
    gROOT.SetBatch(kTRUE)
    gStyle.SetTextFont(132)
	
    dataTag = "Data" if data else "MC"
    folderName = "./results/"+year+"_WG/"
    if (not data): folderName = fileName + "/" +simVer + "/"
    fileName = "trackEff_" + dataTag + "_" + var + "_" + method + "_method.root"


    canvas_list = ["matched", "full", "fail"]
    inputfile = TFile(folderName+fileName)
    print ("Openning", folderName+fileName)
    
    frameDictList = []

    #Get the frames
    for item in inputfile.GetListOfKeys():
        #if any(cvs in item.GetName() for cvs in canvas_list):
        if (item.GetClassName() != "RooPlot"): continue
        if (verbose): print("Got item " + item.GetName() + " of class " + item.GetClassName())        
        frame = RooPlot()
        frame = inputfile.Get(item.GetName())
        name = frame.GetTitle()
        hist = RooHist()
        graph = TGraph()
        text = TPaveText()

        hist = frame.getObject(0)
        graph = frame.getObject(1)
        text = frame.getObject(2)

        frameDict ={"frame": frame,
                    "matched": isMatched(name),
                    "bins": (0,0) if (var=="Full") else binEdges(name),
                    "hist": hist,
                    "graph": graph,
                    "text": text
        }
        frameDictList.append(frameDict)
      
    #Sort into dictionaries based on the same bin
    result = collections.defaultdict(list)  
    for frameDict in frameDictList:
        result[frameDict['bins']].append(frameDict)
    
    sortedList = list(result.values())
    #Check there are only two entries per list
    for frDictList in sortedList:
        if (len(frDictList)!=2): 
            print("ERROR: somehow there are more/less than two frames per bins, check.")
            return


    #Now loop over the bins and plot each function+data
    for i,frDictList in enumerate (sortedList):
        #create canvas
        canvas = TCanvas("c_"+method, "c_"+method, 10,10,650,600)
        canvas.SetBottomMargin(0.15)
        canvas.SetRightMargin(0.05)
        canvas.SetLeftMargin(0.15)   
        canvas.SetTopMargin(0.0675)        

        #Draw the hist
        h_all=ConvertRooHistToTH1D(frDictList[1]['hist'])

        frDictList[1]['hist'].GetXaxis().SetRangeUser(rangeMass(method)[0],rangeMass(method)[1])
        frDictList[1]['hist'].GetYaxis().SetRangeUser(0,h_all.GetMaximum()*1.25)
    
        frDictList[1]['hist'].SetLineWidth(2)
        frDictList[1]['hist'].SetTitle("")

        #h_all.GetXaxis().SetTitleFont(13)
        #h_all.GetXaxis().SetTitleSize(0.05)
        #h_all.GetXaxis().SetTitle("m_{#mu^{+}#mu^{-}} [MeV]")
        #h_all.GetYaxis().SetTitleFont(13)
        #h_all.GetYaxis().SetTitleSize(0.05)
        #h_all.GetYaxis().SetTitle("Entries (a.u.)")
        frDictList[1]['hist'].GetXaxis().SetTitle("m_{#mu^{+}#mu^{-}} [MeV]")
        frDictList[1]['hist'].GetXaxis().SetTitleOffset(1.1)
        frDictList[1]['hist'].GetXaxis().SetTitleFont(132)
        frDictList[1]['hist'].GetXaxis().SetLabelSize(0.05)
        frDictList[1]['hist'].GetXaxis().SetLabelFont(132)
        frDictList[1]['hist'].GetXaxis().SetTitleSize(0.06)
        frDictList[1]['hist'].GetXaxis().SetNdivisions(505)
        frDictList[1]['hist'].GetYaxis().SetTitle("Entries (a.u.)")
        frDictList[1]['hist'].GetYaxis().SetTitleFont(132)
        frDictList[1]['hist'].GetYaxis().SetLabelFont(132)
        frDictList[1]['hist'].GetYaxis().SetTitleOffset(1.25)
        frDictList[1]['hist'].GetYaxis().SetLabelSize(0.05)
        frDictList[1]['hist'].GetYaxis().SetTitleSize(0.06)
        #h_all.GetYaxis().SetTitleFont(13)
        frDictList[1]['hist'].Draw()
        
        #Draw background part
        xmin = h_all.GetXaxis().GetXmin()
        xmax = h_all.GetXaxis().GetXmax()
        ymin = h_all.GetBinContent(1) #Not technically ymin, but an y at xmin
        ymax = h_all.GetBinContent(100) #Not technically ymin, but an y at xmax

        bkg = TF1("bkg","expo", xmin, xmax)
        b = parsePaveText(frDictList[1]['text'],"tau")[0]
        
        a = log(
            (ymin-ymax)/(exp(xmin*b)-exp(xmax*b))
            )
        bkg.SetParameters(a,b)

        bkg_entries = parsePaveText(frDictList[1]['text'],"background_yield")[0]
        h_bkg = TH1D("h_bkg","h_bkg",100,xmin,xmax)
        h_bkg.FillRandom("bkg",int(bkg_entries))
        h_bkg.GetXaxis().SetRangeUser(rangeMass(method)[0],rangeMass(method)[1])
        #h_bkg.Draw("SAMEAL")

        h_all.Add(h_bkg,-1.0) #Note it is a sig
        h_all.GetXaxis().SetRangeUser(rangeMass(method)[0],rangeMass(method)[1])
        h_all.SetLineWidth(2)
        h_all.SetLineColor(kRed+2)
        h_all.Draw("SAMEAL")


        #Draw main tag
        mainTag()
        #Draw the method
        methodTex(method)

        canvas.SetTitle("")
        canvas.SaveAs(outputPath+"/"+var+"_bin"+str(i) +"_"+method+".root","root")
        canvas.Print(outputPath+"/"+var+"_bin"+str(i) +"_"+method+".eps")
        #canvas.Print(outputPath+"/"+var+"_bin"+str(i) +"_"+method+".pdf")
        #canvas.SaveSource(outputPath+"/"+var+"_bin"+str(i) +"_"+method+".C")

    #get object from rootfile
    # canvas_list = ["matched", "full", "fail"]
    # RooFrame *f = _file0->GetObject("frame_J_psi_1S_M_81423c40",f)
    # f->Print()
    # RooHist *h = frame->getHist("h_MatchedETA0")
    # TGraph *g = f->getCurve("totalShape_Norm[J_psi_1S_M]")
    # TPaveText* pt = (TPaveText*)f->getObject(2)

    return

#Plot()
Plot(method="Long",var="Full")
Plot(method="Velo",var="Full")
Plot(method="T",var="Full")