%%cpp -d
#include "TFile.h"
#include "TROOT.h"
#include "TH1F.h"
#include "TCanvas.h"
#include "TSystem.h"
#include <iostream>

#include "RooWorkspace.h"
#include "RooSimultaneous.h"
#include "RooAbsData.h"

#include "RooStats/ModelConfig.h"
#include "RooStats/FeldmanCousins.h"
#include "RooStats/ToyMCSampler.h"
#include "RooStats/PointSetInterval.h"
#include "RooStats/ConfidenceBelt.h"

#include "RooStats/RooStatsUtils.h"
#include "RooStats/ProfileLikelihoodTestStat.h"

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

bool useProof = false; // flag to control whether to use Proof
int nworkers = 0;      // number of workers (default use all available cores)

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

double confidenceLevel = 0.95;

double additionalToysFac = 0.5;
int nPointsToScan = 20; // number of steps in the parameter of interest
int nToyMC = 200;       // number of toys used to define the expected limit and band

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 *inputFile = TFile::Open(filename);

RooWorkspace *w = (RooWorkspace *)inputFile->Get(workspaceName);

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

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

cout << "Found data and ModelConfig:" << endl;
mc->Print();

RooRealVar *firstPOI = (RooRealVar *)mc->GetParametersOfInterest()->first();
/*  firstPOI->setMin(0);*/
/*  firstPOI->setMax(10);*/

FeldmanCousins fc(*data, *mc);
fc.SetConfidenceLevel(confidenceLevel);
fc.AdditionalNToysFactor(additionalToysFac); // improve sampling that defines confidence belt

fc.SetNBins(nPointsToScan); // set how many points per parameter of interest to scan
fc.CreateConfBelt(true);    // save the information in the belt for plotting

ToyMCSampler *toymcsampler = (ToyMCSampler *)fc.GetTestStatSampler();
ProfileLikelihoodTestStat *testStat = dynamic_cast<ProfileLikelihoodTestStat *>(toymcsampler->GetTestStatistic());

if (!mc->GetPdf()->canBeExtended()) {
   if (data->numEntries() == 1)
      fc.FluctuateNumDataEntries(false);
   else
      cout << "Not sure what to do about this model" << endl;
}

if (useProof) {
   ProofConfig pc(*w, nworkers, "", false);
   toymcsampler->SetProofConfig(&pc); // enable proof
}

if (mc->GetGlobalObservables()) {
   cout << "will use global observables for unconditional ensemble" << endl;
   mc->GetGlobalObservables()->Print();
   toymcsampler->SetGlobalObservables(*mc->GetGlobalObservables());
}

PointSetInterval *interval = fc.GetInterval();
ConfidenceBelt *belt = fc.GetConfidenceBelt();

cout << "\n95% interval on " << firstPOI->GetName() << " is : [" << interval->LowerLimit(*firstPOI) << ", "
     << interval->UpperLimit(*firstPOI) << "] " << endl;

RooArgSet tmpPOI(*firstPOI);
double observedUL = interval->UpperLimit(*firstPOI);
firstPOI->setVal(observedUL);
double obsTSatObsUL = fc.GetTestStatSampler()->EvaluateTestStatistic(*data, tmpPOI);

RooDataSet *parameterScan = (RooDataSet *)fc.GetPointsToScan();
RooArgSet *tmpPoint;

TH1F *histOfThresholds =
   new TH1F("histOfThresholds", "", parameterScan->numEntries(), firstPOI->getMin(), firstPOI->getMax());
histOfThresholds->GetXaxis()->SetTitle(firstPOI->GetName());
histOfThresholds->GetYaxis()->SetTitle("Threshold");

for (Int_t i = 0; i < parameterScan->numEntries(); ++i) {
   tmpPoint = (RooArgSet *)parameterScan->get(i)->clone("temp");
   // cout <<"get threshold"<<endl;
   double arMax = belt->GetAcceptanceRegionMax(*tmpPoint);
   double poiVal = tmpPoint->getRealValue(firstPOI->GetName());
   histOfThresholds->Fill(poiVal, arMax);
}
TCanvas *c1 = new TCanvas();
c1->Divide(2);
c1->cd(1);
histOfThresholds->SetMinimum(0);
histOfThresholds->Draw();
c1->cd(2);

std::unique_ptr<RooAbsReal> nll{mc->GetPdf()->createNLL(*data)};
std::unique_ptr<RooAbsReal> profile{nll->createProfile(*mc->GetParametersOfInterest())};
firstPOI->setVal(0.);
profile->getVal(); // this will do fit and set nuisance parameters to profiled values
RooArgSet *poiAndNuisance = new RooArgSet();
if (mc->GetNuisanceParameters())
   poiAndNuisance->add(*mc->GetNuisanceParameters());
poiAndNuisance->add(*mc->GetParametersOfInterest());
w->saveSnapshot("paramsToGenerateData", *poiAndNuisance);
RooArgSet *paramsToGenerateData = (RooArgSet *)poiAndNuisance->snapshot();
cout << "\nWill use these parameter points to generate pseudo data for bkg only" << endl;
paramsToGenerateData->Print("v");

RooArgSet unconditionalObs;
unconditionalObs.add(*mc->GetObservables());
unconditionalObs.add(*mc->GetGlobalObservables()); // comment this out for the original conditional ensemble

double CLb = 0;
double CLbinclusive = 0;

