#include "TStopwatch.h"
#include "TCanvas.h"
#include "TROOT.h"

#include "RooPlot.h"
#include "RooAbsPdf.h"
#include "RooWorkspace.h"
#include "RooDataSet.h"
#include "RooGlobalFunc.h"
#include "RooFitResult.h"
#include "RooRandom.h"
#include "RooAbsReal.h"

#include "RooStats/RooStatsUtils.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/LikelihoodInterval.h"
#include "RooStats/LikelihoodIntervalPlot.h"
#include "RooStats/BayesianCalculator.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/MCMCInterval.h"
#include "RooStats/MCMCIntervalPlot.h"
#include "RooStats/ProposalHelper.h"
#include "RooStats/SimpleInterval.h"
#include "RooStats/FeldmanCousins.h"
#include "RooStats/PointSetInterval.h"
#include "RooStats/ToyMCSampler.h"
#include "RooStats/ProfileLikelihoodTestStat.h"

using namespace std;
using namespace RooFit;
using namespace RooStats;

ModelConfig* makeMyModel (const char* name, RooWorkspace& ws) {
  // derived from data
  ws.factory("sig[2,0,20]"); // POI

  // predefined nuisances
  ws.factory("bg_a[5,0,20]");

  ws.factory("Poisson::pdf_a(na[10,0,100],sum::mu_a(sig,bg_a))");

  // nuisance PDFs (systematics)
  ws.factory("Lognormal::l_bg_a(bg_a,nom_bg_a[5,0,20],sum::kappa_bg_a(1,d_bg_a[0.10]))");

  // model
  ws.factory("PROD::model(pdf_a,l_bg_a)");

  // observables
  ws.defineSet("obs","na");

  // parameters of interest
  ws.defineSet("poi","sig");

  // nuisance parameters
  ws.defineSet("nuis","bg_a");

  // prior (for Bayesian calculation)
  ws.factory("Uniform::prior(sig)");

  // model config
  ModelConfig* modelConfig = new ModelConfig(name);
  modelConfig->SetWorkspace(ws);
  modelConfig->SetPdf("model");
  modelConfig->SetPriorPdf("prior");
  modelConfig->SetParametersOfInterest(*(ws.set("poi")));
  modelConfig->SetNuisanceParameters(*(ws.set("nuis")));
  modelConfig->SetObservables(*(ws.set("obs")));
  ws.import(*modelConfig);
  return modelConfig;
}

double BayesianUpperLimit (RooAbsData& data, ModelConfig& modelConfig, double CL = 0.95) {
  BayesianCalculator calculator (data, modelConfig);
  calculator.SetConfidenceLevel(CL);
  calculator.SetLeftSideTailFraction (0); // UL
  SimpleInterval* interval = calculator.GetInterval();
  calculator.SetScanOfPosterior(40);
  RooPlot * plot = calculator.GetPosteriorPlot();
  plot->Draw(); 
  return interval->UpperLimit();
}

void runModel(){
  RooWorkspace* ws = new RooWorkspace("ws");
  ModelConfig* modelConfig = makeMyModel ("test", *ws);
  RooDataSet data ("data","",*(modelConfig->GetObservables()));
  ws->var("na")->setVal(10);
  data.add( *(modelConfig->GetObservables()));
  ws->import (data); // not really needed for your macro

  // get Bayesian Limit
  double cl95Bayesian = BayesianUpperLimit (data, *modelConfig);
  cout << "Bayesian UL: " << cl95Bayesian << endl;

  // clean up
  delete ws;
  delete modelConfig;

}