%%cpp -d

#include "TFile.h"
#include "TROOT.h"
#include "TCanvas.h"
#include "TMath.h"
#include "TSystem.h"
#include "RooWorkspace.h"
#include "RooAbsData.h"

#include "RooStats/ModelConfig.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/MCMCInterval.h"
#include "RooStats/MCMCIntervalPlot.h"
#include "RooStats/SequentialProposal.h"
#include "RooStats/ProposalHelper.h"
#include "RooStats/ProposalHelper.h"
#include "RooFitResult.h"

using namespace RooFit;
using namespace RooStats;

struct BayesianMCMCOptions {

   double confLevel = 0.95;
   int intervalType = 2; // type of interval (0 is shortest, 1 central, 2 upper limit)
   double maxPOI = -999; // force different values of POI for doing the scan (default is given value)
   double minPOI = -999;
   int numIters = 100000;    // number of iterations
   int numBurnInSteps = 100; // number of burn in steps to be ignored
};

BayesianMCMCOptions optMCMC;

const char *infile = "";
const char *workspaceName = "combined";
const char *modelConfigName = "ModelConfig";
const char *dataName = "obsData";

const char *filename = "";
if (!strcmp(infile, "")) {
   filename = "results/example_combined_GaussExample_model.root";
   bool fileExist = !gSystem->AccessPathName(filename); // note opposite return code
   // if file does not exists generate with histfactory
   if (!fileExist) {
      // Normally this would be run on the command line
      cout << "will run standard hist2workspace example" << endl;
      gROOT->ProcessLine(".! prepareHistFactory .");
      gROOT->ProcessLine(".! hist2workspace config/example.xml");
      cout << "\n\n---------------------" << endl;
      cout << "Done creating example input" << endl;
      cout << "---------------------\n\n" << endl;
   }

} else
   filename = infile;

TFile *file = TFile::Open(filename);

if (!file) {
   cout << "StandardRooStatsDemoMacro: Input file " << filename << " is not found" << endl;
   return;
}

RooWorkspace *w = (RooWorkspace *)file->Get(workspaceName);
if (!w) {
   cout << "workspace not found" << endl;
   return;
}

ModelConfig *mc = (ModelConfig *)w->obj(modelConfigName);

RooAbsData *data = w->data(dataName);

if (!data || !mc) {
   w->Print();
   cout << "data or ModelConfig was not found" << endl;
   return;
}

/*

RooFitResult* fit = mc->GetPdf()->fitTo(*data,Save());
ProposalHelper ph;
ph.SetVariables((RooArgSet&)fit->floatParsFinal());
ph.SetCovMatrix(fit->covarianceMatrix());
ph.SetUpdateProposalParameters(kTRUE); // auto-create mean vars and add mappings
ph.SetCacheSize(100);
ProposalFunction* pf = ph.GetProposalFunction();
*/

SequentialProposal sp(0.1);

MCMCCalculator mcmc(*data, *mc);
mcmc.SetConfidenceLevel(optMCMC.confLevel); // 95% interval

mcmc.SetProposalFunction(sp);
mcmc.SetNumIters(optMCMC.numIters);             // Metropolis-Hastings algorithm iterations
mcmc.SetNumBurnInSteps(optMCMC.numBurnInSteps); // first N steps to be ignored as burn-in

if (optMCMC.intervalType == 0)
   mcmc.SetIntervalType(MCMCInterval::kShortest); // for shortest interval (not really needed)
if (optMCMC.intervalType == 1)
   mcmc.SetLeftSideTailFraction(0.5); // for central interval
if (optMCMC.intervalType == 2)
   mcmc.SetLeftSideTailFraction(0.); // for upper limit

RooRealVar *firstPOI = (RooRealVar *)mc->GetParametersOfInterest()->first();
if (optMCMC.minPOI != -999)
   firstPOI->setMin(optMCMC.minPOI);
if (optMCMC.maxPOI != -999)
   firstPOI->setMax(optMCMC.maxPOI);

MCMCInterval *interval = mcmc.GetInterval();

auto c1 = new TCanvas("IntervalPlot");
MCMCIntervalPlot plot(*interval);
plot.Draw();

TCanvas *c2 = new TCanvas("extraPlots");
const RooArgSet *list = mc->GetNuisanceParameters();
if (list->getSize() > 1) {
   double n = list->getSize();
   int ny = TMath::CeilNint(sqrt(n));
   int nx = TMath::CeilNint(double(n) / ny);
   c2->Divide(nx, ny);
}

int iPad = 1; // iPad, that's funny
for (auto *nuis : static_range_cast<RooRealVar *>(*mc->GetNuisanceParameters())) {
   c2->cd(iPad++);
   plot.DrawChainScatter(*firstPOI, *nuis);
}

cout << "\n>>>> RESULT : " << optMCMC.confLevel * 100 << "% interval on " << firstPOI->GetName() << " is : ["
     << interval->LowerLimit(*firstPOI) << ", " << interval->UpperLimit(*firstPOI) << "] " << endl;

gPad = c1;

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