/*****************************************************************************
 * 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
// Class RooAcceptReject is a generic toy monte carlo generator implement
// the accept/reject sampling technique on any positively valued function.
// The RooAcceptReject generator is used by the various generator context
// classes to take care of generation of observables for which p.d.fs
// do not define internal methods
// END_HTML
//


#include "RooFit.h"
#include "Riostream.h"

#include "RooAcceptReject.h"
#include "RooAcceptReject.h"
#include "RooAbsReal.h"
#include "RooCategory.h"
#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooRandom.h"
#include "RooErrorHandler.h"

#include "TString.h"
#include "TIterator.h"
#include "RooMsgService.h"
#include "TClass.h"
#include "TFoam.h"
#include "RooRealBinding.h"
#include "RooNumGenFactory.h"
#include "RooNumGenConfig.h"

#include <assert.h>

using namespace std;

ClassImp(RooAcceptReject)
  ;


//_____________________________________________________________________________
void RooAcceptReject::registerSampler(RooNumGenFactory& fact)
{
  // Register RooIntegrator1D, is parameters and capabilities with RooNumIntFactory 
  RooRealVar nTrial0D("nTrial0D","Number of trial samples for cat-only generation",100,0,1e9) ;
  RooRealVar nTrial1D("nTrial1D","Number of trial samples for 1-dim generation",1000,0,1e9) ;
  RooRealVar nTrial2D("nTrial2D","Number of trial samples for 2-dim generation",100000,0,1e9) ;
  RooRealVar nTrial3D("nTrial3D","Number of trial samples for N-dim generation",10000000,0,1e9) ;

  RooAcceptReject* proto = new RooAcceptReject ;
  fact.storeProtoSampler(proto,RooArgSet(nTrial0D,nTrial1D,nTrial2D,nTrial3D)) ;
}



//_____________________________________________________________________________
RooAcceptReject::RooAcceptReject(const RooAbsReal &func, const RooArgSet &genVars, const RooNumGenConfig& config, Bool_t verbose, const RooAbsReal* maxFuncVal) :
  RooAbsNumGenerator(func,genVars,verbose,maxFuncVal), _nextCatVar(0), _nextRealVar(0)
{
  // Initialize an accept-reject generator for the specified distribution function,
  // which must be non-negative but does not need to be normalized over the
  // variables to be generated, genVars. The function and its dependents are
  // cloned and so will not be disturbed during the generation process.

  _minTrialsArray[0] = static_cast<Int_t>(config.getConfigSection("RooAcceptReject").getRealValue("nTrial0D")) ;
  _minTrialsArray[1] = static_cast<Int_t>(config.getConfigSection("RooAcceptReject").getRealValue("nTrial1D")) ;
  _minTrialsArray[2] = static_cast<Int_t>(config.getConfigSection("RooAcceptReject").getRealValue("nTrial2D")) ;
  _minTrialsArray[3] = static_cast<Int_t>(config.getConfigSection("RooAcceptReject").getRealValue("nTrial3D")) ;

  _realSampleDim = _realVars.getSize() ;
  TIterator* iter = _catVars.createIterator() ;
  RooAbsCategory* cat ;
  _catSampleMult = 1 ;
  while((cat=(RooAbsCategory*)iter->Next())) {
    _catSampleMult *=  cat->numTypes() ;
  }
  delete iter ;


  // calculate the minimum number of trials needed to estimate our integral and max value
  if (!_funcMaxVal) {

    if(_realSampleDim > 3) {
      _minTrials= _minTrialsArray[3]*_catSampleMult;
      coutW(Generation) << fName << "::" << ClassName() << ": WARNING: generating " << _realSampleDim
			<< " variables with accept-reject may not be accurate" << endl;
    }
    else {
      _minTrials= _minTrialsArray[_realSampleDim]*_catSampleMult;
    }
  } else {
    // No trials needed if we know the maximum a priori
    _minTrials=0 ;
  }

  // Need to fix some things here
  if (_realSampleDim > 1) {
    coutW(Generation) << "RooAcceptReject::ctor(" << fName 
		      << ") WARNING: performing accept/reject sampling on a p.d.f in " << _realSampleDim << " dimensions "
		      << "without prior knowledge on maximum value of p.d.f. Determining maximum value by taking " << _minTrials 
		      << " trial samples. If p.d.f contains sharp peaks smaller than average distance between trial sampling points"
		      << " these may be missed and p.d.f. may be sampled incorrectly." << endl ;
  }
  if (_minTrials>10000) {
    coutW(Generation) << "RooAcceptReject::ctor(" << fName << "): WARNING: " << _minTrials << " trial samples requested by p.d.f for " 
		      << _realSampleDim << "-dimensional accept/reject sampling, this may take some time" << endl ;
  }

  // print a verbose summary of our configuration, if requested
  if(_verbose) {
    coutI(Generation) << fName << "::" << ClassName() << ":" << endl
		      << "  Initializing accept-reject generator for" << endl << "    ";
    _funcClone->printStream(ccoutI(Generation),kName,kSingleLine);
    if (_funcMaxVal) {
      ccoutI(Generation) << "  Function maximum provided, no trial sampling performed" << endl ;
    } else {
      ccoutI(Generation) << "  Real sampling dimension is " << _realSampleDim << endl;
      ccoutI(Generation) << "  Category sampling multiplier is " << _catSampleMult << endl ;
      ccoutI(Generation) << "  Min sampling trials is " << _minTrials << endl;
    }
    if (_catVars.getSize()>0) {
      ccoutI(Generation) << "  Will generate category vars "<< _catVars << endl ;
    }
    if (_realVars.getSize()>0) {
      ccoutI(Generation) << "  Will generate real vars " << _realVars << endl ;
    }
  }
  // create iterators for the new sets
  _nextCatVar= _catVars.createIterator();
  _nextRealVar= _realVars.createIterator();
  assert(0 != _nextCatVar && 0 != _nextRealVar);

  // initialize our statistics
  _maxFuncVal= 0;
  _funcSum= 0;
  _totalEvents= 0;
  _eventsUsed= 0;
}



//_____________________________________________________________________________
RooAcceptReject::~RooAcceptReject() 
{
  // Destructor
  delete _nextCatVar;
  delete _nextRealVar;
}



//_____________________________________________________________________________
const RooArgSet *RooAcceptReject::generateEvent(UInt_t remaining, Double_t& resampleRatio) 
{
  // Return a pointer to a generated event. The caller does not own the event and it
  // will be overwritten by a subsequent call. The input parameter 'remaining' should
  // contain your best guess at the total number of subsequent events you will request.

  // are we actually generating anything? (the cache always contains at least our function value)
  const RooArgSet *event= _cache->get();
  if(event->getSize() == 1) return event;

  if (!_funcMaxVal) {
    // Generation with empirical maximum determination

    // first generate enough events to get reasonable estimates for the integral and
    // maximum function value

    while(_totalEvents < _minTrials) {
      addEventToCache();

      // Limit cache size to 1M events
      if (_cache->numEntries()>1000000) {
	coutI(Generation) << "RooAcceptReject::generateEvent: resetting event cache" << endl ;
	_cache->reset() ;
	_eventsUsed = 0 ;
      }
    }
    
    event= 0;
    Double_t oldMax2(_maxFuncVal);
    while(0 == event) {
      // Use any cached events first
      if (_maxFuncVal>oldMax2) {
	cxcoutD(Generation) << "RooAcceptReject::generateEvent maxFuncVal has changed, need to resample already accepted events by factor" 
			    << oldMax2 << "/" << _maxFuncVal << "=" << oldMax2/_maxFuncVal << endl ;	
	resampleRatio=oldMax2/_maxFuncVal ;
      }
      event= nextAcceptedEvent();
      if(event) break;
      // When we have used up the cache, start a new cache and add
      // some more events to it.      
      _cache->reset();
      _eventsUsed= 0;
      // Calculate how many more events to generate using our best estimate of our efficiency.
      // Always generate at least one more event so we don't get stuck.
      if(_totalEvents*_maxFuncVal <= 0) {
	coutE(Generation) << "RooAcceptReject::generateEvent: cannot estimate efficiency...giving up" << endl;
	return 0;
      }

      Double_t eff= _funcSum/(_totalEvents*_maxFuncVal);
      Long64_t extra= 1 + (Long64_t)(1.05*remaining/eff);
      cxcoutD(Generation) << "RooAcceptReject::generateEvent: adding " << extra << " events to the cache, eff = " << eff << endl;
      Double_t oldMax(_maxFuncVal);
      while(extra--) {
	addEventToCache();
	if((_maxFuncVal > oldMax)) {
	  cxcoutD(Generation) << "RooAcceptReject::generateEvent: estimated function maximum increased from "
			      << oldMax << " to " << _maxFuncVal << endl;
	  oldMax = _maxFuncVal ;
	  // Trim cache here
	}
      }
    }

    // Limit cache size to 1M events
    if (_eventsUsed>1000000) {
      _cache->reset() ;
      _eventsUsed = 0 ;
    }

  } else {
    // Generation with a priori maximum knowledge
    _maxFuncVal = _funcMaxVal->getVal() ;
    
    // Generate enough trials to produce a single accepted event
    event = 0 ;
    while(0==event) {
      addEventToCache() ;
      event = nextAcceptedEvent() ;
    }

  }
  return event;
}



//_____________________________________________________________________________
const RooArgSet *RooAcceptReject::nextAcceptedEvent() 
{
  // Scan through events in the cache which have not been used yet,
  // looking for the first accepted one which is added to the specified
  // container. Return a pointer to the accepted event, or else zero
  // if we use up the cache before we accept an event. The caller does
  // not own the event and it will be overwritten by a subsequent call.

  const RooArgSet *event = 0;
  while((event= _cache->get(_eventsUsed))) {    
    _eventsUsed++ ;
    // accept this cached event?
    Double_t r= RooRandom::uniform();
    if(r*_maxFuncVal > _funcValPtr->getVal()) {
      //cout << " event number " << _eventsUsed << " has been rejected" << endl ;
      continue;
    }
    //cout << " event number " << _eventsUsed << " has been accepted" << endl ;
    // copy this event into the output container
    if(_verbose && (_eventsUsed%1000==0)) {
      cerr << "RooAcceptReject: accepted event (used " << _eventsUsed << " of "
	   << _cache->numEntries() << " so far)" << endl;
    }
    break;
  }  
  //cout << "accepted event " << _eventsUsed << " of " << _cache->numEntries() << endl ;
  return event;
}



//_____________________________________________________________________________
void RooAcceptReject::addEventToCache() 
{
  // Add a trial event to our cache and update our estimates
  // of the function maximum value and integral.

  // randomize each discrete argument
  _nextCatVar->Reset();
  RooCategory *cat = 0;
  while((cat= (RooCategory*)_nextCatVar->Next())) cat->randomize();

  // randomize each real argument
  _nextRealVar->Reset();
  RooRealVar *real = 0;
  while((real= (RooRealVar*)_nextRealVar->Next())) real->randomize();

  // calculate and store our function value at this new point
  Double_t val= _funcClone->getVal();
  _funcValPtr->setVal(val);

  // Update the estimated integral and maximum value. Increase our
  // maximum estimate slightly to give a safety margin with a
  // corresponding loss of efficiency.
  if(val > _maxFuncVal) _maxFuncVal= 1.05*val;
  _funcSum+= val;

  // fill a new entry in our cache dataset for this point
  _cache->fill();
  _totalEvents++;

  if (_verbose &&_totalEvents%10000==0) {
    cerr << "RooAcceptReject: generated " << _totalEvents << " events so far." << endl ;
  }

}

Double_t RooAcceptReject::getFuncMax() 
{
  // Empirically determine maximum value of function by taking a large number
  // of samples. The actual number depends on the number of dimensions in which
  // the sampling occurs

  // Generate the minimum required number of samples for a reliable maximum estimate
  while(_totalEvents < _minTrials) {
    addEventToCache();

    // Limit cache size to 1M events
    if (_cache->numEntries()>1000000) {
      coutI(Generation) << "RooAcceptReject::getFuncMax: resetting event cache" << endl ;
      _cache->reset() ;
      _eventsUsed = 0 ;
    }
  }  

  return _maxFuncVal ;
}