/*****************************************************************************
 * 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 RooGenContext implement a universal generator context for all
// RooAbsPdf classes that do not have or need a specialized generator
// context. This generator context queries the input p.d.f which observables
// it can generate internally and delegates generation of those observables
// to the p.d.f if it deems that safe. The other observables are generated
// use a RooAcceptReject sampling technique.
// END_HTML
//


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

#include "RooGenContext.h"
#include "RooGenContext.h"
#include "RooAbsPdf.h"
#include "RooDataSet.h"
#include "RooRealIntegral.h"
#include "RooAcceptReject.h"
#include "RooRealVar.h"
#include "RooDataHist.h"
#include "RooErrorHandler.h"
#include "RooNumGenConfig.h"
#include "RooNumGenFactory.h"

#include "TString.h"
#include "TIterator.h"
#include "TClass.h"



using namespace std;

ClassImp(RooGenContext)
  ;



//_____________________________________________________________________________
RooGenContext::RooGenContext(const RooAbsPdf &model, const RooArgSet &vars,
			     const RooDataSet *prototype, const RooArgSet* auxProto,
			     Bool_t verbose, const RooArgSet* forceDirect) :  
  RooAbsGenContext(model,vars,prototype,auxProto,verbose),
  _cloneSet(0), _pdfClone(0), _acceptRejectFunc(0), _generator(0),
  _maxVar(0), _uniIter(0), _updateFMaxPerEvent(0) 
{
  // Initialize a new context for generating events with the specified
  // variables, using the specified PDF model. A prototype dataset (if provided)
  // is not cloned and still belongs to the caller. The contents and shape
  // of this dataset can be changed between calls to generate() as long as the
  // expected columns to be copied to the generated dataset are present.
  // Any argument supplied in the forceDirect RooArgSet are always offered
  // for internal generation to the p.d.f., even if this is deemed unsafe by
  // the logic of RooGenContext.

  cxcoutI(Generation) << "RooGenContext::ctor() setting up event generator context for p.d.f. " << model.GetName() 
			<< " for generation of observable(s) " << vars ;
  if (prototype) ccxcoutI(Generation) << " with prototype data for " << *prototype->get() ;
  if (auxProto && auxProto->getSize()>0)  ccxcoutI(Generation) << " with auxiliary prototypes " << *auxProto ;
  if (forceDirect && forceDirect->getSize()>0)  ccxcoutI(Generation) << " with internal generation forced for observables " << *forceDirect ;
  ccxcoutI(Generation) << endl ;


  // Clone the model and all nodes that it depends on so that this context
  // is independent of any existing objects.
  RooArgSet nodes(model,model.GetName());
  _cloneSet= (RooArgSet*) nodes.snapshot(kTRUE);
  if (!_cloneSet) {
    coutE(Generation) << "RooGenContext::RooGenContext(" << GetName() << ") Couldn't deep-clone PDF, abort," << endl ;
    RooErrorHandler::softAbort() ;
  }

  // Find the clone in the snapshot list
  _pdfClone = (RooAbsPdf*)_cloneSet->find(model.GetName());
  _pdfClone->setOperMode(RooAbsArg::ADirty,kTRUE) ;

  // Optionally fix RooAddPdf normalizations
  if (prototype&&_pdfClone->dependsOn(*prototype->get())) {
    RooArgSet fullNormSet(vars) ;
    fullNormSet.add(*prototype->get()) ;
    _pdfClone->fixAddCoefNormalization(fullNormSet) ;
  }
      
  // Analyze the list of variables to generate...
  _isValid= kTRUE;
  TIterator *iterator= vars.createIterator();
  TIterator *servers= _pdfClone->serverIterator();
  const RooAbsArg *tmp = 0;
  const RooAbsArg *arg = 0;
  while((_isValid && (tmp= (const RooAbsArg*)iterator->Next()))) {
    // is this argument derived?
    if(!tmp->isFundamental()) {
      coutE(Generation) << "RooGenContext::ctor(): cannot generate values for derived \""  << tmp->GetName() << "\"" << endl;
      _isValid= kFALSE;
      continue;
    }
    // lookup this argument in the cloned set of PDF dependents
    arg= (const RooAbsArg*)_cloneSet->find(tmp->GetName());
    if(0 == arg) {
      coutI(Generation) << "RooGenContext::ctor() WARNING model does not depend on \"" << tmp->GetName() 
			  << "\" which will have uniform distribution" << endl;
      _uniformVars.add(*tmp);
    }
    else {

      // does the model depend on this variable directly, ie, like "x" in
      // f(x) or f(x,g(x,y)) or even f(x,x) ?      
      const RooAbsArg *direct= arg ;
      if (forceDirect==0 || !forceDirect->find(direct->GetName())) {
	if (!_pdfClone->isDirectGenSafe(*arg)) {
	  cxcoutD(Generation) << "RooGenContext::ctor() observable " << arg->GetName() << " has been determined to be unsafe for internal generation" << endl;
	  direct=0 ;
	}
      }
      
      // does the model depend indirectly on this variable through an lvalue chain?	
      // otherwise, this variable will have to be generated with accept/reject
      if(direct) { 
	_directVars.add(*arg);
      } else {
	_otherVars.add(*arg);
      }
    }
  }
  delete servers;
  delete iterator;
  if(!isValid()) {
    coutE(Generation) << "RooGenContext::ctor() constructor failed with errors" << endl;
    return;
  }

  // At this point directVars are all variables that are safe to be generated directly
  //               otherVars are all variables that are _not_ safe to be generated directly
  if (_directVars.getSize()>0) {
    cxcoutD(Generation) << "RooGenContext::ctor() observables " << _directVars << " are safe for internal generator (if supported by p.d.f)" << endl ;
  }
  if (_otherVars.getSize()>0) {
    cxcoutD(Generation) << "RooGenContext::ctor() observables " << _otherVars << " are NOT safe for internal generator (if supported by p.d.f)" << endl ;
  }

  // If PDF depends on prototype data, direct generator cannot use static initialization
  // in initGenerator()
  Bool_t staticInitOK = !_pdfClone->dependsOn(_protoVars) ;
  if (!staticInitOK) {
    cxcoutD(Generation) << "RooGenContext::ctor() Model depends on supplied protodata observables, static initialization of internal generator will not be allowed" << endl ;
  }
  
  // Can the model generate any of the direct variables itself?
  RooArgSet generatedVars;
  if (_directVars.getSize()>0) {
    _code= _pdfClone->getGenerator(_directVars,generatedVars,staticInitOK);
    
    cxcoutD(Generation) << "RooGenContext::ctor() Model indicates that it can internally generate observables " 
			  << generatedVars << " with configuration identifier " << _code << endl ;
  } else {
    _code = 0 ;
  }

  // Move variables which cannot be generated into the list to be generated with accept/reject
  _directVars.remove(generatedVars);
  _otherVars.add(_directVars);

  // Update _directVars to only include variables that will actually be directly generated
  _directVars.removeAll();
  _directVars.add(generatedVars);

  cxcoutI(Generation) << "RooGenContext::ctor() Context will" ;
  if (_directVars.getSize()>0) ccxcoutI(Generation) << " let model internally generate observables " << _directVars ;
  if (_directVars.getSize()>0 && _otherVars.getSize()>0)  ccxcoutI(Generation) << " and" ;
  if (_otherVars.getSize()>0)  ccxcoutI(Generation) << " generate variables " << _otherVars << " with accept/reject sampling" ;
  if (_uniformVars.getSize()>0) ccxcoutI(Generation) << ", non-observable variables " << _uniformVars << " will be generated with flat distribution" ;
  ccxcoutI(Generation) << endl ;

  // initialize the accept-reject generator
  RooArgSet *depList= _pdfClone->getObservables(_theEvent);
  depList->remove(_otherVars);

  TString nname(_pdfClone->GetName()) ;
  nname.Append("_AccRej") ;
  TString ntitle(_pdfClone->GetTitle()) ;
  ntitle.Append(" (Accept/Reject)") ;
  

  RooArgSet* protoDeps = model.getObservables(_protoVars) ;
  

  if (_protoVars.getSize()==0) {

    // No prototype variables
    
    if(depList->getSize()==0) {
      // All variable are generated with accept-reject
      
      // Check if PDF supports maximum finding
      Int_t maxFindCode = _pdfClone->getMaxVal(_otherVars) ;
      if (maxFindCode != 0) {
	coutI(Generation) << "RooGenContext::ctor() no prototype data provided, all observables are generated with numerically and "
			      << "model supports analytical maximum findin:, can provide analytical pdf maximum to numeric generator" << endl ;
	Double_t maxVal = _pdfClone->maxVal(maxFindCode) / _pdfClone->getNorm(_theEvent) ;
	_maxVar = new RooRealVar("funcMax","function maximum",maxVal) ;
	cxcoutD(Generation) << "RooGenContext::ctor() maximum value returned by RooAbsPdf::maxVal() is " << maxVal << endl ;
      }
    }

    if (_otherVars.getSize()>0) {
      _pdfClone->getVal(&vars) ; // WVE debug
      _acceptRejectFunc= new RooRealIntegral(nname,ntitle,*_pdfClone,*depList,&vars);
      cxcoutI(Generation) << "RooGenContext::ctor() accept/reject sampling function is " << _acceptRejectFunc->GetName() << endl ;
    } else {
      _acceptRejectFunc = 0 ;
    }

  } else {

    // Generation _with_ prototype variable
    depList->remove(_protoVars,kTRUE,kTRUE) ;
    _acceptRejectFunc= (RooRealIntegral*) _pdfClone->createIntegral(*depList,vars) ;
    cxcoutI(Generation) << "RooGenContext::ctor() accept/reject sampling function is " << _acceptRejectFunc->GetName() << endl ;
    
    // Check if PDF supports maximum finding for the entire phase space
    RooArgSet allVars(_otherVars);
    allVars.add(_directVars);
    Int_t maxFindCode = _pdfClone->getMaxVal(allVars) ;
    if (maxFindCode != 0) {
      // Special case: PDF supports max-finding in otherVars, no need to scan other+proto space for maximum
      coutI(Generation) << "RooGenContext::ctor() prototype data provided, and "
                        << "model supports analytical maximum finding in the full phase space: " 
                        << "can provide analytical pdf maximum to numeric generator" << endl ;
      _maxVar = new RooRealVar("funcMax","function maximum",_pdfClone->maxVal(maxFindCode)) ;
    } else {
      maxFindCode = _pdfClone->getMaxVal(_otherVars) ;
      if (maxFindCode != 0) {
         _updateFMaxPerEvent = maxFindCode ;
         coutI(Generation) << "RooGenContext::ctor() prototype data provided, and "
                           << "model supports analytical maximum finding in the variables which are not"
                           << " internally generated. Can provide analytical pdf maximum to numeric "
                           << "generator" << endl;
         cxcoutD(Generation) << "RooGenContext::ctor() maximum value must be reevaluated for each "
                             << "event with configuration code " << maxFindCode << endl ;
         _maxVar = new RooRealVar("funcMax","function maximum",1) ;    
      }
    }

    if (!_maxVar) {
      
      // Regular case: First find maximum in other+proto space
      RooArgSet otherAndProto(_otherVars) ;

      otherAndProto.add(*protoDeps) ;
      
      if (_otherVars.getSize()>0) {      

	cxcoutD(Generation) << "RooGenContext::ctor() prototype data provided, observables are generated numericaly no " 
			      << "analytical estimate of maximum function value provided by model, must determine maximum value through initial sampling space "
			      << "of accept/reject observables plus prototype observables: " << otherAndProto << endl ;

	// Calculate maximum in other+proto space if there are any accept/reject generated observables
	RooAbsNumGenerator* maxFinder = RooNumGenFactory::instance().createSampler(*_acceptRejectFunc,otherAndProto,RooArgSet(_protoVars),
										   *model.getGeneratorConfig(),_verbose) ;
// 	RooAcceptReject maxFinder(*_acceptRejectFunc,otherAndProto,RooNumGenConfig::defaultConfig(),_verbose) ;
	Double_t max = maxFinder->getFuncMax() ;
	_maxVar = new RooRealVar("funcMax","function maximum",max) ;

	if (max==0) {	  
	  coutE(Generation) << "RooGenContext::ctor(" << model.GetName() 
			    << ") ERROR: generating conditional p.d.f. which requires prior knowledge of function maximum, " 
			    << "but chosen numeric generator (" << maxFinder->IsA()->GetName() << ") does not support maximum finding" << endl ;	
	  delete maxFinder ;	
	  throw string("RooGenContext::ctor()") ;	  
	}	
	delete maxFinder ;	

	cxcoutD(Generation) << "RooGenContext::ctor() maximum function value found through initial sampling is " << max << endl ;
      }
    }
      
  }

  if (_acceptRejectFunc && _otherVars.getSize()>0) {
    _generator = RooNumGenFactory::instance().createSampler(*_acceptRejectFunc,_otherVars,RooArgSet(_protoVars),*model.getGeneratorConfig(),_verbose,_maxVar) ;    
    cxcoutD(Generation) << "RooGenContext::ctor() creating MC sampling generator " << _generator->IsA()->GetName() << "  from function for observables " << _otherVars << endl ;
    //_generator= new RooAcceptReject(*_acceptRejectFunc,_otherVars,RooNumGenConfig::defaultConfig(),_verbose,_maxVar);
  } else {
    _generator = 0 ;
  }

  delete protoDeps ;
  delete depList;
  _otherVars.add(_uniformVars);
}


//_____________________________________________________________________________
RooGenContext::~RooGenContext() 
{
  // Destructor.

  // Clean up the cloned objects used in this context.
  delete _cloneSet;

  // Clean up our accept/reject generator
  if (_generator) delete _generator;
  if (_acceptRejectFunc) delete _acceptRejectFunc;
  if (_maxVar) delete _maxVar ;
  if (_uniIter) delete _uniIter ;
}



//_____________________________________________________________________________
void RooGenContext::attach(const RooArgSet& args) 
{
  // Attach the cloned model to the event buffer we will be filling.
  
  _pdfClone->recursiveRedirectServers(args,kFALSE);
  if (_acceptRejectFunc) {
    _acceptRejectFunc->recursiveRedirectServers(args,kFALSE) ; // WVE DEBUG
  }

  // Attach the RooAcceptReject generator the event buffer
  if (_generator) {
    _generator->attachParameters(args) ;
  }

}



//_____________________________________________________________________________
void RooGenContext::initGenerator(const RooArgSet &theEvent) 
{
  // Perform one-time initialization of the generator context

  RooFIter iter = theEvent.fwdIterator() ;
  RooAbsArg* arg ;
  while((arg=iter.next())) {
    arg->setOperMode(RooAbsArg::ADirty) ;
  }

  attach(theEvent) ;

  // Reset the cloned model's error counters.
  _pdfClone->resetErrorCounters();

  // Initialize the PDFs internal generator
  if (_directVars.getSize()>0) {
    cxcoutD(Generation) << "RooGenContext::initGenerator() initializing internal generator of model with code " << _code << endl ;
    _pdfClone->initGenerator(_code) ;
  }

  // Create iterator for uniform vars (if any)
  if (_uniformVars.getSize()>0) {
    _uniIter = _uniformVars.createIterator() ;
  }
}


//_____________________________________________________________________________
void RooGenContext::generateEvent(RooArgSet &theEvent, Int_t remaining) 
{
  // Generate one event. The 'remaining' integer is not used other than
  // for printing messages 

  if(_otherVars.getSize() > 0) {
    // call the accept-reject generator to generate its variables

    if (_updateFMaxPerEvent!=0) {
      Double_t max = _pdfClone->maxVal(_updateFMaxPerEvent)/_pdfClone->getNorm(_otherVars) ;
      cxcoutD(Generation) << "RooGenContext::initGenerator() reevaluation of maximum function value is required for each event, new value is  " << max << endl ;
      _maxVar->setVal(max) ;
    }

    if (_generator) {
      Double_t resampleRatio(1) ;
      const RooArgSet *subEvent= _generator->generateEvent(remaining,resampleRatio);
      if (resampleRatio<1) {
	coutI(Generation) << "RooGenContext::generateEvent INFO: accept/reject generator requests resampling of previously produced events by factor " 
			  << resampleRatio << " due to increased maximum weight" << endl ; 
	resampleData(resampleRatio) ;
      }
      if(0 == subEvent) {
	coutE(Generation) << "RooGenContext::generateEvent ERROR accept/reject generator failed" << endl ;
	return;
      }
      theEvent.assignValueOnly(*subEvent) ;
      //theEvent= *subEvent;
      
    }
  }

  // Use the model's optimized generator, if one is available.
  // The generator writes directly into our local 'event' since we attached it above.
  if(_directVars.getSize() > 0) {
    _pdfClone->generateEvent(_code);
  }

  // Generate uniform variables (non-dependents)  
  if (_uniIter) {
    _uniIter->Reset() ;
    RooAbsArg* uniVar ;
    while((uniVar=(RooAbsArg*)_uniIter->Next())) {
      RooAbsLValue* arglv = dynamic_cast<RooAbsLValue*>(uniVar) ;
      if (!arglv) {
	coutE(Generation) << "RooGenContext::generateEvent(" << GetName() << ") ERROR: uniform variable " << uniVar->GetName() << " is not an lvalue" << endl ;
	RooErrorHandler::softAbort() ;
      }
      arglv->randomize() ;
    }
    theEvent = _uniformVars ;
  }

}


//_____________________________________________________________________________
void RooGenContext::printMultiline(ostream &os, Int_t content, Bool_t verbose, TString indent) const
{
  // Printing interface

  RooAbsGenContext::printMultiline(os,content,verbose,indent);
  os << indent << " --- RooGenContext --- " << endl ;
  os << indent << "Using PDF ";
  _pdfClone->printStream(os,kName|kArgs|kClassName,kSingleLine,indent);
  
  if(verbose) {
    os << indent << "Use PDF generator for " << _directVars << endl ;
    os << indent << "Use MC sampling generator " << (_generator ? _generator->IsA()->GetName() : "<none>") << " for " << _otherVars << endl ;
    if (_protoVars.getSize()>0) {
      os << indent << "Prototype observables are " << _protoVars << endl ;
    }
  }
}