/////////////////////////////////////////////////////////////////////////
//
// 'LIKELIHOOD AND MINIMIZATION' RooFit tutorial macro #604
// 
// Fitting with constraints
//
//
// 07/2008 - Wouter Verkerke 
//
/////////////////////////////////////////////////////////////////////////

#ifndef __CINT__
#include "RooGlobalFunc.h"
#endif
#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "RooConstVar.h"
#include "RooPolynomial.h"
#include "RooAddPdf.h"
#include "RooProdPdf.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "TCanvas.h"
#include "TAxis.h"
#include "TH1.h"
using namespace RooFit ;


void rf604_constraints()
{

  // C r e a t e   m o d e l  a n d   d a t a s e t 
  // ----------------------------------------------

  // Construct a Gaussian p.d.f
  RooRealVar x("x","x",-10,10) ;

  RooRealVar m("m","m",0,-10,10) ;
  RooRealVar s("s","s",2,0.1,10) ;
  RooGaussian gauss("gauss","gauss(x,m,s)",x,m,s) ;

  // Construct a flat p.d.f (polynomial of 0th order)
  RooPolynomial poly("poly","poly(x)",x) ;

  RooRealVar ngauss ("ngauss", "ngauss", 100, 0, 200);
  RooRealVar npoly  ("npoly",  "npoly", 100, 0, 200);

  RooExtendPdf egauss ("egauss", "egauss", gauss, ngauss);
  RooExtendPdf epoly  ("epoly",  "epoly", poly, npoly);

  RooAddPdf model("model","model",RooArgSet(egauss,epoly)) ;

  // Generate small dataset for use in fitting below
  RooDataSet* d = model.generate(x,50) ;

  // Fit model (without use of constraint term)
  RooFitResult* r1 = model.fitTo(*d,Save(), Extended()) ;

  RooAbsReal* nll = model.createNLL(*d,NumCPU(2)) ;
  RooPlot* frame1 = ngauss.frame(Bins(10),Title("LL and profileLL in frac")) ;
  nll->plotOn(frame1,ShiftToZero()) ;

  // Multiply constraint with p.d.f
  RooGaussian fconstraint("fconstraint","fconstraint",ngauss,RooConst(ngauss.getVal()),RooConst(0.01)) ;

  RooProdPdf modelc("modelc","model with constraint",RooArgSet(model,fconstraint)) ;

  // Fit modelc with constraint term on parameter f
  RooFitResult* r2 = modelc.fitTo(*d,Constrain(ngauss),Save(), Extended()) ;
  RooAbsReal* nll2 = modelc.createNLL(*d,NumCPU(2),Constrain(ngauss)) ;
  RooPlot* frame2 = ngauss.frame(Bins(10),Title("LL and profileLL in frac")) ;
  nll2->plotOn(frame2,ShiftToZero(),LineColor(kRed)) ;

  TCanvas *c2 = new TCanvas("c2", "c2", 900, 900); c2->cd(); 
  frame1->Draw();
  frame2->Draw("same");

 
}