%%cpp -d
#include "RooGlobalFunc.h"
#include <cstdlib>
#include "TMatrixDSym.h"
#include "RooMultiVarGaussian.h"
#include "RooArgList.h"
#include "RooRealVar.h"
#include "TH2F.h"
#include "TCanvas.h"
#include "RooAbsReal.h"
#include "RooFitResult.h"
#include "TStopwatch.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/MetropolisHastings.h"
#include "RooStats/MarkovChain.h"
#include "RooStats/ConfInterval.h"
#include "RooStats/MCMCInterval.h"
#include "RooStats/MCMCIntervalPlot.h"
#include "RooStats/ModelConfig.h"
#include "RooStats/ProposalHelper.h"
#include "RooStats/ProposalFunction.h"
#include "RooStats/PdfProposal.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/LikelihoodIntervalPlot.h"
#include "RooStats/LikelihoodInterval.h"

using std::cout, std::endl;
using namespace RooFit;
using namespace RooStats;

Int_t dim = 4;
Int_t nPOI = 2;

TStopwatch t;
t.Start();

RooArgList xVec;
RooArgList muVec;
RooArgSet poi;

Int_t i, j;
RooRealVar *x;
RooRealVar *mu_x;
for (i = 0; i < dim; i++) {
   char *name = Form("x%d", i);
   x = new RooRealVar(name, name, 0, -3, 3);
   xVec.add(*x);

   char *mu_name = Form("mu_x%d", i);
   mu_x = new RooRealVar(mu_name, mu_name, 0, -2, 2);
   muVec.add(*mu_x);
}

for (i = 0; i < nPOI; i++) {
   poi.add(*muVec.at(i));
}

TMatrixDSym cov(dim);
for (i = 0; i < dim; i++) {
   for (j = 0; j < dim; j++) {
      if (i == j)
         cov(i, j) = 3.;
      else
         cov(i, j) = 1.0;
   }
}

RooMultiVarGaussian mvg("mvg", "mvg", xVec, muVec, cov);

std::unique_ptr<RooDataSet> data{mvg.generate(xVec, 100)};

RooWorkspace *w = new RooWorkspace("MVG");
ModelConfig modelConfig(w);
modelConfig.SetPdf(mvg);
modelConfig.SetParametersOfInterest(poi);

std::unique_ptr<RooFitResult> fit{mvg.fitTo(*data, Save(true))};
ProposalHelper ph;
ph.SetVariables((RooArgSet &)fit->floatParsFinal());
ph.SetCovMatrix(fit->covarianceMatrix());
ph.SetUpdateProposalParameters(true);
ph.SetCacheSize(100);
ProposalFunction *pdfProp = ph.GetProposalFunction();

MCMCCalculator mc(*data, modelConfig);
mc.SetConfidenceLevel(0.95);
mc.SetNumBurnInSteps(100);
mc.SetNumIters(10000);
mc.SetNumBins(50);
mc.SetProposalFunction(*pdfProp);

MCMCInterval *mcInt = mc.GetInterval();
RooArgList *poiList = mcInt->GetAxes();

ProfileLikelihoodCalculator plc(*data, modelConfig);
plc.SetConfidenceLevel(0.95);
LikelihoodInterval *plInt = (LikelihoodInterval *)plc.GetInterval();

MCMCIntervalPlot mcPlot(*mcInt);

TCanvas *c1 = new TCanvas();
mcPlot.SetLineColor(kGreen);
mcPlot.SetLineWidth(2);
mcPlot.Draw();

LikelihoodIntervalPlot plPlot(plInt);
plPlot.Draw("same");

if (poiList->getSize() == 1) {
   RooRealVar *p = (RooRealVar *)poiList->at(0);
   Double_t ll = mcInt->LowerLimit(*p);
   Double_t ul = mcInt->UpperLimit(*p);
   cout << "MCMC interval: [" << ll << ", " << ul << "]" << endl;
}

if (poiList->getSize() == 2) {
   RooRealVar *p0 = (RooRealVar *)poiList->at(0);
   RooRealVar *p1 = (RooRealVar *)poiList->at(1);
   TCanvas *scatter = new TCanvas();
   Double_t ll = mcInt->LowerLimit(*p0);
   Double_t ul = mcInt->UpperLimit(*p0);
   cout << "MCMC interval on p0: [" << ll << ", " << ul << "]" << endl;
   ll = mcInt->LowerLimit(*p1);
   ul = mcInt->UpperLimit(*p1);
   cout << "MCMC interval on p1: [" << ll << ", " << ul << "]" << endl;

   // MCMC interval on p0: [-0.2, 0.6]
   // MCMC interval on p1: [-0.2, 0.6]

   mcPlot.DrawChainScatter(*p0, *p1);
   scatter->Update();
}

t.Print();

gROOT->GetListOfCanvases()->Draw()