using namespace RooFit;
using namespace std;

RooAbsPdf *getModelPdf(RooRealVar& m, string component);

void myRooFit_minimalExample(){
    
    // invariant mass
    RooRealVar *m = new RooRealVar("m", "m (GeV)", 60., 120.);
    
    // plot frames
    RooPlot *f1 = m->frame(Title("Signal template and created pdf"));
    RooPlot *f2 = m->frame(Title("Background template and created pdf"));
    RooPlot *f3 = m->frame(Title("Plot of template based model pdf fitted to toy data"));
    
    // get pdfs to create toy data
    RooAbsPdf *model = getModelPdf(*m, "full");
    RooAbsPdf *sig   = getModelPdf(*m, "signal");
    RooAbsPdf *bkg   = getModelPdf(*m, "background");
    
    // create toy data
    Double_t Ndata = 1000;
    RooDataSet *rdsModel = model->generate(*m, Ndata);
    // create a binned version of the toy data (120 bins)
    RooDataHist *rdhModel = new RooDataHist("rdhModel", "binned data", *m,
                                        Import(*(rdsModel->createHistogram(m->GetName(), 120))));
    
    // get both signal parts seperately
    RooAbsPdf  *sigBW   = getModelPdf(*m, "breitwigner");
    RooAbsPdf  *sigGaus = getModelPdf(*m, "gaus");
    
    // create templates
    Double_t Ntemplate = 10000; // (more statistics than data)
    RooDataSet *rdsSigTemplate = sigBW->generate(*m, 0.82*Ntemplate);
    RooDataSet *rdsBkgTemplate = bkg->generate(*m, 0.18*Ntemplate);
    
    // create binned signal template for histogram based interpolation
    RooDataHist *rdhSigTemplate = new RooDataHist( "SigPdf", "signal model", *m, 
                                     Import(*(rdsSigTemplate->createHistogram(m->GetName(), 60))));
    
    // create pdfs:
    // signal: use second order interpolation
    RooAbsPdf *SigHistPdf = new RooHistPdf( "SigHistPdf", "Histogram based signal p.d.f.",
                                        *m, *rdhSigTemplate, 2);
    // convolve this again with the gaussian
    RooAbsPdf *SigPdf     = new RooNumConvPdf("SigPdf", "Template based Signal pdf convolved with gaussian", 
                                                *m, *sigGaus, *SigHistPdf);
    
    // background: use kernel estimator with rho = 1/sqrt(12)
    RooAbsPdf *BkgPdf = new RooKeysPdf( "BkgPdf", "Kernel estimated background p.d.f.",
                                        *m, *rdsBkgTemplate, RooKeysPdf::MirrorBoth, 0.2887);
    
    // add both pdf's for full fit model
    Double_t f = 0.76; // slightly different fraction than created
    RooRealVar *nbkg  = new RooRealVar("nbkg", "Background events", (1-f)*Ndata, 0.01*Ndata, 0.5*Ndata);
    RooRealVar *nsig  = new RooRealVar("nsig", "Signal events", f*Ndata, 0.5*Ndata, Ndata);
    RooAbsPdf *fitPdf = new RooAddPdf("fitPdf", "Complete fit pdf", 
                                        RooArgList(*SigPdf, *BkgPdf), RooArgList(*nsig, *nbkg));
    
    // now perform an extended likelihood fit to the binned! data
    RooFitResult* fitResult = NULL;
    fitResult = fitPdf->fitTo(*rdhModel, FitOptions("I"), Extended(kTRUE), Timer(kTRUE), Save());
    
    // some plots
    rdhSigTemplate->plotOn(f1, MarkerColor(kRed));
    SigPdf->plotOn(f1, LineColor(kRed+2), LineWidth(3));
    rdsBkgTemplate->plotOn(f2, MarkerColor(kBlue), Binning(60));
    BkgPdf->plotOn(f2, LineColor(kBlue+2), LineWidth(3));
    rdsModel->plotOn(f3, MarkerColor(kBlack));
    fitPdf->plotOn(f3, LineColor(kBlue));
    fitPdf->plotOn(f3, Components("BkgPdf"), LineColor(kBlue), LineStyle(kDashed));
    fitPdf->paramOn(f3, Layout(0.6, 0.9, 0.9));
    
    TCanvas *c1 = new TCanvas("c1", "test canvas", 1500, 500);
    c1->Divide(3, 1);
    c1->cd(1); f1->Draw();
    c1->cd(2); f2->Draw();
    c1->cd(3); f3->Draw();
    
    return ;
}

RooAbsPdf *getModelPdf(RooRealVar& m, string component){
    
    // model: breit wigner convoluted with gaussian plus chebychev background
    
    Double_t Ntot = 1000;    // total events
    Double_t f = 0.8;        // signal fraction
    
    RooRealVar *nbkg = new RooRealVar("nbkg", "Background events", (1-f)*Ntot, 0.01*Ntot, 0.5*Ntot);
    RooRealVar *nsig = new RooRealVar("nsig", "Signal events", f*Ntot, 0.5*Ntot, Ntot);
    
    // construct gaussian and breit wigner
    RooRealVar *mbw = new RooRealVar("mbw", "Mean breit wigner", 90., 85., 95.);
    RooRealVar *sbw = new RooRealVar("sbw", "Width breit wigner", 2.495, 2., 4.);
    RooRealVar *mg  = new RooRealVar("mg", "Mean gaussian", 0., -1., 1.);
    RooRealVar *sg  = new RooRealVar("sg", "Width gaussian", 1., 0.1, 2.);
    RooAbsPdf  *breitwigner = new RooBreitWigner("breitwigner", "Breit Wigner Peak", m, *mbw, *sbw);
    RooAbsPdf  *gaussian    = new RooGaussian("gaussian", "Gaussian smearing pdf", m, *mg, *sg);
    
    // convoluted signal
    RooAbsPdf  *signal = new RooNumConvPdf("signal", "Z peak shape", m, *gaussian, *breitwigner);
    
    // background checbychev
    RooRealVar *a0 = new RooRealVar("a0","a0",0.5, 0., 2.) ;
    RooRealVar *a1 = new RooRealVar("a1","a1",-0.2 ,-2., 2.) ;
    RooAbsPdf  *background = new RooChebychev("background", "Chebychev background", m, RooArgSet(*a0,*a1));
    
    // complete model
    RooAbsPdf  *model = new RooAddPdf("model", "Signal+Background", RooArgList(*signal, *background), RooArgList(*nsig, *nbkg));
    
    if(!component.compare("signal")) return signal;
    if(!component.compare("background")) return background;
    if(!component.compare("gaus")) return gaussian;
    if(!component.compare("breitwigner")) return breitwigner;
    if(!component.compare("full")) return model;
    else{
        cout << "no parameter given while calling getModelPdf(), complete model returned!" << endl;
        return model;
    }
}