//////////////////////////////////////////////////////////////////////////
//
// 'SPECIAL PDFS' RooFit tutorial macro #701
// 
// Unbinned maximum likelihood fit of an efficiency eff(x) function to 
// a dataset D(x,cut), where cut is a category encoding a selection, of which
// the efficiency as function of x should be described by eff(x)
//
// 07/2008 - Wouter Verkerke 
// 04/2009 - Max Baak
//
/////////////////////////////////////////////////////////////////////////

#ifndef __CINT__
#include "RooGlobalFunc.h"
#endif
#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "RooFormulaVar.h"
#include "RooProdPdf.h"
#include "RooEfficiency.h"
#include "RooPolynomial.h"
#include "RooCategory.h"
#include "TCanvas.h"
#include "RooPlot.h"
using namespace RooFit ;


void pullsefficiencyfit()
{
  gStyle->SetOptFit(1);
  gStyle->SetOptStat(0);
  gStyle->SetOptTitle(0);


  // C o n s t r u c t   e f f i c i e n c y   f u n c t i o n   e ( x ) 
  // -------------------------------------------------------------------

  // Declare variables x,mean,sigma with associated name, title, initial value and allowed range
  RooRealVar x("x","x",0,60) ;

  // Efficiency function eff(x;a,b) 
  RooRealVar a("a","a",0.99,0,10) ;
  RooRealVar b("b","b",20,0,50) ;
  RooRealVar c("c","c",3,-10,10) ;
  //RooFormulaVar effFunc("effFunc","(1-a)+a*cos((x-c)/b)",RooArgList(a,b,c,x)) ;
  RooFormulaVar effFunc("effFunc","a/(1+exp(-(x-b)/c))",RooArgList(a,b,c,x)) ;

  RooRealVar wtest("wtest","wtest",-1,3);

  // C o n s t r u c t   c o n d i t i o n a l    e f f i c i e n c y   p d f   E ( c u t | x ) 
  // ------------------------------------------------------------------------------------------

  // Acceptance state cut (1 or 0)
  RooCategory cut("cut","cutr") ;
  cut.defineType("accept",1) ;
  cut.defineType("reject",0) ;

  // Construct efficiency p.d.f eff(cut|x)
  RooEfficiency effPdf("effPdf","effPdf",effFunc,cut,"accept") ;



  // G e n e r a t e   d a t a   ( x ,   c u t )   f r o m   a   t o y   m o d e l 
  // -----------------------------------------------------------------------------

  // Construct global shape p.d.f shape(x) and product model(x,cut) = eff(cut|x)*shape(x) 
  // (These are _only_ needed to generate some toy MC here to be used later)
  //RooPolynomial shapePdf("shapePdf","shapePdf",x,RooConst(-0.095)) ;
  RooPolynomial shapePdf("shapePdf","shapePdf",x) ;
  RooProdPdf model("model","model",shapePdf,Conditional(effPdf,cut)) ;

  TH1D* hist0 = new TH1D("hist0","hist0",100,-6,6);
  TH1D* hist1 = new TH1D("hist1","hist1",100,-6,6);
  TH1D* hist2 = new TH1D("hist2","hist2",100,-6,6);
  TH1D* hist3 = new TH1D("hist3","hist3",100,-6,6);

  // -------------------------------------------------------------------------------
  for (int k=0; k<1000; ++k)
  {
    a.setVal(0.99);
    b.setVal(20.);
    c.setVal(3.);
    a.setConstant(false);
    b.setConstant(false);
    c.setConstant(false);

    // Generate some toy data from model
    RooDataSet* data = model.generate(RooArgSet(x,cut),10000) ;
    
    // Generate event weights
    RooDataSet* ds = new RooDataSet("ds","ds",RooArgSet(wtest));
    for (Int_t i=0 ; i<data->numEntries() ; i++) { 
      double myrand = 4.*RooRandom::uniform()-1.; 
      wtest.setVal(myrand); 
      ds->add(RooArgSet(wtest)) ; 
    }
    data->merge(ds); 
    
    // Instruct dataset d in interpret w as event weight rather than as observable
    data->setWeightVar(wtest) ;
    data->Print("v");
    
    // F i t   c o n d i t i o n a l   e f f i c i e n c y   p d f   t o   d a t a 
    // --------------------------------------------------------------------------
    
    // Fit conditional efficiency p.d.f to data
    //effPdf.fitTo(*data,ConditionalObservables(x)) ;
    effPdf.fitTo(*data,ConditionalObservables(x),SumW2Error(kTRUE)) ;
    
    hist1->Fill( (a.getVal()-0.99)/a.getError() );
    hist2->Fill( (b.getVal()-20.)/b.getError() );
    hist3->Fill( (c.getVal()-3.)/c.getError() );
    
    delete data;
    delete ds;
  }
  
  // -------------------------------------------------------------------------------
  
  TCanvas* c1 = new TCanvas("weightfix","weightfix",1000,400) ;

  c1->Divide(3);
  c1->cd(1);
  hist1->Draw();
  c1->cd(2);
  hist2->Draw();
  c1->cd(3);
  hist3->Draw();
  //c1->cd(4);
  //hist1->Draw();

  c1->Draw();

}