/// \file
/// \ingroup tutorial_roostats
/// \notebook
/// Neutrino Oscillation Example from Feldman & Cousins
///
/// This tutorial shows a more complex example using the FeldmanCousins utility
/// to create a confidence interval for a toy neutrino oscillation experiment.
/// The example attempts to faithfully reproduce the toy example described in Feldman & Cousins'
/// original paper, Phys.Rev.D57:3873-3889,1998.
///
/// \macro_image
/// \macro_output
/// \macro_code
///
/// \author Kyle Cranmer

#include "RooGlobalFunc.h"
#include "RooStats/ConfInterval.h"
#include "RooStats/FeldmanCousins.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/UniformProposal.h"
#include "RooStats/LikelihoodIntervalPlot.h"
#include "RooStats/MCMCIntervalPlot.h"
#include "RooStats/MCMCInterval.h"

#include "RooDataSet.h"
#include "RooDataHist.h"
#include "RooRealVar.h"
#include "RooConstVar.h"
#include "RooAddition.h"
#include "RooProduct.h"
#include "RooProdPdf.h"
#include "RooAddPdf.h"

#include "TROOT.h"
#include "RooPolynomial.h"
#include "RooRandom.h"

#include "RooNLLVar.h"
#include "RooProfileLL.h"

#include "RooPlot.h"

#include "TCanvas.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TTree.h"
#include "TMarker.h"
#include "TStopwatch.h"

#include <iostream>

using namespace RooFit;
using namespace RooStats;

void FC_2D_test_2(bool doFeldmanCousins = true)

