from __future__ import print_function, division
import os
from ROOT import RooDataHist, RooArgSet, RooArgList, RooRealVar, RooBDecay, RooEffProd, RooTruthModel, RooFFTConvPdf, RooBinning, RooRealConstant, RooPlot,RooPolyVar, TCanvas, TAxis, RooExponential, TH1D, TFile, RooFit, RooHistPdf,RooGenericPdf, RooTruthModel, RooDecay, RooFormulaVar, RooGaussian, RooSimultaneous, RooCategory, RooDataSet
from ROOT import kRed, kDashed, gPad

def fit():

    const = RooRealConstant.value
       
    t = RooRealVar("t", "t", -10., 30.)
    tau = RooRealVar("tau", "tau", 1.54, 0.1, 5.)
    tau2 = RooRealVar("tau2", "tau2", 0.9, 0.1, 3.)

    mean = RooRealVar("mean", "mean", 0)
    sigma = RooRealVar("sigma", "sigma", 0.5)
    
    gaussian_resol = RooGaussian("gauss", "gauss", t, mean, sigma)
    

    datahist = gaussian_resol.generateBinned(t, 100000)
    acceptance = RooFormulaVar("spline", "spline", "0.3+0.02*t", RooArgList(t))
    resol = RooHistPdf("res", "res", t, datahist)
    
    true = RooTruthModel("true", "true", t)
    model1 =  RooDecay("model1", "model1", t, tau, true, RooDecay.SingleSided)
    model2 = RooDecay("model2", "model2", t, tau2, true, RooDecay.SingleSided)
    cat = RooCategory("cat", "cat")
    cat.defineType("long")
    cat.defineType("short")


    model = RooSimultaneous("model", "model", cat)
    model.addPdf(model1, "long")
    model.addPdf(model2, "short")
    
    model1conv = RooFFTConvPdf("model1Conv", "model1Conv", t, model1, resol)
    model1Eff = RooEffProd("model1Eff", "model1Eff", model1conv, acceptance)
    model2conv = RooFFTConvPdf("model2Conv", "model2Conv", t, model2, resol)
    model2Eff = RooEffProd("model2Eff", "model2Eff", model2conv, acceptance)
 

    modelconv = RooFFTConvPdf("modelConv", "modelConv", t, model, resol)
    modelEff = RooEffProd("modelEff", "modelEff", modelconv, acceptance)
       
    frame1 = t.frame()
    data1 = model1Eff.generate(RooArgSet(t), 10000)
    data2 = model2Eff.generate(RooArgSet(t), 20000)
  
    data = RooDataSet("data", "data", RooArgSet(t), RooFit.Index(cat), RooFit.Import("long" , data1), RooFit.Import("short", data2))

    modelEff.fitTo(data)
    data.plotOn(frame1)
    modelEff.plotOn(frame1)
 
    c = TCanvas("rf703_effpdfprod", "rf703_effpdfprod", 600, 600)
    
    c.cd()
    gPad.SetLeftMargin(0.15)
    frame1.GetYaxis().SetTitleOffset(1.4)
    frame1.Draw()
    c.SaveAs('test.pdf')


if __name__ == '__main__':
    fit()