TH1F *histOfUL = new TH1F("histOfUL", "", 100, 0, firstPOI->getMax());
histOfUL->GetXaxis()->SetTitle("Upper Limit (background only)");
histOfUL->GetYaxis()->SetTitle("Entries");
for (int imc = 0; imc < nToyMC; ++imc) {

   // set parameters back to values for generating pseudo data
   //    cout << "\n get current nuis, set vals, print again" << endl;
   w->loadSnapshot("paramsToGenerateData");
   //    poiAndNuisance->Print("v");

   std::unique_ptr<RooDataSet> toyData;
   // now generate a toy dataset for the main measurement
   if (!mc->GetPdf()->canBeExtended()) {
      if (data->numEntries() == 1)
         toyData = std::unique_ptr<RooDataSet>{mc->GetPdf()->generate(*mc->GetObservables(), 1)};
      else
         cout << "Not sure what to do about this model" << endl;
   } else {
      //      cout << "generating extended dataset"<<endl;
      toyData = std::unique_ptr<RooDataSet>{mc->GetPdf()->generate(*mc->GetObservables(), Extended())};
   }

   // generate global observables
   // need to be careful for simpdf.
   // In ROOT 5.28 there is a problem with generating global observables
   // with a simultaneous PDF.  In 5.29 there is a solution with
   // RooSimultaneous::generateSimGlobal, but this may change to
   // the standard generate interface in 5.30.

   RooSimultaneous *simPdf = dynamic_cast<RooSimultaneous *>(mc->GetPdf());
   if (!simPdf) {
      std::unique_ptr<RooDataSet> one{mc->GetPdf()->generate(*mc->GetGlobalObservables(), 1)};
      const RooArgSet *values = one->get();
      std::unique_ptr<RooArgSet> allVars{mc->GetPdf()->getVariables()};
      allVars->assign(*values);
   } else {
      std::unique_ptr<RooDataSet> one{simPdf->generateSimGlobal(*mc->GetGlobalObservables(), 1)};
      const RooArgSet *values = one->get();
      std::unique_ptr<RooArgSet> allVars{mc->GetPdf()->getVariables()};
      allVars->assign(*values);
   }

   // get test stat at observed UL in observed data
   firstPOI->setVal(observedUL);
   double toyTSatObsUL = fc.GetTestStatSampler()->EvaluateTestStatistic(*toyData, tmpPOI);
   //    toyData->get()->Print("v");
   //    cout <<"obsTSatObsUL " <<obsTSatObsUL << "toyTS " << toyTSatObsUL << endl;
   if (obsTSatObsUL < toyTSatObsUL) // not sure about <= part yet
      CLb += (1.) / nToyMC;
   if (obsTSatObsUL <= toyTSatObsUL) // not sure about <= part yet
      CLbinclusive += (1.) / nToyMC;

   // loop over points in belt to find upper limit for this toy data
   double thisUL = 0;
   for (Int_t i = 0; i < parameterScan->numEntries(); ++i) {
      tmpPoint = (RooArgSet *)parameterScan->get(i)->clone("temp");
      double arMax = belt->GetAcceptanceRegionMax(*tmpPoint);
      firstPOI->setVal(tmpPoint->getRealValue(firstPOI->GetName()));
      //   double thisTS = profile->getVal();
      double thisTS = fc.GetTestStatSampler()->EvaluateTestStatistic(*toyData, tmpPOI);

      //   cout << "poi = " << firstPOI->getVal()
      // << " max is " << arMax << " this profile = " << thisTS << endl;
      //      cout << "thisTS = " << thisTS<<endl;
      if (thisTS <= arMax) {
         thisUL = firstPOI->getVal();
      } else {
         break;
      }
   }

   histOfUL->Fill(thisUL);

   // for few events, data is often the same, and UL is often the same
   //    cout << "thisUL = " << thisUL<<endl;
}
histOfUL->Draw();
c1->SaveAs("two-sided_upper_limit_output.pdf");

/*
SamplingDistPlot sampPlot;
int indexInScan = 0;
tmpPoint = (RooArgSet*) parameterScan->get(indexInScan)->clone("temp");
firstPOI->setVal( tmpPoint->getRealValue(firstPOI->GetName()) );
toymcsampler->SetParametersForTestStat(tmpPOI);
SamplingDistribution* samp = toymcsampler->GetSamplingDistribution(*tmpPoint);
sampPlot.AddSamplingDistribution(samp);
sampPlot.Draw();
   */

Double_t *bins = histOfUL->GetIntegral();
TH1F *cumulative = (TH1F *)histOfUL->Clone("cumulative");
cumulative->SetContent(bins);
double band2sigDown = 0, band1sigDown = 0, bandMedian = 0, band1sigUp = 0, band2sigUp = 0;
for (int i = 1; i <= cumulative->GetNbinsX(); ++i) {
   if (bins[i] < RooStats::SignificanceToPValue(2))
      band2sigDown = cumulative->GetBinCenter(i);
   if (bins[i] < RooStats::SignificanceToPValue(1))
      band1sigDown = cumulative->GetBinCenter(i);
   if (bins[i] < 0.5)
      bandMedian = cumulative->GetBinCenter(i);
   if (bins[i] < RooStats::SignificanceToPValue(-1))
      band1sigUp = cumulative->GetBinCenter(i);
   if (bins[i] < RooStats::SignificanceToPValue(-2))
      band2sigUp = cumulative->GetBinCenter(i);
}
cout << "-2 sigma  band " << band2sigDown << endl;
cout << "-1 sigma  band " << band1sigDown << " [Power Constraint)]" << endl;
cout << "median of band " << bandMedian << endl;
cout << "+1 sigma  band " << band1sigUp << endl;
cout << "+2 sigma  band " << band2sigUp << endl;

// print out the interval on the first Parameter of Interest
cout << "\nobserved 95% upper-limit " << interval->UpperLimit(*firstPOI) << endl;
cout << "CLb strict [P(toy>obs|0)] for observed 95% upper-limit " << CLb << endl;
cout << "CLb inclusive [P(toy>=obs|0)] for observed 95% upper-limit " << CLbinclusive << endl;

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