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 matchTag(match,full):    #This order is crucial
    if (match): return "Matched"
    if (full):  return "Full"
    else:       return "Failed"    

def matchTagDraw(match, full):
    tex = TLatex()
    tex.SetNDC(True)
    tex.SetTextFont(132)
    tex.SetTextSize(0.06)
    tex.SetTextAlign(13)  
    tex.DrawLatex(0.205, 0.805, matchTag(match,full))

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 getCanvas(method):
    canvas = TCanvas("c_"+method, "c_"+method, 10,10,650,600)
    canvas.SetBottomMargin(0.15)
    canvas.SetRightMargin(0.05)
    canvas.SetLeftMargin(0.16)   
    canvas.SetTopMargin(0.0675)   
    return canvas
    
def designRooHist(hist,method):
    #Set range
    hist.GetXaxis().SetRangeUser(rangeMass(method)[0],rangeMass(method)[1])
    
    #Set line design
    hist.SetLineWidth(2)
    hist.SetLineColor(kBlack)

    #Set title
    hist.SetTitle("")

    #Set X axis
    hist.GetXaxis().SetTitle("m_{#mu^{+}#mu^{-}} [MeV]")
    hist.GetXaxis().SetNdivisions(505)
    hist.GetXaxis().SetTitleFont(132)
    hist.GetXaxis().SetTitleSize(0.06)
    hist.GetXaxis().SetTitleOffset(1.1)
    hist.GetXaxis().SetLabelFont(132)
    hist.GetXaxis().SetLabelSize(0.05)

    #Set Y axis
    hist.GetYaxis().SetTitle("Entries (a.u.)")
    hist.GetYaxis().SetNdivisions(505)
    hist.GetYaxis().SetTitleFont(132)
    hist.GetYaxis().SetTitleSize(0.06)
    hist.GetYaxis().SetTitleOffset(1.3)
    hist.GetYaxis().SetLabelFont(132)
    hist.GetYaxis().SetLabelSize(0.05)
    return

def getHistLimits(hist):
    xmin = hist.GetXaxis().GetXmin()
    xmax = hist.GetXaxis().GetXmax()
    ymin = hist.GetBinContent(1) #Not technically ymin, but an y at xmin
    ymax = hist.GetBinContent(100) #Not technically ymin, but an y at xmax
    return xmin, xmax, ymin, ymax

def designSignal(hist, method):
    hist.GetXaxis().SetRangeUser(rangeMass(method)[0],rangeMass(method)[1])
    hist.SetLineWidth(2)
    hist.SetLineStyle(7)
    hist.SetLineColor(kRed+2)
    return

def getPlotName(outputPath, 
                var, bin, method, matchTag):
    return outputPath+"/"+var+"_bin"+str(bin) +"_"+method+ "_" + matchTag

def plotYield(histRoo, bkg_entries, bkg_exp,     
              method, match, full,
              plotName  ): 
    #create canvases
    canvas = getCanvas(method)      
   
    #Draw the hist
    h_all = ConvertRooHistToTH1D(histRoo)
    designRooHist(frDictList[match]['hist'],method)
    hist.GetYaxis().SetRangeUser(0,h_all.GetMaximum()*1.25)
    histRoo.Draw()  

    #Get the hist values at edges and get the background histogram
    xmin, xmax, ymin, ymax = getHistLimits(h_all)

    bkg = TF1("bkg","expo", xmin, xmax)
        
    a = log(
        (ymin-ymax)/(exp(xmin*bkg_exp)-exp(xmax*bkg_exp))
        )
    bkg.SetParameters(a,bkg_exp)
    print (a,bkg_exp)
    
    #Fill the bkg histogram
    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")

    #Get signal
    h_all.Add(h_bkg,-1.0) #Substract the h_bkg from h_all 
    designSignal(h_all, method)
    h_all.Draw("SAMEAL")

    #Draw main tag
    mainTag()
    #Draw the method
    methodTex(method)
    #Draw matched/failed
    matchTagDraw(match,full)

    canvas.SetTitle("")
    canvas.SaveAs(plotName+".root","root")
    canvas.Print(plotName+".eps")
    return


def Plot(simVer = "Sim09h", year = "2018_25ns", data = True,
         method = "Long", var = "ETA", simFit = True,
         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)
        
        if ( (simFit) and isMatched(name)): 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
    #Matched and Failed has to be treated seperately as the failed ones don't have TPaveText
    for i,frDictList in enumerate (sortedList):
        if (simFit):
            #This is needed only cause full sample doesn't have a pavetext in failed sample... sigh
            b_M = parsePaveText(frDictList[0]['text'],"tau")[0]
            b_F = parsePaveText(frDictList[0]['text'],"tauF")[0]

            bkg_entries_M = parsePaveText(frDictList[0]['text'],"background_yield")[0]
            bkg_eff = parsePaveText(frDictList[0]['text'],"efficiency_bkg")[0]
            bkg_entries_F = bkg_entries_M/bkg_eff - bkg_entries_M #eff = M/(M+F)
                
            plotName = getPlotName(outputPath,var,i,method,matchTag(True,False))
            plotYield(frDictList[0]['hist'],bkg_entries_M,b_M,method,True,False,plotName)

            plotName = getPlotName(outputPath,var,i,method,matchTag(False,False))
            plotYield(frDictList[1]['hist'],bkg_entries_F,b_F,method,False,False,plotName)
                
        else:    
            for matched in [True,False]:
                b = parsePaveText(frDictList[matched]['text'],"tau")[0]
                bkg_entries = parsePaveText(frDictList[matched]['text'],"background_yield")[0]
                plotName = getPlotName(outputPath,var,i,method,matchTag(False,True))
                plotYield(frDictList[matched]['hist'],method,bkg_entries,b,matched,True,outputPath)

    return

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