using namespace RooFit;
using namespace RooStats;

void bayes_exercise_2() {
  //Open the rootfile and get the workspace from the exercise_0
  TFile *fInput = new TFile("Workspace_mumufit.root");
  RooWorkspace *w = (RooWorkspace*)fInput->Get("w");
  w->Print();

  // You can set constant parameters that are known
  // If you leave them floating, the fit procedure will determine their uncertainty
  // Right now we will fix all the nuisance parameters just to speed up the computing time
  w->var("meanJpsi")->setConstant(1);
  w->var("sigmaJpsi")->setConstant(1);
  w->var("alphaJpsi")->setConstant(1);
  w->var("nJpsi")->setConstant(1);
  w->var("NJpsi")->setConstant(1);
  w->var("meanpsi2S")->setConstant(1);
  w->var("Nbkg")->setConstant(1);
  w->var("a1")->setConstant(1);
  w->var("a2")->setConstant(1);
  w->var("a3")->setConstant(1);


  // Configure the model, we need both the S+B and the B only models
  ModelConfig sbModel = ModelConfig();
  sbModel.SetWorkspace(*w);
  sbModel.SetPdf("totPDF");
  sbModel.SetName("S+B Model");
  RooRealVar varpoi = (*(w->var("Npsi")));
  varpoi.setRange(0.,20.);  // this is mostly for plotting
  RooArgSet poi(varpoi);
  sbModel.SetParametersOfInterest(RooArgSet(varpoi));

  sbModel.Print();

  // the B only model
  ModelConfig bModel = ModelConfig();
  bModel.SetWorkspace(*w);
  bModel.SetPdf("totPDF");
  bModel.SetName( (TString::Format("%s_with_poi_0",sbModel.GetName()).Data()));
  varpoi.setVal(0);
  bModel.SetSnapshot(poi);
  bModel.SetParametersOfInterest(RooArgSet(varpoi));

  bModel.Print();

  // Example using the BayesianCalculator
  // Now we also need to specify a prior in the ModelConfig
  // To be quicker, we'll use the PDF factory facility of RooWorkspace
  // Careful! For simplicity, we are using a flat prior, but this doesn't mean it's the best choice!
  w->factory("Uniform::prior(Npsi)");
  sbModel.SetPriorPdf(*(w->pdf("prior")));

  // Construct the bayesian calculator
  BayesianCalculator bc(*(w->data("data")), sbModel);
  bc.SetConfidenceLevel(0.95);
  bc.SetLeftSideTailFraction(0.); // for upper limit

  std::cout << "\n\nBefore bc.GetInterval() \n\n";
  auto bcInterval = bc.GetInterval();
  std::cout << "\n\nAfter  bc.GetInterval() \n\n";


  // Now let's see what the bayesian limit is
  std::cout << "Bayesian upper limit on Npsi = " << bcInterval->UpperLimit() << "\n";


  // Plot now the result of the scan 

  // First the CLs
  // auto freq_plot = HypoTestInverterPlot("HTI_Result_Plot","Frequentist scan result for Npsi",result);
  // Then the Bayesian posterior
  auto bc_plot = bc.GetPosteriorPlot();

  // Plot in a new canvas with style
  TCanvas *dataCanvas = new TCanvas("dataCanvas");
  //dataCanvas->Divide(2,1);
  dataCanvas->SetLogy(0);
  //dataCanvas->cd(1);
  //freq_plot.Draw("2CL");
  //dataCanvas->cd(2);
  bc_plot->Draw();
  dataCanvas->SaveAs("bayes_exercise_2.png");
  dataCanvas->SaveAs("bayes_exercise_2.pdf");

}