import ROOT
from ROOT import RooRealVar, RooArgList, RooFormulaVar, RooAddPdf, RooFit, RooExponential, RooGaussian, RooFFTConvPdf

fileP = ROOT.TFile.Open("nphistwx.root")
histograms = fileP.Get("histNP")

var = ROOT.RooRealVar("tau", "Pseudoproper Lifetime ", -1, 2, "ps")
def pdf_to_th(name, model, variable, binning):
    """Converts RooFit pdf into a histogram
    Main purpose of this function is to correctly normalize
    by bin width, RooFit divides by bin width.
    """
    root_hist = model.createHistogram(name, variable, Extended=False, Scaling=False, Binning=binning)
    root_hist.SetDirectory(0)
    ROOT.SetOwnership(root_hist, True)
 
    for i in range(root_hist.GetNbinsX()):
        val = root_hist.GetBinContent(i)*total_events*bin_width
        root_hist.SetBinContent(i, val)
        #output1 = ROOT.TFile("pdf.root", "RECREATE")
        #root_hist.Write()
        #output1.Close()
        bin_contents = [root_hist.GetBinContent(i+1) for i in range(root_hist.GetNbinsX())]
    #print(bin_contents)

    return root_hist
def roofit_draw(plotName, var, data, model, npar, extra_labels, components = {}):
    """Plots result of RooFit fit with possibility to plot specific model components.

    Arguments:
        var (``RooRealVar``): Variable used in the fit.
        data (``RooDataHist``): Hist to which model wasfitted to.
        model (``RooAddPdf``): Model which was fitted.
        result (``RooFitResult``): Result of the fit.
        extra_labels (``list`` of `` strings``): Extra labels to be
            added to the plot
        components (``string``): Components of the model to be plotted
            in addition to the complete model
    """

    ROOT.gROOT.SetStyle('ATLAS')
    ROOT.gROOT.SetBatch(True)
    c = ROOT.TCanvas(plotName)
    ROOT.SetOwnership(c, True)
    c.SetCanvasSize(1000, 1000)
    c.cd()
    main_pad = ROOT.TPad('mainpad', 'mainpad', 0, 0.3, 1, 1 )
    ROOT.SetOwnership(main_pad, True)
    main_pad.SetBottomMargin(0)
    main_pad.Draw()
    main_pad.cd()
    frame = var.frame()
    ROOT.SetOwnership(frame, True)
    data.plotOn(frame, ROOT.RooFit.Name("data"))
    model.plotOn(frame, ROOT.RooFit.Name("model"), ROOT.RooFit.LineColor(ROOT.kRed))
    #chi2_var = ROOT.RooChi2Var("chi2", "chi2", model, data, False, ROOT.RooDataHist.SumW2)
    #chi2_value = chi2_var.getVal()
    chi2 = frame.chiSquare("model", "data", npar)
    #ndof = dataHist.numEntries() - npar
   # chi2= chi2_value/ndof
    #col = 0
    #cols = [ROOT.kGreen+2, ROOT.kBlue, ROOT.kOrange+1, ROOT.kCyan]
    #for name, comp in components.items():
         #model.plotOn(frame, ROOT.RooFit.Name(name), ROOT.RooFit.Components(comp),
                     #ROOT.RooFit.LineColor(cols[col]),
                     #ROOT.RooFit.LineStyle(ROOT.kDotted))
         #col+=1
    frame.Draw()

    frame.GetYaxis().SetRangeUser(0, frame.GetMaximum()*1.5)

    leg = ROOT.TLegend(0.7, 0.63, 0.91, 0.85)
    ROOT.SetOwnership(leg, True)
    leg.SetBorderSize(0)
    leg.SetFillStyle(0)
    leg.AddEntry("data", "Data", "LP")
    leg.AddEntry("model", "model", "L")
    for name, comp in components.items():
        leg.AddEntry(name, name, "L")
    leg.Draw()

    ltx = ROOT.TLatex()
    ROOT.SetOwnership(ltx, True)
    ltx.SetNDC()
    ltx.SetTextColor(ROOT.kBlack)
    pos_y = 0.85
    ltx.DrawLatex(0.2, pos_y, "#chi^{{2}}/dof={chi2:.3f}".format(chi2=chi2))
    for lbl in extra_labels:
        pos_y-=0.05
        ltx.DrawLatex(0.2, pos_y, lbl)


    c.cd()
    ratio_pad = ROOT.TPad('ratiopad', 'ratiopad', 0, 0, 1, 0.3)
    ROOT.SetOwnership(ratio_pad, True)
    ratio_pad.SetBottomMargin(0.35)
    ratio_pad.SetTopMargin(0)
    ratio_pad.Draw()
    ratio_pad.cd()
    ratio = data.createHistogram('data_hist', var)
    ROOT.SetOwnership(ratio, True)
    ratio.SetDirectory(0)
    tmp = pdf_to_th('mc_hist', model, var, data.getBinnings()[0])


    for i in range(ratio.GetNbinsX()):
        val = ratio.GetBinContent(i)
        err = ratio.GetBinError(i)
        valmc = tmp.GetBinContent(i)
        if val != 0:
        #if abs(valmc) > 1e-30:
            ratio.SetBinContent(i, val/val)
            #ratio.SetBinError(i, err/valmc)
        else:
            ratio.SetBinContent(i, 0)
            ratio.SetBinError(i, 0)
    tmp.GetYaxis().SetRangeUser(-2, 10)
    tmp.GetYaxis().SetTitle("data/model")
    tmp.GetYaxis().SetTitleOffset(0.56)
    tmp.GetXaxis().SetTitleSize(0.12)
    tmp.GetXaxis().SetLabelSize(0.12)
    tmp.GetYaxis().SetTitleSize(0.12)
    tmp.GetYaxis().SetLabelSize(0.12)
    tmp.SetNdivisions(505, 'Y')
    tmp.Divide(tmp)
    tmp.SetLineColor(ROOT.kRed)
    tmp.Draw("hist")
    ratio.Draw("same")

    # if path contains directories, check if they exist
    # if not, create them
    if "/" in plotName:
        dirName = plotName[:plotName.rindex("/")]
        if not os.path.exists(dirName):
            os.makedirs(dirName)

    # Supress info messages about creating plots
    ignoreLevel = ROOT.gErrorIgnoreLevel
    ROOT.gErrorIgnoreLevel = ROOT.kWarning
    c.SaveAs(plotName+".png")
    ROOT.gErrorIgnoreLevel = ignoreLevel