{

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

   // Taken from Feldman & Cousins paper, Phys.Rev.D57:3873-3889,1998.
   // e-Print: physics/9711021 (see page 13.)
   
   // Make signal model model

  double bm_min(5.25), bm_max(5.55);

  RooRealVar  Bmass("Bmass","m(B_{s}) (GeV/c^{2})", 5.25,5.55);

  //----------------POI---------------------------//
  RooRealVar  mean("mean","mean", 5.367,5.25,5.55);
  RooRealVar  sigma1("sigma1","sigma1",0.02541,0.0,0.05);
 //-----------------------------------------------//
 
  RooRealVar  sigma2("sigma2","sigma2",0.04818);
  RooRealVar sigM_frac("sigM_frac","fraction of Gaussians",0.55026);
  RooRealVar alpha_1("alpha_1", "alpha_1",1.9828364);
  RooRealVar alpha_2("alpha_2", "alpha_2",-1.9279166);
  RooRealVar n_1("n_1", "n_1",1.0);
  RooRealVar n_2("n_2", "n_2",10.0);

  RooRealVar mean_pb("mean_pb","mean of the Peaking Bkg", 5.4);
  RooRealVar sigma1_pb("sigma1_pb","sigma of the Peaking Bkg",0.0761);
  RooRealVar sigma2_pb("sigma2_pb","sigma of the Peaking Bkg",0.052);
  RooRealVar sigM_frac_pb("sigM_frac_pb","fraction",0.679);
  RooRealVar alpha1_pb("alpha1_pb", "alpha_1",1.35);
  RooRealVar alpha2_pb("alpha2_pb", "alpha_2",-0.26);
  RooRealVar n1_pb("n1_pb", "n_1",0.50);
  RooRealVar n2_pb("n2_pb", "n_2",9.9);
  RooRealVar fsfd("fsfd","fsfd",0.0411);

  RooCBShape cbs_1("cbs_1", "cbs_1", Bmass, mean, sigma1, alpha_1, n_1);  ///DCB PDF
  RooCBShape cbs_2("cbs_2", "cbs_2", Bmass, mean, sigma2, alpha_2, n_2);

  RooAddPdf sigG("sigG","sigG", RooArgList(cbs_1,cbs_2), RooArgList(sigM_frac) );

  RooCBShape cbs1_pb("cbs1_pb", "cbs1_pb", Bmass, mean_pb, sigma1_pb, alpha1_pb, n1_pb);  ///DCB Peaking Bkg PDF
  RooCBShape cbs2_pb("cbs2_pb", "cbs2_pb", Bmass, mean_pb, sigma2_pb, alpha2_pb, n2_pb);

  RooAddPdf sigG_pb("sigG_pb","sigG_pb", RooArgList(cbs1_pb,cbs2_pb), RooArgList(sigM_frac_pb) ); ///Peaking Bkg Model

  RooConstVar lambda("lambda","slope",-3.10);
  RooExponential bkgE("bkgE","exponential PDF", Bmass, lambda);


  RooConstVar nsigG("nsigG", "signal events", 157);
  RooConstVar nbkgE("nbkgE", "number of exponential bkg events", 68);
  RooProduct nsigG_pb("nsigG_pb","Peaking Bkg events",RooArgList(nsigG,fsfd));


  // total model
  RooAddPdf model("model","S+Exp+peakBkg", RooArgList(sigG,bkgE,sigG_pb), RooArgList(nsigG,nbkgE,nsigG_pb));    ///Total Model

  model.Print("t");

   // for debugging, check model tree
  model.printCompactTree();
  model.graphVizTree("model.dot");

   // turn off some messages
   RooMsgService::instance().setStreamStatus(0, kFALSE);
   RooMsgService::instance().setStreamStatus(1, kFALSE);
   RooMsgService::instance().setStreamStatus(2, kFALSE);

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

   // n events in data to data, simply sum of sig+bkg
   Int_t nEventsData = nbkgE.getVal() + nsigG.getVal() +  nsigG.getVal()*fsfd.getVal() ;
   cout << "generate toy data with nEvents = " << nEventsData << endl;

   // adjust random seed to get a toy dataset similar to one in paper.
   // Found by trial and error (3 trials, so not very "fine tuned")

   RooRandom::randomGenerator()->SetSeed(3);

   // create a toy dataset
   RooDataSet *data = model.generate(RooArgSet(Bmass), nEventsData);

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

   //------------------------- make some plots------------------------//

   TCanvas *c = new TCanvas("c");
   c->Divide(2, 2);

   // --------------------------plot the PDF------------------------//

   c->cd(1);
   TH1 *hh = sigG.createHistogram("hh", Bmass, Binning(40)/*, YVar(sigma1, Binning(40)), Scaling(kFALSE)*/);
   hh->SetLineColor(kBlue);
   hh->SetTitle("True Signal Model");
   //hh->Draw("surf");
   hh->Draw();

   // ------------------plot the data with the best fit------------ //

   c->cd(2);
   RooPlot *Bframe = Bmass.frame();
   data->plotOn(Bframe);
   model.fitTo(*data, Extended());
   model.plotOn(Bframe);
   model.plotOn(Bframe, RooFit::Name("sigG"), Components("sigG"), /*FillStyle(3013), FillColor(kGreen), DrawOption("F")*/LineColor(kGreen) ); //pdf for only signal
   model.plotOn(Bframe, RooFit::Name("sigG_pb"), Components("sigG_pb"), /*FillStyle(3013), FillColor(kOrange), DrawOption("F"),*/ LineColor(kOrange)); //pdf for only signal
   model.plotOn(Bframe, RooFit::Name("bkgE"), Components("bkgE"),LineStyle(9), LineColor(kRed),LineWidth(5) );
   model.plotOn(Bframe);
   Bframe->SetTitle("toy data with best fit model (and sig+bkg components)");
   Bframe->Draw();

   // ------------------plot the likelihood function--------------- //

   c->cd(3);
   RooNLLVar nll("nll", "nll", model, *data, Extended());
   RooProfileLL pll("pll", "", nll, RooArgSet(mean, sigma1));
   TH1 *hhh = pll.createHistogram("hhh", sigma1, Binning(20), YVar(mean, Binning(20)), Scaling(kFALSE));
   hhh->SetLineColor(kRed);
   hhh->SetTitle("Likelihood Function");
   hhh->Draw("surf");

   //c->Update();


   // -------------------------------------------------------------- //
   // show use of Feldman-Cousins utility in RooStats
   // set the distribution creator, which encodes the test statistic

   RooArgSet parameters(mean, sigma1);
   RooWorkspace *w = new RooWorkspace();

   ModelConfig modelConfig;
   modelConfig.SetWorkspace(*w);
   modelConfig.SetPdf(model);
   modelConfig.SetParametersOfInterest(parameters);

   RooStats::FeldmanCousins fc(*data, modelConfig);
   //fc.SetTestSize(.1); // set size of test
   fc.SetTestSize(.05); // 
   fc.UseAdaptiveSampling(true);
   fc.FluctuateNumDataEntries(false);
   //fc.SetNBins(10); // number of points to test per parameter
   fc.SetNBins(20); // 

   // use the Feldman-Cousins tool
   ConfInterval *interval = 0;
   if (doFeldmanCousins)
      interval = fc.GetInterval();

   // -------------------------------------------------------------- //
   // show use of ProfileLikeihoodCalculator utility in RooStats

   RooStats::ProfileLikelihoodCalculator plc(*data, modelConfig);
   plc.SetTestSize(.1);
   double CL = 0.683;
   plc.SetConfidenceLevel(CL);
   ConfInterval *plcInterval = plc.GetInterval();

   RooStats::ProfileLikelihoodCalculator plc1(*data, modelConfig);
   plc1.SetTestSize(.1);
   double CL1 = 0.90;
   plc1.SetConfidenceLevel(CL1);
   ConfInterval *plcInterval1 = plc1.GetInterval();

   RooStats::ProfileLikelihoodCalculator plc2(*data, modelConfig);
   plc2.SetTestSize(.1);
   double CL2 = 0.95;
   plc2.SetConfidenceLevel(CL2);
   ConfInterval *plcInterval2 = plc2.GetInterval();


   // -------------------- make plot of resulting interval-------------------//

   c->cd(4);

   // first plot a small dot for every point tested
   if (doFeldmanCousins) {
      RooDataHist *parameterScan = (RooDataHist *)fc.GetPointsToScan();
      TH2F *hist = (TH2F *)parameterScan->createHistogram("mean:sigma1", 30, 30);
      //hist->Draw();
      TH2F *forContour = (TH2F *)hist->Clone();

      // now loop through the points and put a marker if it's in the interval
      RooArgSet *tmpPoint;
      // loop over points to test
      for (Int_t i = 0; i < parameterScan->numEntries(); ++i) {
         // get a parameter point from the list of points to test.
         tmpPoint = (RooArgSet *)parameterScan->get(i)->clone("temp");

         if (interval) {
            if (interval->IsInInterval(*tmpPoint)) {
               forContour->SetBinContent(
                  hist->FindBin(tmpPoint->getRealValue("mean"), tmpPoint->getRealValue("sigma1")), 1);
            } else {
               forContour->SetBinContent(
                  hist->FindBin(tmpPoint->getRealValue("mean"), tmpPoint->getRealValue("sigma1")), 0);
            }
         }

         delete tmpPoint;
      }

      if (interval) {
         //Double_t level = 0.5;  //CL
	 Double_t level = 0.683;  //CL
         forContour->SetContour(1, &level);
         forContour->SetLineWidth(2);
         forContour->SetLineColor(kRed);
         forContour->Draw("cont2");
      }
   }


////////////////////////////////////////////////  LikelihoodIntervalPlot  //////////////////////////////////////////////////

   LikelihoodIntervalPlot plotInt((LikelihoodInterval *)plcInterval);
   plotInt.SetTitle("Confidence Intervals");
   //plotInt.SetLineWidth(2);
   plotInt.SetLineColor(kGreen);

   LikelihoodIntervalPlot plotInt1((LikelihoodInterval *)plcInterval1);
   plotInt1.SetTitle("90% Confidence Intervals");
   //plotInt1.SetLineStyle(10);
   plotInt1.SetLineColor(kOrange);

   LikelihoodIntervalPlot plotInt2((LikelihoodInterval *)plcInterval2);
   plotInt2.SetTitle("95% Confidence Intervals");
   //plotInt2.SetLineStyle(5);
   plotInt2.SetLineColor(kCyan);

   plotInt.Draw("same");
   plotInt1.Draw("same");
   plotInt2.Draw("same");


  //dataCanvas->Update();

   /// print timing info
   t.Stop();
   t.Print();
}