    /*****************************************************************************
     * Project: RooFit                                                           *
     * Package: RooFitCore                                                       *
     *    File: $Id: RooMCIntegrator.cc,v 1.22 2005/06/20 15:44:54 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] --
    // RooMCIntegrator implements an adaptive multi-dimensional Monte Carlo
    // numerical integration, following the VEGAS algorithm originally described
    // in G. P. Lepage, J. Comp. Phys. 27, 192(1978). This implementation is
    // based on a C version from the 0.9 beta release of the GNU scientific library.

    #include "RooFit.h"

    #include "RooMCIntegrator.h"
    #include "RooMCIntegrator.h"
    #include "RooArgSet.h"
    #include "RooNumber.h"
    #include "RooAbsArg.h"
    #include "RooNumIntFactory.h"
    #include "RooRealVar.h"
    #include "RooCategory.h"

    #include <math.h>
    #include <assert.h>

    ClassImp(RooMCIntegrator)
    ;

    // Register this class with RooNumIntFactory
    static void registerMCIntegrator(RooNumIntFactory& fact)
    {
      // Construct default configuration
      RooCategory samplingMode("samplingMode","Sampling Mode") ;
      samplingMode.defineType("Importance",RooMCIntegrator::Importance) ;
      samplingMode.defineType("ImportanceOnly",RooMCIntegrator::ImportanceOnly) ;
      samplingMode.defineType("Stratified",RooMCIntegrator::Stratified) ;
      samplingMode.setIndex(RooMCIntegrator::Importance) ;

      RooCategory genType("genType","Generator Type") ;
      genType.defineType("QuasiRandom",RooMCIntegrator::QuasiRandom) ;
      genType.defineType("PseudoRandom",RooMCIntegrator::PseudoRandom) ;
      genType.setIndex(RooMCIntegrator::QuasiRandom) ;

      RooCategory verbose("verbose","Verbose flag") ;
      verbose.defineType("true",1) ;
      verbose.defineType("false",0) ;
      verbose.setIndex(0) ;

      RooRealVar alpha("alpha","Grid structure constant",1.5) ;
      RooRealVar nRefineIter("nRefineIter","Number of refining iterations",5) ;
      RooRealVar nRefinePerDim("nRefinePerDim","Number of refining samples (per dimension)",1000) ;
      RooRealVar nIntPerDim("nIntPerDim","Number of integration samples (per dimension)",5000) ;

      // Create prototype integrator
      RooMCIntegrator* proto = new RooMCIntegrator() ;

      // Register prototype and default config with factory
      fact.storeProtoIntegrator(proto,RooArgSet(samplingMode,genType,verbose,alpha,nRefineIter,nRefinePerDim,nIntPerDim)) ;

      // Make this method the default for all N>2-dim integrals
      RooNumIntConfig::defaultConfig().methodND().setLabel(proto->IsA()->GetName()) ;
    }
    static Bool_t dummy = RooNumIntFactory::instance().registerInitializer(&registerMCIntegrator) ;

     RooMCIntegrator::RooMCIntegrator()
    {
      // Dummy default ctor
    }

     RooMCIntegrator::RooMCIntegrator(const RooAbsFunc& function, SamplingMode mode,
    				 GeneratorType genType, Bool_t verbose) :
      RooAbsIntegrator(function), _grid(function), _verbose(verbose),
      _alpha(1.5),  _mode(mode), _genType(genType),
      _nRefineIter(5),_nRefinePerDim(1000),_nIntegratePerDim(5000)
    {
      // check that our grid initialized without errors
      if(!(_valid= _grid.isValid())) return;
      if(_verbose) _grid.Print();
    }

     RooMCIntegrator::RooMCIntegrator(const RooAbsFunc& function, const RooNumIntConfig& config) :
      RooAbsIntegrator(function), _grid(function)
    {
      const RooArgSet& configSet = config.getConfigSection(IsA()->GetName()) ;
      _verbose = (Bool_t) configSet.getCatIndex("verbose",0) ;
      _alpha = configSet.getRealValue("alpha",1.5) ;
      _mode = (SamplingMode) configSet.getCatIndex("samplingMode",Importance) ;
      _genType = (GeneratorType) configSet.getCatIndex("genType",QuasiRandom) ;
      _nRefineIter = (Int_t) configSet.getRealValue("nRefineIter",5) ;
      _nRefinePerDim = (Int_t) configSet.getRealValue("nRefinePerDim",1000) ;
      _nIntegratePerDim = (Int_t) configSet.getRealValue("nIntPerDim",5000) ;

      // check that our grid initialized without errors
      if(!(_valid= _grid.isValid())) return;
      if(_verbose) _grid.Print();
    }

     RooAbsIntegrator* RooMCIntegrator::clone(const RooAbsFunc& function, const RooNumIntConfig& config) const
    {
      return new RooMCIntegrator(function,config) ;
    }

     RooMCIntegrator::~RooMCIntegrator() {
    }

     Bool_t RooMCIntegrator::checkLimits() const {
      // Check if we can integrate over the current domain.

      return _grid.initialize(*integrand());
    }

     Double_t RooMCIntegrator::integral(const Double_t* /*yvec*/) {
      // Evaluate the integral using a fixed number of calls to evaluate the integrand
      // equal to about 10k per dimension. Use the first 5k calls to refine the grid
      // over 5 iterations of 1k calls each, and the remaining 5k calls for a single
      // high statistics integration.
      _timer.Start(kTRUE);
      vegas(AllStages,_nRefinePerDim*_grid.getDimension(),_nRefineIter);
      Double_t ret = vegas(ReuseGrid,_nIntegratePerDim*_grid.getDimension(),1);
      return ret ;
    }

     Double_t RooMCIntegrator::vegas(Stage stage, UInt_t calls, UInt_t iterations, Double_t *absError) {
      // Perform one step of Monte Carlo integration using the specified number of iterations
      // with (approximately) the specified number of integrand evaluation calls per iteration.
      // Use the VEGAS algorithm, starting from the specified stage. Returns the best estimate
      // of the integral. Also sets *absError to the estimated absolute error of the integral
      // estimate if absError is non-zero.

      //cout << "VEGAS stage = " << stage << " calls = " << calls << " iterations = " << iterations << endl ;

      // reset the grid to its initial state if we are starting from scratch
      if(stage == AllStages) _grid.initialize(*_function);

      // reset the results of previous calculations on this grid, but reuse the grid itself.
      if(stage <= ReuseGrid) {
        _wtd_int_sum = 0;
        _sum_wgts = 0;
        _chi_sum = 0;
        _it_num = 1;
        _samples = 0;
      }

      // refine the results of previous calculations on the current grid.
      if(stage <= RefineGrid) {
        UInt_t bins = RooGrid::maxBins;
        UInt_t boxes = 1;
        UInt_t dim(_grid.getDimension());

        // select the sampling mode for the next step
        if(_mode != ImportanceOnly) {
          // calculate the largest number of equal subdivisions ("boxes") along each
          // axis that results in an average of no more than 2 integrand calls per cell
          boxes = (UInt_t)floor(TMath::Power(calls/2.0,1.0/dim));
          // use stratified sampling if we are allowed enough calls (or equivalently,
          // if the dimension is low enough)
          _mode = Importance;
          if (2*boxes >= RooGrid::maxBins) {
    	_mode = Stratified;
    	// adjust the number of bins and boxes to give an integral number >= 1 of boxes per bin
    	Int_t box_per_bin= (boxes > RooGrid::maxBins) ? boxes/RooGrid::maxBins : 1;
    	bins= boxes/box_per_bin;
    	if(bins > RooGrid::maxBins) bins= RooGrid::maxBins;
    	boxes = box_per_bin * bins;
    	if(_verbose) cout << "RooMCIntegrator: using stratified sampling with " << bins << " bins and "
    			  << box_per_bin << " boxes/bin" << endl;
          }
          else {
    	if(_verbose) cout << "RooMCIntegrator: using importance sampling with " << bins << " bins and "
    			  << boxes << " boxes" << endl;
          }
        }

        // calculate the total number of n-dim boxes for this step
        Double_t tot_boxes = TMath::Power((Double_t)boxes,(Double_t)dim);

        // increase the total number of calls to get at least 2 calls per box, if necessary
        _calls_per_box = (UInt_t)(calls/tot_boxes);
        if(_calls_per_box < 2) _calls_per_box= 2;
        calls= (UInt_t)(_calls_per_box*tot_boxes);

        // calculate the Jacobean factor: volume/(avg # of calls/bin)
        _jac = _grid.getVolume()*TMath::Power((Double_t)bins,(Double_t)dim)/calls;

        // setup our grid to use the calculated number of boxes and bins
        _grid.setNBoxes(boxes);
        if(bins != _grid.getNBins()) _grid.resize(bins);
      }

      // allocate memory for some book-keeping arrays
      UInt_t *box= _grid.createIndexVector();
      UInt_t *bin= _grid.createIndexVector();
      Double_t *x= _grid.createPoint();

      // loop over iterations for this step
      Double_t cum_int(0),cum_sig(0);
      _it_start = _it_num;
      _chisq = 0.0;
      for (UInt_t it = 0; it < iterations; it++) {
        Double_t intgrl(0),intgrl_sq(0),sig(0),jacbin(_jac);

        _it_num = _it_start + it;

        // reset the values associated with each grid cell
        _grid.resetValues();

        // loop over grid boxes
        _grid.firstBox(box);
        do {
          Double_t m(0),q(0);
          // loop over integrand evaluations within this grid box
          for(UInt_t k = 0; k < _calls_per_box; k++) {
    	// generate a random point in this box
    	Double_t bin_vol(0);
    	_grid.generatePoint(box, x, bin, bin_vol, _genType == QuasiRandom ? kTRUE : kFALSE);
    	// evaluate the integrand at the generated point
    	Double_t fval= jacbin*bin_vol*integrand(x);
    	// update mean and variance calculations
    	Double_t d = fval - m;
    	m+= d / (k + 1.0);
    	q+= d * d * (k / (k + 1.0));
    	// accumulate the results of this evaluation (importance sampling only)
    	if (_mode != Stratified) _grid.accumulate(bin, fval*fval);
          }
          intgrl += m * _calls_per_box;
          Double_t f_sq_sum = q * _calls_per_box ;
          sig += f_sq_sum ;

          // accumulate the results for this grid box (stratified sampling only)
          if (_mode == Stratified) _grid.accumulate(bin, f_sq_sum);

          // print occasional progress messages
          if(_timer.RealTime() > 1) { // wait at least 1 sec since the last message
    	cout << "RooMCIntegrator: still working..." << endl;
    	_timer.Start(kTRUE);
          }
          else {
    	_timer.Start(kFALSE);
          }

        } while(_grid.nextBox(box));

        // compute final results for this iteration
        Double_t wgt;
        sig = sig / (_calls_per_box - 1.0)  ;
        if (sig > 0) {
          wgt = 1.0 / sig;
        }
        else if (_sum_wgts > 0) {
          wgt = _sum_wgts / _samples;
        }
        else {
          wgt = 0.0;
        }
        intgrl_sq = intgrl * intgrl;
        _result = intgrl;
        _sigma  = sqrt(sig);

        if (wgt > 0.0) {
          _samples++ ;
          _sum_wgts += wgt;
          _wtd_int_sum += intgrl * wgt;
          _chi_sum += intgrl_sq * wgt;

          cum_int = _wtd_int_sum / _sum_wgts;
          cum_sig = sqrt (1 / _sum_wgts);

          if (_samples > 1) {
    	_chisq = (_chi_sum - _wtd_int_sum * cum_int)/(_samples - 1.0);
          }
        }
        else {
          cum_int += (intgrl - cum_int) / (it + 1.0);
          cum_sig = 0.0;
        }
        if (_verbose) {
          cout << "=== Iteration " << _it_num << " : I = " << intgrl << " +/- " << sqrt(sig) << endl
    	   << "    Cumulative : I = " << cum_int << " +/- " << cum_sig << "( chi2 = " << _chisq
    	   << ")" << endl;
          // print the grid after the final iteration
          if(it + 1 == iterations) _grid.Print("V");
        }
        _grid.refine(_alpha);
      }

      // cleanup
      delete[] bin;
      delete[] box;
      delete[] x;

      if(absError) *absError = cum_sig;
      return cum_int;
    }