/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitCore                                                       *
 * @(#)root/roofitcore:$Id: RooChi2MCSModule.cxx 28963 2009-06-12 15:47:45Z wouter $
 * Authors:                                                                  *
 *   WV, Wouter Verkerke, UC Santa Barbara, verkerke@slac.stanford.edu       *
 *   DK, David Kirkby,    UC Irvine,         dkirkby@uci.edu                 *
 *                                                                           *
 * Copyright (c) 2000-2005, Regents of the University of California          *
 *                          and Stanford University. All rights reserved.    *
 *                                                                           *
 * Redistribution and use in source and binary forms,                        *
 * with or without modification, are permitted according to the terms        *
 * listed in LICENSE (http://roofit.sourceforge.net/license.txt)             *
 *****************************************************************************/

//////////////////////////////////////////////////////////////////////////////
// 
// BEGIN_HTML
// RooChi2MCSModule is an add-on modules to RooMCStudy that
// calculates the chi-squared of fitted p.d.f with respect to a binned
// version of the data. For each fit the chi-squared, the reduced chi-squared
// the number of degrees of freedom and the probability of the chi-squared
// is store in the summary dataset
// END_HTML
//
//

#include "Riostream.h"

#include "RooDataSet.h"
#include "RooRealVar.h"
#include "TString.h"
#include "RooFit.h"
#include "RooFitResult.h"
#include "RooChi2MCSModule.h"
#include "RooMsgService.h"
#include "RooChi2Var.h"
#include "RooDataHist.h"
#include "TMath.h"
#include "RooGlobalFunc.h"



ClassImp(RooChi2MCSModule)
  ;



//_____________________________________________________________________________
RooChi2MCSModule::RooChi2MCSModule() : 
  RooAbsMCStudyModule("RooChi2MCSModule","RooChi2Module"),
  _data(0), _chi2(0), _ndof(0), _chi2red(0), _prob(0)

{
  // Constructor of module 
}




//_____________________________________________________________________________
RooChi2MCSModule::RooChi2MCSModule(const RooChi2MCSModule& other) : 
  RooAbsMCStudyModule(other), 
  _data(0), _chi2(0), _ndof(0), _chi2red(0), _prob(0)
{
  // Copy constructor
}



//_____________________________________________________________________________
RooChi2MCSModule:: ~RooChi2MCSModule() 
{
  // Destructor

  if (_chi2) {
    delete _chi2 ;
  }
  if (_ndof) {
    delete _ndof ;
  }
  if (_chi2red) {
    delete _chi2red ;
  }
  if (_prob) {
    delete _prob ;
  }
  if (_data) {
    delete _data ;
  }
}



//_____________________________________________________________________________
Bool_t RooChi2MCSModule::initializeInstance()
{
  // Initialize module after attachment to RooMCStudy object

  // Construct variable that holds -log(L) fit with null hypothesis for given parameter
  _chi2     = new RooRealVar("chi2","chi^2",0) ;
  _ndof     = new RooRealVar("ndof","number of degrees of freedom",0) ;   
  _chi2red  = new RooRealVar("chi2red","reduced chi^2",0) ; 
  _prob     = new RooRealVar("prob","prob(chi2,ndof)",0) ;

  // Create new dataset to be merged with RooMCStudy::fitParDataSet
  _data = new RooDataSet("Chi2Data","Additional data for Chi2 study",RooArgSet(*_chi2,*_ndof,*_chi2red,*_prob)) ;

  return kTRUE ;
}



//_____________________________________________________________________________
Bool_t RooChi2MCSModule::initializeRun(Int_t /*numSamples*/) 
{
  // Initialize module at beginning of RooCMStudy run

  _data->reset() ;
  return kTRUE ;
}



//_____________________________________________________________________________
RooDataSet* RooChi2MCSModule::finalizeRun() 
{
  // Return auxiliary dataset with results of chi2 analysis
  // calculations of this module so that it is merged with
  // RooMCStudy::fitParDataSet() by RooMCStudy

  return _data ;
}



//_____________________________________________________________________________
Bool_t RooChi2MCSModule::processAfterFit(Int_t /*sampleNum*/)  
{
  // Bin dataset and calculate chi2 of p.d.f w.r.t binned dataset

  RooAbsData* data = genSample() ;
  RooDataHist* binnedData = dynamic_cast<RooDataHist*>(data) ;
  Bool_t deleteData(kFALSE) ;
  if (!binnedData) {
    deleteData = kTRUE ;
    binnedData = ((RooDataSet*)data)->binnedClone() ;
  }

  RooChi2Var chi2Var("chi2Var","chi2Var",*fitModel(),*binnedData,RooFit::Extended(extendedGen()),RooFit::DataError(RooAbsData::SumW2)) ;

  RooArgSet* floatPars = (RooArgSet*) fitParams()->selectByAttrib("Constant",kFALSE) ;  

  _chi2->setVal(chi2Var.getVal()) ;
  _ndof->setVal(binnedData->numEntries()-floatPars->getSize()-1) ; 
  _chi2red->setVal(_chi2->getVal()/_ndof->getVal()) ;
  _prob->setVal(TMath::Prob(_chi2->getVal(),static_cast<int>(_ndof->getVal()))) ;

  _data->add(RooArgSet(*_chi2,*_ndof,*_chi2red,*_prob)) ;

  if (deleteData) {
    delete binnedData ;
  }
  delete floatPars ;

  return kTRUE ;
}