%%cpp -d
#include "RooDataSet.h"
#include "RooRealVar.h"
#include "RooConstVar.h"
#include "RooBernstein.h"
#include "TCanvas.h"
#include "RooAbsPdf.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include <string>
#include <vector>
#include <stdio.h>
#include <sstream>
#include <iostream>

#include "RooProdPdf.h"
#include "RooAddPdf.h"
#include "RooGaussian.h"
#include "RooProfileLL.h"
#include "RooWorkspace.h"

#include "RooStats/BernsteinCorrection.h"

using namespace RooFit;
using namespace RooStats;

Double_t lowRange = -1, highRange = 5;

RooRealVar x("x", "x", lowRange, highRange);

RooGaussian narrow("narrow", "", x, RooConst(0.), RooConst(.8));
RooGaussian wide("wide", "", x, RooConst(0.), RooConst(2.));
RooAddPdf reality("reality", "", RooArgList(narrow, wide), RooConst(0.8));

std::unique_ptr<RooDataSet> data{reality.generate(x, 1000)};

RooRealVar sigma("sigma", "", 1., 0, 10);
RooGaussian nominal("nominal", "", x, RooConst(0.), sigma);

RooWorkspace *wks = new RooWorkspace("myWorksspace");

wks->import(*data, Rename("data"));
wks->import(nominal);

if (TClass::GetClass("ROOT::Minuit2::Minuit2Minimizer")) {
   // use Minuit2 if ROOT was built with support for it:
   ROOT::Math::MinimizerOptions::SetDefaultMinimizer("Minuit2");
}

Double_t tolerance = 0.05;
BernsteinCorrection bernsteinCorrection(tolerance);
Int_t degree = bernsteinCorrection.ImportCorrectedPdf(wks, "nominal", "x", "data");

if (degree < 0) {
   Error("rs_bernsteinCorrection", "Bernstein correction failed ! ");
   return;
}

cout << " Correction based on Bernstein Poly of degree " << degree << endl;

RooPlot *frame = x.frame();
data->plotOn(frame);

 nominal.fitTo(*data, PrintLevel(0));
nominal.plotOn(frame);

RooAbsPdf *corrected = wks->pdf("corrected");
if (!corrected)
   return;

corrected->fitTo(*data, PrintLevel(0));
corrected->plotOn(frame, LineColor(kRed));

RooAbsPdf *poly = wks->pdf("poly");
if (poly)
   poly->plotOn(frame, LineColor(kGreen), LineStyle(kDashed));

bool checkSamplingDist = true;
int numToyMC = 20; // increase this value for sensible results

TCanvas *c1 = new TCanvas();
if (checkSamplingDist) {
   c1->Divide(1, 2);
   c1->cd(1);
}
frame->Draw();
gPad->Update();

if (checkSamplingDist) {
   // check sampling dist
   ROOT::Math::MinimizerOptions::SetDefaultPrintLevel(-1);
   TH1F *samplingDist = new TH1F("samplingDist", "", 20, 0, 10);
   TH1F *samplingDistExtra = new TH1F("samplingDistExtra", "", 20, 0, 10);
   bernsteinCorrection.CreateQSamplingDist(wks, "nominal", "x", "data", samplingDist, samplingDistExtra, degree,
                                           numToyMC);

   c1->cd(2);
   samplingDistExtra->SetLineColor(kRed);
   samplingDistExtra->Draw();
   samplingDist->Draw("same");
}

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