#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;

// -------------------------------------------------------

void TwoSidedFrequentistUpperLimitWithBands(const char *infile = "", const char *workspaceName = "combined",
                                            const char *modelConfigName = "ModelConfig",
                                            const char *dataName = "obsData")
{

   double confidenceLevel = 0.95;
   // degrade/improve number of pseudo-experiments used to define the confidence belt.
   // value of 1 corresponds to default number of toys in the tail, which is 50/(1-confidenceLevel)
   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

   // -------------------------------------------------------
   // First part is just to access a user-defined file
   // or create the standard example file if it doesn't exist
   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;

   // Try to open the file
   TFile *inputFile = TFile::Open(filename);

   // -------------------------------------------------------
   // Now get the data and workspace

   // get the workspace out of the file
   RooWorkspace *w = (RooWorkspace *)inputFile->Get(workspaceName);

   // get the modelConfig out of the file
   ModelConfig *mc = (ModelConfig *)w->obj(modelConfigName);

   // get the modelConfig out of the file
   RooAbsData *data = w->data(dataName);

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

   // -------------------------------------------------------
   // Now get the POI for convenience
   // you may want to adjust the range of your POI
   RooRealVar *firstPOI = (RooRealVar *)mc->GetParametersOfInterest()->first();
   /*  firstPOI->setMin(0);*/
   /*  firstPOI->setMax(10);*/

   // -------------------------------------------------------
   // create and use the FeldmanCousins tool
   // to find and plot the 95% confidence interval
   // on the parameter of interest as specified
   // in the model config
   // REMEMBER, we will change the test statistic
   // so this is NOT a Feldman-Cousins interval
   FeldmanCousins fc(*data, *mc);
   fc.SetConfidenceLevel(confidenceLevel);
   fc.AdditionalNToysFactor(additionalToysFac); // improve sampling that defines confidence belt
   //  fc.UseAdaptiveSampling(true); // speed it up a bit, but don't use for expected limits
   fc.SetNBins(nPointsToScan); // set how many points per parameter of interest to scan
   fc.CreateConfBelt(true);    // save the information in the belt for plotting

   // -------------------------------------------------------
   // Feldman-Cousins is a unified limit by definition
   // but the tool takes care of a few things for us like which values
   // of the nuisance parameters should be used to generate toys.
   // so let's just change the test statistic and realize this is
   // no longer "Feldman-Cousins" but is a fully frequentist Neyman-Construction.
   //  fc.GetTestStatSampler()->SetTestStatistic(&onesided);
   // ((ToyMCSampler*) fc.GetTestStatSampler())->SetGenerateBinned(true);
   ToyMCSampler *toymcsampler = (ToyMCSampler *)fc.GetTestStatSampler();
   ProfileLikelihoodTestStat *testStat = dynamic_cast<ProfileLikelihoodTestStat *>(toymcsampler->GetTestStatistic());

   // Since this tool needs to throw toy MC the PDF needs to be
   // extended or the tool needs to know how many entries in a dataset
   // per pseudo experiment.
   // In the 'number counting form' where the entries in the dataset
   // are counts, and not values of discriminating variables, the
   // datasets typically only have one entry and the PDF is not
   // extended.
   if (!mc->GetPdf()->canBeExtended()) {
      if (data->numEntries() == 1)
         fc.FluctuateNumDataEntries(false);
      else
         cout << "Not sure what to do about this model" << endl;
   }

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

   // Now get the interval
   PointSetInterval *interval = fc.GetInterval();
   ConfidenceBelt *belt = fc.GetConfidenceBelt();

   // print out the interval on the first Parameter of Interest
   cout << "\n95% interval on " << firstPOI->GetName() << " is : [" << interval->LowerLimit(*firstPOI) << ", "
        << interval->UpperLimit(*firstPOI) << "] " << endl;

   // get observed UL and value of test statistic evaluated there
   RooArgSet tmpPOI(*firstPOI);
   double observedUL = interval->UpperLimit(*firstPOI);
   firstPOI->setVal(observedUL);
   double obsTSatObsUL = fc.GetTestStatSampler()->EvaluateTestStatistic(*data, tmpPOI);

   // Ask the calculator which points were scanned
   RooDataSet *parameterScan = (RooDataSet *)fc.GetPointsToScan();
   RooArgSet *tmpPoint;

   // make a histogram of parameter vs. threshold
   TH1F *histOfThresholds =
      new TH1F("histOfThresholds", "", parameterScan->numEntries(), firstPOI->getMin(), firstPOI->getMax());
   histOfThresholds->GetXaxis()->SetTitle(firstPOI->GetName());
   histOfThresholds->GetYaxis()->SetTitle("Threshold");

   // loop through the points that were tested and ask confidence belt
   // what the upper/lower thresholds were.
   // For FeldmanCousins, the lower cut off is always 0
   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);

   // -------------------------------------------------------
   // Now we generate the expected bands and power-constraint

   // First: find parameter point for mu=0, with conditional MLEs for nuisance parameters
   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;

   // Now we generate background only and find distribution of upper limits
   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
      std::unique_ptr<RooDataSet> one{mc->GetPdf()->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");

   // if you want to see a plot of the sampling distribution for a particular scan point:
   /*
   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();
      */

   // Now find bands and power constraint
   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;
}