#include "TCanvas.h"
#include "RooWorkspace.h"
#include "RooAbsReal.h"
#include "RooRealVar.h"
#include "RooArgSet.h"
#include "RooArgList.h"
#include "RooDataSet.h"
#include "RooPlot.h"
#include "RooStats/BayesianCalculator.h"
#include "RooStats/ModelConfig.h"
#include "RooStats/SimpleInterval.h"
#include "RooStats/SequentialProposal.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/MCMCInterval.h"
#include "RooStats/MCMCIntervalPlot.h"

using namespace RooFit;
using namespace RooStats;

double CalcLimit(){
    if(make_RooFit_workspace()){
         double upperxs = AsymptLimits("workspace.root","myWS");
    }return 0;
}

int make_RooFit_workspace(){
   // read rooFit libraries
   // gSystem->Load("libRooFit");
   
    // create workspace
   RooWorkspace * pWs = new RooWorkspace("myWS");

   ////////////////////////////////////////
   // begin
   //define model functions and pdfs
   //////////////////////////////////////

   // observable: number of events
   pWs->factory( "n[100000.0]" );

   // Event number - parameter of interest
   pWs->factory( "BSMAnz[0.0,50000.0]" );

   // background yield with systematics
   pWs->factory( "nbkg_nom[100000.0,50000,200000]" );

   pWs->factory( "nbkg_kappa[1.20]" );
   pWs->factory( "expr::alpha_nbkg('pow(nbkg_kappa,beta_nbkg)',nbkg_kappa,beta_nbkg[0,-5,5])" );
   pWs->factory( "prod::nbkg(nbkg_nom,alpha_nbkg)" );
   pWs->factory( "Gaussian::constr_nbkg(beta_nbkg,glob_nbkg[0,-5,5],1)" );

   // Add signal and background expectation
   pWs->factory( "sum::yield(BSMAnz,nbkg)" );

   // Core model: Poisson probability with mean signal+bkg
   pWs->factory( "Poisson::model_core(n,yield)" );

   // model with systematics
   pWs->factory( "PROD::model(model_core,constr_nbkg)" );


   // create set of observables (will need it for datasets and ModelConfig later)
   RooRealVar * pObs = pWs->var("n"); // get the pointer to the observable
   RooArgSet obs("observables");
   obs.add(*pObs);

   // create the dataset
   pObs->setVal(100000); // this is your observed data:
   RooDataSet * data = new RooDataSet("data", "data", obs);
   data->add( *pObs );

   // import dataset into workspace
   pWs->import(*data);

   ////////////////////////////////////////
   // end
   //define model functions and pdfs
   //////////////////////////////////////

   ////////////////////////////////////////
   // begin
   //setup model config
   //////////////////////////////////////

   // create set of global observables (need to be defined as constants!)
   pWs->var("glob_nbkg")->setConstant(true);
   //pWs->var("n")->setConstant(true);
   RooArgSet globalObs("global_obs");
   globalObs.add( *pWs->var("glob_nbkg") );
   //globalObs.add( *pWs->var("n") );

   // create set of parameters of interest (POI)
   RooArgSet poi("poi");
   pWs->factory( "Uniform::uniform(BSMAnz)" );
   poi.add( *pWs->var("BSMAnz") );

   // create set of nuisance parameters
   RooArgSet nuis("nuis");
   nuis.add( *pWs->var("beta_nbkg") );

   // fix all other variables in model:
   // everything except observables, POI, and nuisance parameters
   // must be constant
   pWs->var("nbkg_nom")->setConstant(true);
   pWs->var("nbkg_kappa")->setConstant(true);

   // create signal+background Model Config
   RooStats::ModelConfig sbHypo("SbHypo");
   sbHypo.SetWorkspace( *pWs );
   //sbHypo.SetPdf( *pWs->pdf("model_core") );
   sbHypo.SetPdf( *pWs->pdf("model") );
   sbHypo.SetObservables( obs );
   sbHypo.SetGlobalObservables( globalObs );
   sbHypo.SetParametersOfInterest( poi );
   sbHypo.SetNuisanceParameters( nuis );

   // set parameter snapshot that corresponds to the best fit to data
   RooAbsReal * pNll = sbHypo.GetPdf()->createNLL( *data );
   RooAbsReal * pProfile = pNll->createProfile( globalObs ); // do not profile global observables
   pProfile->getVal(); // this will do fit and set POI and nuisance parameters to fitted values
   RooArgSet * pPoiAndNuisance = new RooArgSet("poiAndNuisance");
   pPoiAndNuisance->add(*sbHypo.GetNuisanceParameters());
   pPoiAndNuisance->add(*sbHypo.GetParametersOfInterest());
   sbHypo.SetSnapshot(*pPoiAndNuisance);

   //Create background only hypothesis
   RooRealVar* paraOI = (RooRealVar*) sbHypo.GetParametersOfInterest()->first();

   ModelConfig * bModel = (ModelConfig*) sbHypo.Clone();
   bModel->SetName(TString("BModel"));
   double oldval = paraOI->getVal();
   paraOI->setVal(0);
   bModel->SetSnapshot( *paraOI  );
   paraOI->setVal(oldval);

   delete pProfile;
   delete pNll;
   delete pPoiAndNuisance;

   // import ModelConfig into workspace
   pWs->import( sbHypo );
   pWs->import( *bModel );
   ////////////////////////////////////////
   // end
   //setup model config
   //////////////////////////////////////

   // save workspace to file
   pWs -> SaveAs("workspace.root");
   return 1;


}