# --- Adjust Histogram Binning ---
histograms.Rebin(1)
#model used:


landau_mean = ROOT.RooRealVar("landau_mean", "Landau MPV", 0.005, -0.01, 0.3)
landau_sigma = ROOT.RooRealVar("landau_sigma", "Landau Sigma", 0.15, 0.05, 0.4)
mean1 = RooRealVar("mean1", "Mean of Gaussian 1", 0, -0.1, 0.5)
sigma1 = RooRealVar("sigma1", "Sigma of Gaussian 1", 0.5, 0, 1)
gauss1 = ROOT.RooGaussian("gauss", "Gaussian", var, mean1, sigma1)
landau = ROOT.RooLandau("landau", "Landau Distribution", var, landau_mean, landau_sigma)
frac1 = RooRealVar("frac1", "Fraction of Gauss1", 0.5, 0, 1)
model = ROOT.RooAddPdf("landaugaussian","landugaussian", RooArgList(landau, gauss1), RooArgList(frac1), True)
# Convert histogram to RooDataHist
dataHist = ROOT.RooDataHist("dataHist", "Data Histogram", ROOT.RooArgList(var), histograms)
# Get the total number of entries in dataHist
total_events = dataHist.sumEntries()
bin_width = 0.05
result = model.chi2FitTo(dataHist, ROOT.RooFit.Save(),  ROOT.RooFit.Strategy(2))
ratio=roofit_draw( "lgausspt15w" ,var, dataHist, model, 5, ["MC Non prompt", "Gaussian+Landau Model"], {"Landau Distribution":"landau" ,"Gaussian" : "gauss" } )
fileP.Close()