/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitCore                                                       *
 *    File: $Id: RooGenContext.cc,v 1.45 2005/06/20 15:44:52 wverkerke Exp $
 * 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)             *
 *****************************************************************************/

// -- CLASS DESCRIPTION [AUX] --
// A class description belongs here...


#include "RooFit.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 "TString.h"
#include "TIterator.h"

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.

//   cout << "RooGenContext::ctor(" << model.GetName() << ")" << 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) {
    cout << "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());

  // 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->isDerived()) {
      cout << ClassName() << "::" << GetName() << ": 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) {
      cout << ClassName() << "::" << GetName() << ":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)) {
	  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()) {
    cout << ClassName() << "::" << GetName() << ": constructor failed with errors" << endl;
    return;
  }

  // If PDF depends on prototype data, direct generator cannot use static initialization
  // in initGenerator()
  Bool_t staticInitOK = !_pdfClone->dependsOn(_protoVars) ;
  
  // Can the model generate any of the direct variables itself?
  RooArgSet generatedVars;
  _code= _pdfClone->getGenerator(_directVars,generatedVars,staticInitOK);

  // 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);

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

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

  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) {
	if (verbose) {
	  cout << "RooGenContext::ctor(" << model.GetName() 
	       << ") PDF supports maximum finding, initial sampling phase skipped" << endl ;
	}
	Double_t maxVal = _pdfClone->maxVal(maxFindCode) / _pdfClone->getNorm(_theEvent) ;
	_maxVar = new RooRealVar("funcMax","function maximum",maxVal) ;
      }
    } 

    _pdfClone->getVal(&vars) ; // WVE debug
    _acceptRejectFunc= new RooRealIntegral(nname,ntitle,*_pdfClone,*depList,&vars);

  } else {

    // Generation with prototype variable
    depList->remove(_protoVars,kTRUE,kTRUE) ;
    _acceptRejectFunc= new RooRealIntegral(nname,ntitle,*_pdfClone,*depList,&vars);

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

      // Check if PDF supports maximum finding
      Int_t maxFindCode = _pdfClone->getMaxVal(_otherVars) ;
      if (maxFindCode != 0) {

	// Special case: PDF supports max-finding in otherVars, no need to scan other+proto space for maximum
	if (verbose) {
	  cout << "RooGenContext::ctor(" << model.GetName() 
	       << ") PDF supports maximum finding, initial sampling phase skipped" << endl ;
	}
	_maxVar = new RooRealVar("funcMax","function maximum",1) ;
	_updateFMaxPerEvent = maxFindCode ;
      }
    }


    if (!_maxVar) {

      // Regular case: First find maximum in other+proto space
      RooArgSet otherAndProto(_otherVars) ;
      
      RooArgSet* protoDeps = model.getObservables(_protoVars) ;
      otherAndProto.add(*protoDeps) ;
      delete protoDeps ;
      
      if (_otherVars.getSize()>0) {      
	// Calculate maximum in other+proto space if there are any accept/reject generated observables
	RooAcceptReject maxFinder(*_acceptRejectFunc,otherAndProto,0,_verbose) ;
	Double_t max = maxFinder.getFuncMax() ;
	_maxVar = new RooRealVar("funcMax","function maximum",max) ;
      }
    }
      
  }

  _generator= new RooAcceptReject(*_acceptRejectFunc,_otherVars,_maxVar,_verbose);

  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
  delete _generator;
  delete _acceptRejectFunc;
  if (_maxVar) delete _maxVar ;
  if (_uniIter) delete _uniIter ;
}

 void RooGenContext::initGenerator(const RooArgSet &theEvent) {

  // Attach the cloned model to the event buffer we will be filling.
  _pdfClone->recursiveRedirectServers(theEvent,kFALSE);

  // Attach the RooAcceptReject generator the event buffer
  _generator->attachParameters(theEvent) ;

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

  // Initialize the PDFs internal generator
  if (_directVars.getSize()>0) {
    _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 variables for a new event.

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


    if (_updateFMaxPerEvent!=0) {
      _maxVar->setVal(_pdfClone->maxVal(_updateFMaxPerEvent)/_pdfClone->getNorm(_otherVars)) ;
    }

    const RooArgSet *subEvent= _generator->generateEvent(remaining);
    if(0 == subEvent) {
      cout << ClassName() << "::" << GetName() << ":generate: accept/reject generator failed." << endl;
      return;
    }
    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) {
	cout << "RooGenContext::generateEvent(" << GetName() << ") ERROR: uniform variable " 
	     << uniVar->GetName() << " is not an lvalue" << endl ;
	RooErrorHandler::softAbort() ;
      }
      arglv->randomize() ;
    }
    theEvent = _uniformVars ;
  }
}

 void RooGenContext::printToStream(ostream &os, PrintOption opt, TString indent) const
{
  RooAbsGenContext::printToStream(os,opt,indent);
  if(opt >= Standard) {
    PrintOption less= lessVerbose(opt);
    TString deeper(indent);
    indent.Append("  ");
    os << indent << "Using PDF ";
    _pdfClone->printToStream(os,less,deeper);
    if(opt >= Verbose) {
      os << indent << "Use PDF generator for ";
      _directVars.printToStream(os,less,deeper);
      os << indent << "Use accept/reject for ";
      _otherVars.printToStream(os,less,deeper);
    }
  }
}