double AsymptLimits( std::string filename = "workspace.root",
			  std::string wsname = "myWS"){
   //
   // this function loads a workspace and computes
   // an asymptotic Limit 
   //
   // open file with workspace for reading
   TFile * pInFile = new TFile(filename.c_str(), "read");
   
   // load workspace
   RooWorkspace * pWs = (RooWorkspace *)pInFile->Get(wsname.c_str());
   if (!pWs){
   std::cout << "workspace " << wsname
         << " not found" << std::endl;
   return -1;
   }

   // printout workspace content
   pWs->Print();

   // load and print data from workspace
   RooAbsData * data = pWs->data("data");
   data->Print();

   // load and print S+B Model Config
   RooStats::ModelConfig * pSbHypo = (RooStats::ModelConfig *)pWs->obj("SbHypo");
   RooStats::ModelConfig * bHypo = (RooStats::ModelConfig *)pWs->obj("BModel");
   pSbHypo->Print();
   bHypo->Print();

   RooStats::AsymptoticCalculator ac( *data,*bHypo, *pSbHypo, true);

   ac.SetOneSided(true);

   HypoTestInverter calc(ac);

   // set confidence level (e.g. 95% upper limits)
   calc.SetConfidenceLevel(0.95);
   // for CLS
   bool useCLs = true;
   calc.UseCLs(useCLs);
   calc.SetVerbose(false);

   // profile likelihood test statistics
   ProfileLikelihoodTestStat profll( *pSbHypo->GetPdf() );
   // for CLs (bounded intervals) use one-sided profile likelihood
   if (useCLs) profll.SetOneSided(true);

   RooRealVar* paraOI = (RooRealVar*) pSbHypo->GetParametersOfInterest()->first();

   int npoints = 20;  // number of points to scan
   // min and max (better to choose smaller intervals)
   double poimin = paraOI->getMin();
   double poimax = paraOI->getMax();

   std::cout << "Doing a fixed scan  in interval : " << poimin << " , " << poimax << std::endl;
   calc.SetFixedScan(npoints,poimin,poimax);

   HypoTestInverterResult * r = calc.GetInterval();

   std::cout << "Expected upper limits, using the B (alternate) model : " << std::endl;
   std::cout << " expected limit (median) " << r->GetExpectedUpperLimit(0) << std::endl;
   std::cout << " expected limit (-1 sig) " << r->GetExpectedUpperLimit(-1) << std::endl;
   std::cout << " expected limit (+1 sig) " << r->GetExpectedUpperLimit(1) << std::endl;
   std::cout << " expected limit (-2 sig) " << r->GetExpectedUpperLimit(-2) << std::endl;
   std::cout << " expected limit (+2 sig) " << r->GetExpectedUpperLimit(2) << std::endl;
   std::cout << " observed limit " << r->UpperLimit() << std::endl;
   
   return 0.0;
}