using namespace RooFit;
using namespace RooStats;

void exercise_2() {
  
  RooAbsReal::defaultIntegratorConfig()->method1D().setLabel("RooAdaptiveGaussKronrodIntegrator1D");
  
  //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);
  RooArgSet * vars = sbModel.GetPdf()->getVariables() ;
  //RooRealVar * p = (RooRealVar*) vars->find("Npsi");
  //p->setRange(0., 30.);   // set range
  sbModel.SetParametersOfInterest(RooArgSet(varpoi));

  sbModel.Print();

  // the B only model
  ModelConfig bModel = ModelConfig();
  bModel.SetWorkspace(*w);
  bModel.SetObservables("mass");
  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();

  // First example is with a frequentist approach
  FrequentistCalculator fc(*(w->data("data")), bModel, sbModel);
  fc.SetToys(2500,1500);

  // Create hypotest inverter passing the desired calculator 
  HypoTestInverter calc(fc);

  // set confidence level (e.g. 95% upper limits)
  calc.SetConfidenceLevel(0.95);

  // use CLs
  calc.UseCLs(1);

  // reduce the noise
  calc.SetVerbose(0);

  // Configure ToyMC Samler
  auto toymcs = calc.GetHypoTestCalculator()->GetTestStatSampler();

  // Use profile likelihood as test statistics
  ProfileLikelihoodTestStat profll(*(sbModel.GetPdf()));

  // for CLs (bounded intervals) use one-sided profile likelihood
  profll.SetOneSided(1);

  // set the test statistic to use for toys
  toymcs->SetTestStatistic(&profll);

  int npoints = 8; // Number of points to scan
  // min and max for the scan (better to choose smaller intervals)
  double poimin = ((RooRealVar*)(poi.find("Npsi")))->getMin();
  double poimax = ((RooRealVar*)(poi.find("Npsi")))->getMax();

  // print "Doing a fixed scan  in interval : ", poimin, " , ", poimax
  calc.SetFixedScan(npoints,poimin,poimax);

  auto result = calc.GetInterval(); // This is a HypoTestInverter class object
  double upperLimit = result->UpperLimit();


  // 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

  auto bcInterval = bc.GetInterval();

  // Now let's print the result of the two methods
  // First the CLs
  std::cout << "################\n";
  std::cout << "The observed CLs upper limit is: " <<  upperLimit << "\n";

  // Compute expected limit
  std::cout << " Expected upper limits, using the B (alternate) model : \n"; 
  std::cout << " expected limit (median) " <<  result->GetExpectedUpperLimit(0) << "\n";
  std::cout << " expected limit (-1 sig) " << result->GetExpectedUpperLimit(-1) << "\n";;
  std::cout << " expected limit (+1 sig) " << result->GetExpectedUpperLimit(1) << "\n";;
  std::cout << " expected limit (-2 sig) " << result->GetExpectedUpperLimit(-2) << "\n";;
  std::cout << " expected limit (+2 sig) " << result->GetExpectedUpperLimit(2) << "\n";;
  std::cout << "################\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("exercise_2.png");
  dataCanvas->SaveAs("exercise_2.pdf");

}