/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitCore                                                       *
 * @(#)root/roofitcore:$Id$
 * 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
// RooRandomizeParamMCSModule is an add-on modules to RooMCStudy that
// allows you to randomize input generation parameters. Randomized generation
// parameters can be sampled from a uniform or Gaussian distribution.
// For every randomized parameter, an extra variable is added to 
// RooMCStudy::fitParDataSet() named <tt><parname>_gen</tt> that indicates the actual
// value used for generation for each trial. 
// <p>
// You can also choose to randomize the sum of N parameters, rather
// than a single parameter. In that case common multiplicative scale
// factor is applied to each component to bring the sum to the desired
// target value taken from either uniform or Gaussian sampling. This
// latter option is for example useful if you want to change the total
// number of expected events of an extended p.d.f
// END_HTML
//


#include "Riostream.h"
#include "RooDataSet.h"
#include "RooRealVar.h"
#include "RooRandom.h"
#include "TString.h"
#include "RooFit.h"
#include "RooFitResult.h"
#include "RooAddition.h"
#include "RooMsgService.h"
#include "RooRandomizeParamMCSModule.h"

using namespace std ;

ClassImp(RooRandomizeParamMCSModule)
  ;



////////////////////////////////////////////////////////////////////////////////
/// Constructor

RooRandomizeParamMCSModule::RooRandomizeParamMCSModule() : 
  RooAbsMCStudyModule("RooRandomizeParamMCSModule","RooRandomizeParamMCSModule"), _data(0)
{
}



////////////////////////////////////////////////////////////////////////////////
/// Copy constructor

RooRandomizeParamMCSModule::RooRandomizeParamMCSModule(const RooRandomizeParamMCSModule& other) : 
  RooAbsMCStudyModule(other), 
  _unifParams(other._unifParams),
  _gausParams(other._gausParams), 
  _data(0)
{
}



////////////////////////////////////////////////////////////////////////////////
/// Destructor

RooRandomizeParamMCSModule:: ~RooRandomizeParamMCSModule() 
{
  if (_data) {
    delete _data ;
  }
}



////////////////////////////////////////////////////////////////////////////////
/// Request uniform smearing of param in range [lo,hi] in RooMCStudy
/// generation cycle

void RooRandomizeParamMCSModule::sampleUniform(RooRealVar& param, Double_t lo, Double_t hi) 
{  
  // If we're already attached to a RooMCStudy, check that given param is actual generator model parameter
  // If not attached, this check is repeated at the attachment moment
  if (genParams()) {
    RooRealVar* actualPar = static_cast<RooRealVar*>(genParams()->find(param.GetName())) ;
    if (!actualPar) {
      oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::initializeInstance: variable " << param.GetName() << " is not a parameter of RooMCStudy model and is ignored!" << endl ;
      return ;
    }
  }

  _unifParams.push_back(UniParam(&param,lo,hi)) ;
}



////////////////////////////////////////////////////////////////////////////////
/// Request Gaussian smearing of param in with mean 'mean' and width
/// 'sigma' in RooMCStudy generation cycle

void RooRandomizeParamMCSModule::sampleGaussian(RooRealVar& param, Double_t mean, Double_t sigma) 
{
  // If we're already attached to a RooMCStudy, check that given param is actual generator model parameter
  // If not attached, this check is repeated at the attachment moment
  if (genParams()) {
    RooRealVar* actualPar = static_cast<RooRealVar*>(genParams()->find(param.GetName())) ;
    if (!actualPar) {
      oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::initializeInstance: variable " << param.GetName() << " is not a parameter of RooMCStudy model and is ignored!" << endl ;
      return ;
    }
  }

  _gausParams.push_back(GausParam(&param,mean,sigma)) ;
}




////////////////////////////////////////////////////////////////////////////////
/// Request uniform smearing of sum of parameters in paramSet uniform
/// smearing in range [lo,hi] in RooMCStudy generation cycle.  This
/// option applies a common multiplicative factor to each parameter
/// in paramSet to make the sum of the parameters add up to the
/// sampled value in the range [lo,hi]

void RooRandomizeParamMCSModule::sampleSumUniform(const RooArgSet& paramSet, Double_t lo, Double_t hi) 
{
  // Check that all args are RooRealVars
  RooArgSet okset ;
  TIterator* iter = paramSet.createIterator() ;
  RooAbsArg* arg ;
  while((arg=(RooAbsArg*)iter->Next())) {
    // Check that arg is a RooRealVar
    RooRealVar* rrv = dynamic_cast<RooRealVar*>(arg) ;
    if (!rrv) {
      oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::sampleSumUniform() ERROR: input parameter " << arg->GetName() << " is not a RooRealVar and is ignored" << endl ;
      continue;
    }
    okset.add(*rrv) ;
  }
  delete iter ;

  // If we're already attached to a RooMCStudy, check that given param is actual generator model parameter
  // If not attached, this check is repeated at the attachment moment
  RooArgSet okset2 ;
  if (genParams()) {
    TIterator* psiter = okset.createIterator() ;
    RooAbsArg* arg2 ;
    while ((arg2=(RooAbsArg*)psiter->Next())) {
      RooRealVar* actualVar= static_cast<RooRealVar*>(genParams()->find(arg2->GetName())) ;
      if (!actualVar) {
	oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::sampleSumUniform: variable " << arg2->GetName() << " is not a parameter of RooMCStudy model and is ignored!" << endl ;	
      } else {
	okset2.add(*actualVar) ;
      }
    }    
    delete psiter ;
  } else {

   // If genParams() are not available, skip this check for now
   okset2.add(okset) ;

  }


  _unifParamSets.push_back(UniParamSet(okset2,lo,hi)) ;

}




////////////////////////////////////////////////////////////////////////////////
/// Request gaussian smearing of sum of parameters in paramSet
/// uniform smearing with mean 'mean' and width 'sigma' in RooMCStudy
/// generation cycle.  This option applies a common multiplicative
/// factor to each parameter in paramSet to make the sum of the
/// parameters add up to the sampled value from the
/// gaussian(mean,sigma)

void RooRandomizeParamMCSModule::sampleSumGauss(const RooArgSet& paramSet, Double_t mean, Double_t sigma) 
{
  // Check that all args are RooRealVars
  RooArgSet okset ;
  TIterator* iter = paramSet.createIterator() ;
  RooAbsArg* arg ;
  while((arg=(RooAbsArg*)iter->Next())) {
    // Check that arg is a RooRealVar
    RooRealVar* rrv = dynamic_cast<RooRealVar*>(arg) ;
    if (!rrv) {
      oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::sampleSumGauss() ERROR: input parameter " << arg->GetName() << " is not a RooRealVar and is ignored" << endl ;
      continue;
    }
    okset.add(*rrv) ;
  }

  // If we're already attached to a RooMCStudy, check that given param is actual generator model parameter
  // If not attached, this check is repeated at the attachment moment
  RooArgSet okset2 ;
  if (genParams()) {
    TIterator* psiter = okset.createIterator() ;
    RooAbsArg* arg2 ;
    while ((arg2=(RooAbsArg*)psiter->Next())) {
      RooRealVar* actualVar= static_cast<RooRealVar*>(genParams()->find(arg2->GetName())) ;
      if (!actualVar) {
	oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::sampleSumUniform: variable " << arg2->GetName() << " is not a parameter of RooMCStudy model and is ignored!" << endl ;	
      } else {
	okset2.add(*actualVar) ;
      }
    }    
    delete psiter ;
  } else {

   // If genParams() are not available, skip this check for now
   okset2.add(okset) ;

  }

  _gausParamSets.push_back(GausParamSet(okset,mean,sigma)) ;
  
}




////////////////////////////////////////////////////////////////////////////////
/// Initialize module after attachment to RooMCStudy object

Bool_t RooRandomizeParamMCSModule::initializeInstance()
{
  // Loop over all uniform smearing parameters
  std::list<UniParam>::iterator uiter ;
  for (uiter= _unifParams.begin() ; uiter!= _unifParams.end() ; ++uiter) {

    // Check that listed variable is actual generator model parameter
    RooRealVar* actualPar = static_cast<RooRealVar*>(genParams()->find(uiter->_param->GetName())) ;
    if (!actualPar) {
      oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::initializeInstance: variable " << uiter->_param->GetName() << " is not a parameter of RooMCStudy model and is ignored!" << endl ;
      uiter = _unifParams.erase(uiter) ;
      continue ;
    }
    uiter->_param = actualPar ;

    // Add variable to summary dataset to hold generator value
    TString parName = Form("%s_gen",uiter->_param->GetName()) ;
    TString parTitle = Form("%s as generated",uiter->_param->GetTitle()) ;
    RooRealVar* par_gen = new RooRealVar(parName.Data(),parTitle.Data(),0) ;    
    _genParSet.addOwned(*par_gen) ;
  }
  
  // Loop over all gaussian smearing parameters
  std::list<UniParam>::iterator giter ;
  for (giter= _unifParams.begin() ; giter!= _unifParams.end() ; ++giter) {

    // Check that listed variable is actual generator model parameter
    RooRealVar* actualPar = static_cast<RooRealVar*>(genParams()->find(giter->_param->GetName())) ;
    if (!actualPar) {
      oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::initializeInstance: variable " << giter->_param->GetName() << " is not a parameter of RooMCStudy model and is ignored!" << endl ;
      giter = _unifParams.erase(giter) ;
      continue ;
    }
    giter->_param = actualPar ;

    // Add variable to summary dataset to hold generator value
    TString parName = Form("%s_gen",giter->_param->GetName()) ;
    TString parTitle = Form("%s as generated",giter->_param->GetTitle()) ;
    RooRealVar* par_gen = new RooRealVar(parName.Data(),parTitle.Data(),0) ;    
    _genParSet.addOwned(*par_gen) ;
  }


  // Loop over all uniform smearing set of parameters
  std::list<UniParamSet>::iterator usiter ;
  for (usiter= _unifParamSets.begin() ; usiter!= _unifParamSets.end() ; ++usiter) {
    
    // Check that all listed variables are actual generator model parameters
    RooArgSet actualPSet ;
    TIterator* psiter = usiter->_pset.createIterator() ;
    RooAbsArg* arg ;
    while ((arg=(RooAbsArg*)psiter->Next())) {
      RooRealVar* actualVar= static_cast<RooRealVar*>(genParams()->find(arg->GetName())) ;
      if (!actualVar) {
	oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::initializeInstance: variable " << arg->GetName() << " is not a parameter of RooMCStudy model and is ignored!" << endl ;	
      } else {
	actualPSet.add(*actualVar) ;
      }
    }    
    delete psiter ;
    usiter->_pset.removeAll() ;
    usiter->_pset.add(actualPSet) ;

    // Add variables to summary dataset to hold generator values
    TIterator* iter = usiter->_pset.createIterator() ;
    RooRealVar* param ;
    while((param=(RooRealVar*)iter->Next())) {
      TString parName = Form("%s_gen",param->GetName()) ;
      TString parTitle = Form("%s as generated",param->GetTitle()) ;
      RooRealVar* par_gen = new RooRealVar(parName.Data(),parTitle.Data(),0) ;    
      _genParSet.addOwned(*par_gen) ;          
    }
    delete iter ;
  }
  
  // Loop over all gaussian smearing set of parameters
  std::list<GausParamSet>::iterator ugiter ;
  for (ugiter= _gausParamSets.begin() ; ugiter!= _gausParamSets.end() ; ++ugiter) {

    // Check that all listed variables are actual generator model parameters
    RooArgSet actualPSet ;
    TIterator* psiter = ugiter->_pset.createIterator() ;
    RooAbsArg* arg ;
    while ((arg=(RooAbsArg*)psiter->Next())) {
      RooRealVar* actualVar= static_cast<RooRealVar*>(genParams()->find(arg->GetName())) ;
      if (!actualVar) {
	oocoutW((TObject*)0,InputArguments) << "RooRandomizeParamMCSModule::initializeInstance: variable " << arg->GetName() << " is not a parameter of RooMCStudy model and is ignored!" << endl ;	
      } else {
	actualPSet.add(*actualVar) ;
      }
    }    
    ugiter->_pset.removeAll() ;
    ugiter->_pset.add(actualPSet) ;

    // Add variables to summary dataset to hold generator values
    TIterator* iter = ugiter->_pset.createIterator() ;
    RooRealVar* param ;
    while((param=(RooRealVar*)iter->Next())) {
      TString parName = Form("%s_gen",param->GetName()) ;
      TString parTitle = Form("%s as generated",param->GetTitle()) ;
      RooRealVar* par_gen = new RooRealVar(parName.Data(),parTitle.Data(),0) ;    
      _genParSet.addOwned(*par_gen) ;          
    }
  }
  
  // Create new dataset to be merged with RooMCStudy::fitParDataSet
  _data = new RooDataSet("DeltaLLSigData","Additional data for Delta(-log(L)) study",_genParSet) ;

  return kTRUE ;
}



////////////////////////////////////////////////////////////////////////////////
/// Initialize module at beginning of RooCMStudy run

Bool_t RooRandomizeParamMCSModule::initializeRun(Int_t /*numSamples*/) 
{
  // Clear dataset at beginning of run
  _data->reset() ;
  return kTRUE ;
}



////////////////////////////////////////////////////////////////////////////////
/// Apply all smearings to generator parameters

Bool_t RooRandomizeParamMCSModule::processBeforeGen(Int_t /*sampleNum*/) 
{
  // Apply uniform smearing to all generator parameters for which it is requested
  std::list<UniParam>::iterator uiter ;
  for (uiter= _unifParams.begin() ; uiter!= _unifParams.end() ; ++uiter) {
    Double_t newVal = RooRandom::randomGenerator()->Uniform(uiter->_lo,uiter->_hi) ;        
    oocoutE((TObject*)0,Generation) << "RooRandomizeParamMCSModule::processBeforeGen: applying uniform smearing to generator parameter " 
	 << uiter->_param->GetName() << " in range [" << uiter->_lo << "," << uiter->_hi << "], chosen value for this sample is " << newVal << endl ;
    uiter->_param->setVal(newVal) ;

    RooRealVar* genpar = static_cast<RooRealVar*>(_genParSet.find(Form("%s_gen",uiter->_param->GetName()))) ;
    genpar->setVal(newVal) ;
  }

  // Apply gaussian smearing to all generator parameters for which it is requested
  std::list<GausParam>::iterator giter ;
  for (giter= _gausParams.begin() ; giter!= _gausParams.end() ; ++giter) {
    Double_t newVal = RooRandom::randomGenerator()->Gaus(giter->_mean,giter->_sigma) ;
    oocoutI((TObject*)0,Generation) << "RooRandomizeParamMCSModule::processBeforeGen: applying gaussian smearing to generator parameter " 
	 << giter->_param->GetName() << " with a mean of " << giter->_mean << " and a width of " << giter->_sigma << ", chosen value for this sample is " << newVal << endl ;
    giter->_param->setVal(newVal) ;

    RooRealVar* genpar = static_cast<RooRealVar*>(_genParSet.find(Form("%s_gen",giter->_param->GetName()))) ;
    genpar->setVal(newVal) ;
  }

  // Apply uniform smearing to all sets of generator parameters for which it is requested
  std::list<UniParamSet>::iterator usiter ;
  for (usiter= _unifParamSets.begin() ; usiter!= _unifParamSets.end() ; ++usiter) {

    // Calculate new value for sum 
    Double_t newVal = RooRandom::randomGenerator()->Uniform(usiter->_lo,usiter->_hi) ;        
    oocoutI((TObject*)0,Generation) << "RooRandomizeParamMCSModule::processBeforeGen: applying uniform smearing to sum of set of generator parameters " 
				    <<  usiter->_pset
				    << " in range [" << usiter->_lo << "," << usiter->_hi << "], chosen sum value for this sample is " << newVal << endl ;

    // Determine original value of sum and calculate per-component scale factor to obtain new valye for sum
    RooAddition sumVal("sumVal","sumVal",usiter->_pset) ;
    Double_t compScaleFactor = newVal/sumVal.getVal() ;

    // Apply multiplicative correction to each term of the sum
    TIterator* iter = usiter->_pset.createIterator() ;
    RooRealVar* param ;
    while((param=(RooRealVar*)iter->Next())) {
      param->setVal(param->getVal()*compScaleFactor) ;
      RooRealVar* genpar = static_cast<RooRealVar*>(_genParSet.find(Form("%s_gen",param->GetName()))) ;
      genpar->setVal(param->getVal()) ;
    }
    delete iter ;
  }

  // Apply gaussian smearing to all sets of generator parameters for which it is requested
  std::list<GausParamSet>::iterator gsiter ;
  for (gsiter= _gausParamSets.begin() ; gsiter!= _gausParamSets.end() ; ++gsiter) {
    
    // Calculate new value for sum 
    Double_t newVal = RooRandom::randomGenerator()->Gaus(gsiter->_mean,gsiter->_sigma) ;        
    oocoutI((TObject*)0,Generation) << "RooRandomizeParamMCSModule::processBeforeGen: applying gaussian smearing to sum of set of generator parameters " 
				    << gsiter->_pset
				    << " with a mean of " << gsiter->_mean << " and a width of " << gsiter->_sigma 
				    << ", chosen value for this sample is " << newVal << endl ;

    // Determine original value of sum and calculate per-component scale factor to obtain new valye for sum
    RooAddition sumVal("sumVal","sumVal",gsiter->_pset) ;
    Double_t compScaleFactor = newVal/sumVal.getVal() ;

    // Apply multiplicative correction to each term of the sum
    TIterator* iter = gsiter->_pset.createIterator() ;
    RooRealVar* param ;
    while((param=(RooRealVar*)iter->Next())) {
      param->setVal(param->getVal()*compScaleFactor) ;
      RooRealVar* genpar = static_cast<RooRealVar*>(_genParSet.find(Form("%s_gen",param->GetName()))) ;
      genpar->setVal(param->getVal()) ;
    }
  }
  
  // Store generator values for all modified parameters
  _data->add(_genParSet) ;
  
  return kTRUE ;
}



////////////////////////////////////////////////////////////////////////////////
/// Return auxiliary data of this module so that it is merged with
/// RooMCStudy::fitParDataSet()

RooDataSet* RooRandomizeParamMCSModule::finalizeRun() 
{
  return _data ;
}