/////////////////////////////////////////////////////////////////////////
//
// 'Limit Example' RooStats tutorial macro #101
// author: Kyle Cranmer
// date June. 2009
//
// This tutorial shows an example of creating a simple
// model for a number counting experiment with uncertainty
// on both the background rate and signal efficeincy. We then 
// use a Confidence Interval Calculator to set a limit on the signal.
//
//
/////////////////////////////////////////////////////////////////////////

#ifndef __CINT__
#include "RooGlobalFunc.h"
#endif

#include "RooProfileLL.h"
#include "RooAbsPdf.h"
#include "RooStats/HypoTestResult.h"
#include "RooRealVar.h"
#include "RooPlot.h"
#include "RooDataSet.h"
#include "RooTreeDataStore.h"
#include "TTree.h"
#include "TCanvas.h"
#include "TLine.h"
#include "TStopwatch.h"

#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/UniformProposal.h"
#include "RooStats/FeldmanCousins.h"
#include "RooStats/NumberCountingPdfFactory.h"
#include "RooStats/ConfInterval.h"
#include "RooStats/PointSetInterval.h"
#include "RooStats/LikelihoodInterval.h"
#include "RooStats/LikelihoodIntervalPlot.h"
#include "RooStats/RooStatsUtils.h"
#include "RooStats/ModelConfig.h"
#include "RooStats/MCMCInterval.h"
#include "RooStats/MCMCIntervalPlot.h"
#include "RooStats/ProposalFunction.h"
#include "RooStats/ProposalHelper.h"
#include "RooFitResult.h"


// use this order for safety on library loading
using namespace RooFit ;
using namespace RooStats ;


void rs101_limitexampleNoUnc()
{
  /////////////////////////////////////////
  // An example of setting a limit in a number counting experiment with uncertainty on background and signal
  /////////////////////////////////////////

  // to time the macro
  TStopwatch t;
  t.Start();

  /////////////////////////////////////////
  // The Model building stage
  /////////////////////////////////////////
  RooWorkspace* wspace = new RooWorkspace();
  wspace->factory("Poisson::countingModel(obs[0,30], sum(s[0,0,30],b[0,0,200]))"); // counting model
  wspace->Print();

  RooAbsPdf* modelWithConstraints = wspace->pdf("countingModel"); // get the model
  RooRealVar* obs = wspace->var("obs"); // get the observable
  RooRealVar* s = wspace->var("s"); // get the signal we care about
  RooRealVar* b = wspace->var("b"); // get the background and set it to a constant. 
  b->setConstant();

  obs->setVal(0);
  RooDataSet* data = new RooDataSet("exampleData", "exampleData", RooArgSet(*obs));
  data->add(*obs);

  // not necessary
  modelWithConstraints->fitTo(*data);

  // Now let's make some confidence intervals for s, our parameter of interest
  RooArgSet paramOfInterest(*s);

  ModelConfig modelConfig(new RooWorkspace());
  modelConfig.SetPdf(*modelWithConstraints);
  modelConfig.SetParametersOfInterest(paramOfInterest);


  // First, let's use a Calculator based on the Profile Likelihood Ratio
  //ProfileLikelihoodCalculator plc(*data, *modelWithConstraints, paramOfInterest); 
  ProfileLikelihoodCalculator plc(*data, modelConfig); 
  plc.SetTestSize(.05);
  ConfInterval* lrint = plc.GetInterval();  // that was easy.

  // Let's make a plot
  TCanvas* dataCanvas = new TCanvas("dataCanvas");
  dataCanvas->Divide(2,1);

  dataCanvas->cd(1);
  LikelihoodIntervalPlot plotInt((LikelihoodInterval*)lrint);
  plotInt.SetTitle("Profile Likelihood Ratio and Posterior for S");
  plotInt.Draw();

  // Get Lower and Upper limits from Profile Calculator
  cout << "Profile lower limit on s = " << ((LikelihoodInterval*) lrint)->LowerLimit(*s) << endl;
  cout << "Profile upper limit on s = " << ((LikelihoodInterval*) lrint)->UpperLimit(*s) << endl;

  RooFitResult* fit = modelWithConstraints->fitTo(*data, Save(true));

  // Third, use a Calculator based on Markov Chain monte carlo
  // Before configuring the calculator, let's make a ProposalFunction
  // that will achieve a high acceptance rate
  ProposalHelper ph;
  ph.SetVariables((RooArgSet&)fit->floatParsFinal());
  ph.SetCovMatrix(fit->covarianceMatrix());
  ph.SetUpdateProposalParameters(true);
  ph.SetCacheSize(100);
  ProposalFunction* pdfProp = ph.GetProposalFunction();  // that was easy

  MCMCCalculator mc(*data, modelConfig);
  mc.SetNumIters(20000); // steps to propose in the chain
  mc.SetTestSize(.05); // 95% CL
  mc.SetNumBurnInSteps(100); // ignore first N steps in chain as "burn in"
  mc.SetProposalFunction(*pdfProp);
  mc.SetLeftSideTailFraction(0.0);  // find a "central" interval
  MCMCInterval* mcInt = (MCMCInterval*)mc.GetInterval();  // that was easy

  // Plot MCMC interval and print some statistics
  MCMCIntervalPlot mcPlot(*mcInt);
  mcPlot.SetLineColor(kMagenta);
  mcPlot.SetLineWidth(2);
  mcPlot.Draw("same");

  double mcul = mcInt->UpperLimit(*s);
  double mcll = mcInt->LowerLimit(*s);
  cout << "MCMC lower limit on s = " << mcll << endl;
  cout << "MCMC upper limit on s = " << mcul << endl;
  cout << "MCMC Actual confidence level: "
     << mcInt->GetActualConfidenceLevel() << endl;

  // 3-d plot of the parameter points
  dataCanvas->cd(2);
  // also plot the points in the markov chain
  TTree& chain =  ((RooTreeDataStore*) mcInt->GetChainAsDataSet()->store())->tree();
  chain.SetMarkerStyle(6);
  chain.SetMarkerColor(kRed);


  delete wspace;
  delete lrint;
  delete mcInt;
  delete data;


}