%%cpp -d
#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooUniform.h"
#include "RooDataSet.h"
#include "RooDataHist.h"
#include "RooHistFunc.h"
#include "RooRealSumPdf.h"
#include "RooParamHistFunc.h"
#include "RooHistConstraint.h"
#include "RooProdPdf.h"
#include "RooPlot.h"

#include "TCanvas.h"
#include "TPaveText.h"

#include <iostream>
#include <memory>

using namespace RooFit;

RooRealVar x("x", "x", -20, 20);
x.setBins(25);

RooRealVar meanG("meanG", "meanG", 1, -10, 10);
RooRealVar sigG("sigG", "sigG", 1.5, -10, 10);
RooGaussian g("g", "Gauss", x, meanG, sigG);
RooUniform u("u", "Uniform", x);

std::unique_ptr<RooDataSet> sigData(g.generate(x, 50));
std::unique_ptr<RooDataSet> bkgData(u.generate(x, 1000));

RooDataSet sumData("sumData", "Gauss + Uniform", x);
sumData.append(*sigData);
sumData.append(*bkgData);

std::unique_ptr<RooDataHist> dh_sig( g.generateBinned(x, 50) );
std::unique_ptr<RooDataHist> dh_bkg( u.generateBinned(x, 10000) );

RooHistFunc p_h_sig("p_h_sig","p_h_sig",x,*dh_sig);
RooHistFunc p_h_bkg("p_h_bkg","p_h_bkg",x,*dh_bkg);

RooRealVar Asig0("Asig","Asig",1,0.01,5000);
RooRealVar Abkg0("Abkg","Abkg",1,0.01,5000);

RooRealSumPdf model0("model0","model0",
    RooArgList(p_h_sig,p_h_bkg),
    RooArgList(Asig0,Abkg0),
    true);

RooParamHistFunc p_ph_sig1("p_ph_sig","p_ph_sig",*dh_sig, x);
RooParamHistFunc p_ph_bkg1("p_ph_bkg","p_ph_bkg",*dh_bkg, x);

RooRealVar Asig1("Asig","Asig",1,0.01,5000);
RooRealVar Abkg1("Abkg","Abkg",1,0.01,5000);

RooRealSumPdf model_tmp("sp_ph", "sp_ph",
    RooArgList(p_ph_sig1,p_ph_bkg1),
    RooArgList(Asig1,Abkg1),
    true);

RooHistConstraint hc_sig("hc_sig","hc_sig",p_ph_sig1);
RooHistConstraint hc_bkg("hc_bkg","hc_bkg",p_ph_bkg1);

RooProdPdf model1("model1","model1",RooArgSet(hc_sig,hc_bkg),Conditional(model_tmp,x));

RooParamHistFunc p_ph_sig2("p_ph_sig2", "p_ph_sig2", *dh_sig, x);
RooParamHistFunc p_ph_bkg2("p_ph_bkg2", "p_ph_bkg2", *dh_bkg, x, &p_ph_sig2, true);

RooRealVar Asig2("Asig","Asig",1,0.01,5000);
RooRealVar Abkg2("Abkg","Abkg",1,0.01,5000);

RooRealSumPdf model2_tmp("sp_ph","sp_ph",
    RooArgList(p_ph_sig2,p_ph_bkg2),
    RooArgList(Asig2,Abkg2),
    true);

RooHistConstraint hc_sigbkg("hc_sigbkg","hc_sigbkg",RooArgSet(p_ph_sig2,p_ph_bkg2));

RooProdPdf model2("model2","model2",hc_sigbkg, RooFit::Conditional(model2_tmp,x));

auto result0 = model0.fitTo(sumData, PrintLevel(0), Save());
auto result1 = model1.fitTo(sumData, PrintLevel(0), Save());
auto result2 = model2.fitTo(sumData, PrintLevel(0), Save());


TCanvas* can = new TCanvas("can", "", 1500, 600);
can->Divide(3,1);

TPaveText pt(-19.5, 1, -2, 25);
pt.SetFillStyle(0);
pt.SetBorderSize(0);


can->cd(1);
auto frame = x.frame(Title("No template uncertainties"));

sumData.plotOn(frame);
model0.plotOn(frame, LineColor(kBlue), VisualizeError(*result0));

sumData.plotOn(frame);

model0.plotOn(frame, LineColor(kBlue));
model0.plotOn(frame, Components(p_h_sig), LineColor(kAzure));
model0.plotOn(frame, Components(p_h_bkg), LineColor(kRed));
model0.paramOn(frame);

sigData->plotOn(frame, MarkerColor(kBlue));
frame->Draw();

for (auto text : {
  "No template uncertainties",
  "are taken into account.",
  "This leads to low errors",
  "for the parameters A, since",
  "the only source of errors",
  "are the data statistics."}) {
  pt.AddText(text);
}
pt.DrawClone();


can->cd(2);
frame = x.frame(Title("Barlow Beeston for Sig & Bkg separately"));
sumData.plotOn(frame);
model1.plotOn(frame, LineColor(kBlue), VisualizeError(*result1));

sumData.plotOn(frame);
model1.plotOn(frame, LineColor(kBlue));
model1.plotOn(frame, Components(p_ph_sig1), LineColor(kAzure));
model1.plotOn(frame, Components(p_ph_bkg1), LineColor(kRed));
model1.paramOn(frame, Parameters(RooArgSet(Asig1, Abkg1)));

sigData->plotOn(frame, MarkerColor(kBlue));
frame->Draw();

pt.Clear();
for (auto text : {
  "With gamma parameters, the",
  "signal & background templates",
  "can adapt to the data.",
  "Note how the blue signal",
  "template changes its shape.",
  "This leads to higher errors",
  "of the scale parameters A."}) {
  pt.AddText(text);
}
pt.DrawClone();

can->cd(3);
frame = x.frame(Title("Barlow Beeston light for (Sig+Bkg)"));
sumData.plotOn(frame);
model2.plotOn(frame, LineColor(kBlue), VisualizeError(*result2));

sumData.plotOn(frame);
model2.plotOn(frame, LineColor(kBlue));
model2.plotOn(frame, Components(p_ph_sig2), LineColor(kAzure));
model2.plotOn(frame, Components(p_ph_bkg2), LineColor(kRed));
model2.paramOn(frame, Parameters(RooArgSet(Asig2, Abkg2)));

sigData->plotOn(frame, MarkerColor(kBlue));
frame->Draw();

pt.Clear();
for (auto text : {
  "When signal and background",
  "template share one gamma para-",
  "meter per bin, they adapt less.",
  "The errors of the A parameters",
  "also shrink slightly."}) {
  pt.AddText(text);
}
pt.DrawClone();


std::cout << "Asig [normal ] = " << Asig0.getVal() << " +/- " << Asig0.getError() << std::endl;
std::cout << "Asig [BB     ] = " << Asig1.getVal() << " +/- " << Asig1.getError() << std::endl;
std::cout << "Asig [BBlight] = " << Asig2.getVal() << " +/- " << Asig2.getError() << std::endl;

%jsroot on
gROOT->GetListOfCanvases()->Draw()