/*****************************************************************************
 * 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 
// RooBinnedGenContext is an efficient implementation of the
// generator context specific for binned pdfs
// END_HTML
//


#include "RooFit.h"

#include "Riostream.h"


#include "RooMsgService.h"
#include "RooBinnedGenContext.h"
#include "RooAbsPdf.h"
#include "RooRealVar.h"
#include "RooDataHist.h"
#include "RooDataSet.h"
#include "RooRandom.h"

using namespace std;

ClassImp(RooBinnedGenContext)
;
  

//_____________________________________________________________________________
RooBinnedGenContext::RooBinnedGenContext(const RooAbsPdf &model, const RooArgSet &vars, 
				   const RooDataSet *prototype, const RooArgSet* auxProto,
				   Bool_t verbose) :
  RooAbsGenContext(model,vars,prototype,auxProto,verbose)
{
  // Constructor

  cxcoutI(Generation) << "RooBinnedGenContext::ctor() setting up event special generator context for sum 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 ;
  ccxcoutI(Generation) << endl ;

  // Constructor. Build an array of generator contexts for each product component PDF
  _pdfSet = (RooArgSet*) RooArgSet(model).snapshot(kTRUE) ;
  _pdf = (RooAbsPdf*) _pdfSet->find(model.GetName()) ;
  _pdf->setOperMode(RooAbsArg::ADirty,kTRUE) ;

  // Fix normalization set of this RooAddPdf
  if (prototype) 
    {
      RooArgSet coefNSet(vars) ;
      coefNSet.add(*prototype->get()) ;
      _pdf->fixAddCoefNormalization(coefNSet) ;
    }

  _pdf->recursiveRedirectServers(*_theEvent) ;
  _vars = _pdf->getObservables(vars) ;

  // If pdf has boundary definitions, follow those for the binning 
  RooFIter viter = _vars->fwdIterator() ;
  RooAbsArg* var ;
  while((var=viter.next())) {
    RooRealVar* rvar = dynamic_cast<RooRealVar*>(var) ;
    if (rvar) {
      list<Double_t>* binb = model.binBoundaries(*rvar,rvar->getMin(),rvar->getMax()) ;
      delete binb ;
    }
  }


  // Create empty RooDataHist
  _hist = new RooDataHist("genData","genData",*_vars) ;

  _expectedData = kFALSE ;
}


//_____________________________________________________________________________
RooBinnedGenContext::~RooBinnedGenContext()
{
  // Destructor. Delete all owned subgenerator contexts

  delete _vars ;
  delete _pdfSet ;
  delete _hist ;
}



//_____________________________________________________________________________
void RooBinnedGenContext::attach(const RooArgSet& args) 
{
  // Attach given set of variables to internal p.d.f. clone

  _pdf->recursiveRedirectServers(args) ;
}



//_____________________________________________________________________________
void RooBinnedGenContext::initGenerator(const RooArgSet &theEvent)
{
  // One-time initialization of generator contex. Attach theEvent
  // to internal p.d.f clone and forward initialization call to 
  // the component generators

  _pdf->recursiveRedirectServers(theEvent) ;

}


//_____________________________________________________________________________
void RooBinnedGenContext::setExpectedData(Bool_t flag) 
{
  _expectedData = flag ;
}


//_____________________________________________________________________________
RooDataSet *RooBinnedGenContext::generate(Double_t nEvt, Bool_t /*skipInit*/, Bool_t extended)
{
  
  // Scale to number of events and introduce Poisson fluctuations
  _hist->reset() ;

  Double_t nEvents = nEvt ;
  
  if (nEvents<=0) {
    if (!_pdf->canBeExtended()) {
      coutE(InputArguments) << "RooAbsPdf::generateBinned(" << GetName() 
			    << ") ERROR: No event count provided and p.d.f does not provide expected number of events" << endl ;
      return 0 ;
    } else {
      // Don't round in expectedData mode
      if (_expectedData || extended) {
	nEvents = _pdf->expectedEvents(_vars) ;
      } else {
	nEvents = Int_t(_pdf->expectedEvents(_vars)+0.5) ;
      }
    }
  } 

  // Sample p.d.f. distribution
  _pdf->fillDataHist(_hist,_vars,1,kTRUE) ;  
  
  // Output container
  RooRealVar weight("weight","weight",0,1e9) ;
  RooArgSet tmp(*_vars) ;
  tmp.add(weight) ;
  RooDataSet* wudata = new RooDataSet("wu","wu",tmp,RooFit::WeightVar("weight")) ;

  vector<int> histOut(_hist->numEntries()) ;
  Double_t histMax(-1) ;
  Int_t histOutSum(0) ;
  for (int i=0 ; i<_hist->numEntries() ; i++) {
    _hist->get(i) ;
    if (_expectedData) { 

      // Expected data, multiply p.d.f by nEvents
      Double_t w=_hist->weight()*nEvents ;
      wudata->add(*_hist->get(),w) ;

    } else if (extended) {

      // Extended mode, set contents to Poisson(pdf*nEvents)
      Double_t w = RooRandom::randomGenerator()->Poisson(_hist->weight()*nEvents) ;
      wudata->add(*_hist->get(),w) ;
      
    } else {

      // Regular mode, fill array of weights with Poisson(pdf*nEvents), but to not fill
      // histogram yet.
      if (_hist->weight()>histMax) {
	histMax = _hist->weight() ;
      }
      histOut[i] = RooRandom::randomGenerator()->Poisson(_hist->weight()*nEvents) ;
      histOutSum += histOut[i] ;
    }
  }


  if (!_expectedData && !extended) {

    // Second pass for regular mode - Trim/Extend dataset to exact number of entries

    // Calculate difference between what is generated so far and what is requested
    Int_t nEvtExtra = abs(Int_t(nEvents)-histOutSum) ;
    Int_t wgt = (histOutSum>nEvents) ? -1 : 1 ;

    // Perform simple binned accept/reject procedure to get to exact event count
    while(nEvtExtra>0) {

      Int_t ibinRand = RooRandom::randomGenerator()->Integer(_hist->numEntries()) ;
      _hist->get(ibinRand) ;
      Double_t ranY = RooRandom::randomGenerator()->Uniform(histMax) ;

      if (ranY<_hist->weight()) {
	if (wgt==1) {
	  histOut[ibinRand]++ ;
	} else {
	  // If weight is negative, prior bin content must be at least 1
	  if (histOut[ibinRand]>0) {
	    histOut[ibinRand]-- ;
	  } else {
	    continue ;
	  }
	}
	nEvtExtra-- ;
      }
    }

    // Transfer working array to histogram
    for (int i=0 ; i<_hist->numEntries() ; i++) {
      _hist->get(i) ;
      wudata->add(*_hist->get(),histOut[i]) ;      
    }    

  } 

  return wudata ;

}



//_____________________________________________________________________________
void RooBinnedGenContext::generateEvent(RooArgSet&, Int_t)
{
  // this method is not implemented for this context
  assert(0) ;
}



//_____________________________________________________________________________
void RooBinnedGenContext::printMultiline(ostream &os, Int_t content, Bool_t verbose, TString indent) const 
{
  // Print the details of the context

  RooAbsGenContext::printMultiline(os,content,verbose,indent) ;
  os << indent << "--- RooBinnedGenContext ---" << endl ;
  os << indent << "Using PDF ";
  _pdf->printStream(os,kName|kArgs|kClassName,kSingleLine,indent);
}
 RooBinnedGenContext.cxx:1
 RooBinnedGenContext.cxx:2
 RooBinnedGenContext.cxx:3
 RooBinnedGenContext.cxx:4
 RooBinnedGenContext.cxx:5
 RooBinnedGenContext.cxx:6
 RooBinnedGenContext.cxx:7
 RooBinnedGenContext.cxx:8
 RooBinnedGenContext.cxx:9
 RooBinnedGenContext.cxx:10
 RooBinnedGenContext.cxx:11
 RooBinnedGenContext.cxx:12
 RooBinnedGenContext.cxx:13
 RooBinnedGenContext.cxx:14
 RooBinnedGenContext.cxx:15
 RooBinnedGenContext.cxx:16
 RooBinnedGenContext.cxx:17
 RooBinnedGenContext.cxx:18
 RooBinnedGenContext.cxx:19
 RooBinnedGenContext.cxx:20
 RooBinnedGenContext.cxx:21
 RooBinnedGenContext.cxx:22
 RooBinnedGenContext.cxx:23
 RooBinnedGenContext.cxx:24
 RooBinnedGenContext.cxx:25
 RooBinnedGenContext.cxx:26
 RooBinnedGenContext.cxx:27
 RooBinnedGenContext.cxx:28
 RooBinnedGenContext.cxx:29
 RooBinnedGenContext.cxx:30
 RooBinnedGenContext.cxx:31
 RooBinnedGenContext.cxx:32
 RooBinnedGenContext.cxx:33
 RooBinnedGenContext.cxx:34
 RooBinnedGenContext.cxx:35
 RooBinnedGenContext.cxx:36
 RooBinnedGenContext.cxx:37
 RooBinnedGenContext.cxx:38
 RooBinnedGenContext.cxx:39
 RooBinnedGenContext.cxx:40
 RooBinnedGenContext.cxx:41
 RooBinnedGenContext.cxx:42
 RooBinnedGenContext.cxx:43
 RooBinnedGenContext.cxx:44
 RooBinnedGenContext.cxx:45
 RooBinnedGenContext.cxx:46
 RooBinnedGenContext.cxx:47
 RooBinnedGenContext.cxx:48
 RooBinnedGenContext.cxx:49
 RooBinnedGenContext.cxx:50
 RooBinnedGenContext.cxx:51
 RooBinnedGenContext.cxx:52
 RooBinnedGenContext.cxx:53
 RooBinnedGenContext.cxx:54
 RooBinnedGenContext.cxx:55
 RooBinnedGenContext.cxx:56
 RooBinnedGenContext.cxx:57
 RooBinnedGenContext.cxx:58
 RooBinnedGenContext.cxx:59
 RooBinnedGenContext.cxx:60
 RooBinnedGenContext.cxx:61
 RooBinnedGenContext.cxx:62
 RooBinnedGenContext.cxx:63
 RooBinnedGenContext.cxx:64
 RooBinnedGenContext.cxx:65
 RooBinnedGenContext.cxx:66
 RooBinnedGenContext.cxx:67
 RooBinnedGenContext.cxx:68
 RooBinnedGenContext.cxx:69
 RooBinnedGenContext.cxx:70
 RooBinnedGenContext.cxx:71
 RooBinnedGenContext.cxx:72
 RooBinnedGenContext.cxx:73
 RooBinnedGenContext.cxx:74
 RooBinnedGenContext.cxx:75
 RooBinnedGenContext.cxx:76
 RooBinnedGenContext.cxx:77
 RooBinnedGenContext.cxx:78
 RooBinnedGenContext.cxx:79
 RooBinnedGenContext.cxx:80
 RooBinnedGenContext.cxx:81
 RooBinnedGenContext.cxx:82
 RooBinnedGenContext.cxx:83
 RooBinnedGenContext.cxx:84
 RooBinnedGenContext.cxx:85
 RooBinnedGenContext.cxx:86
 RooBinnedGenContext.cxx:87
 RooBinnedGenContext.cxx:88
 RooBinnedGenContext.cxx:89
 RooBinnedGenContext.cxx:90
 RooBinnedGenContext.cxx:91
 RooBinnedGenContext.cxx:92
 RooBinnedGenContext.cxx:93
 RooBinnedGenContext.cxx:94
 RooBinnedGenContext.cxx:95
 RooBinnedGenContext.cxx:96
 RooBinnedGenContext.cxx:97
 RooBinnedGenContext.cxx:98
 RooBinnedGenContext.cxx:99
 RooBinnedGenContext.cxx:100
 RooBinnedGenContext.cxx:101
 RooBinnedGenContext.cxx:102
 RooBinnedGenContext.cxx:103
 RooBinnedGenContext.cxx:104
 RooBinnedGenContext.cxx:105
 RooBinnedGenContext.cxx:106
 RooBinnedGenContext.cxx:107
 RooBinnedGenContext.cxx:108
 RooBinnedGenContext.cxx:109
 RooBinnedGenContext.cxx:110
 RooBinnedGenContext.cxx:111
 RooBinnedGenContext.cxx:112
 RooBinnedGenContext.cxx:113
 RooBinnedGenContext.cxx:114
 RooBinnedGenContext.cxx:115
 RooBinnedGenContext.cxx:116
 RooBinnedGenContext.cxx:117
 RooBinnedGenContext.cxx:118
 RooBinnedGenContext.cxx:119
 RooBinnedGenContext.cxx:120
 RooBinnedGenContext.cxx:121
 RooBinnedGenContext.cxx:122
 RooBinnedGenContext.cxx:123
 RooBinnedGenContext.cxx:124
 RooBinnedGenContext.cxx:125
 RooBinnedGenContext.cxx:126
 RooBinnedGenContext.cxx:127
 RooBinnedGenContext.cxx:128
 RooBinnedGenContext.cxx:129
 RooBinnedGenContext.cxx:130
 RooBinnedGenContext.cxx:131
 RooBinnedGenContext.cxx:132
 RooBinnedGenContext.cxx:133
 RooBinnedGenContext.cxx:134
 RooBinnedGenContext.cxx:135
 RooBinnedGenContext.cxx:136
 RooBinnedGenContext.cxx:137
 RooBinnedGenContext.cxx:138
 RooBinnedGenContext.cxx:139
 RooBinnedGenContext.cxx:140
 RooBinnedGenContext.cxx:141
 RooBinnedGenContext.cxx:142
 RooBinnedGenContext.cxx:143
 RooBinnedGenContext.cxx:144
 RooBinnedGenContext.cxx:145
 RooBinnedGenContext.cxx:146
 RooBinnedGenContext.cxx:147
 RooBinnedGenContext.cxx:148
 RooBinnedGenContext.cxx:149
 RooBinnedGenContext.cxx:150
 RooBinnedGenContext.cxx:151
 RooBinnedGenContext.cxx:152
 RooBinnedGenContext.cxx:153
 RooBinnedGenContext.cxx:154
 RooBinnedGenContext.cxx:155
 RooBinnedGenContext.cxx:156
 RooBinnedGenContext.cxx:157
 RooBinnedGenContext.cxx:158
 RooBinnedGenContext.cxx:159
 RooBinnedGenContext.cxx:160
 RooBinnedGenContext.cxx:161
 RooBinnedGenContext.cxx:162
 RooBinnedGenContext.cxx:163
 RooBinnedGenContext.cxx:164
 RooBinnedGenContext.cxx:165
 RooBinnedGenContext.cxx:166
 RooBinnedGenContext.cxx:167
 RooBinnedGenContext.cxx:168
 RooBinnedGenContext.cxx:169
 RooBinnedGenContext.cxx:170
 RooBinnedGenContext.cxx:171
 RooBinnedGenContext.cxx:172
 RooBinnedGenContext.cxx:173
 RooBinnedGenContext.cxx:174
 RooBinnedGenContext.cxx:175
 RooBinnedGenContext.cxx:176
 RooBinnedGenContext.cxx:177
 RooBinnedGenContext.cxx:178
 RooBinnedGenContext.cxx:179
 RooBinnedGenContext.cxx:180
 RooBinnedGenContext.cxx:181
 RooBinnedGenContext.cxx:182
 RooBinnedGenContext.cxx:183
 RooBinnedGenContext.cxx:184
 RooBinnedGenContext.cxx:185
 RooBinnedGenContext.cxx:186
 RooBinnedGenContext.cxx:187
 RooBinnedGenContext.cxx:188
 RooBinnedGenContext.cxx:189
 RooBinnedGenContext.cxx:190
 RooBinnedGenContext.cxx:191
 RooBinnedGenContext.cxx:192
 RooBinnedGenContext.cxx:193
 RooBinnedGenContext.cxx:194
 RooBinnedGenContext.cxx:195
 RooBinnedGenContext.cxx:196
 RooBinnedGenContext.cxx:197
 RooBinnedGenContext.cxx:198
 RooBinnedGenContext.cxx:199
 RooBinnedGenContext.cxx:200
 RooBinnedGenContext.cxx:201
 RooBinnedGenContext.cxx:202
 RooBinnedGenContext.cxx:203
 RooBinnedGenContext.cxx:204
 RooBinnedGenContext.cxx:205
 RooBinnedGenContext.cxx:206
 RooBinnedGenContext.cxx:207
 RooBinnedGenContext.cxx:208
 RooBinnedGenContext.cxx:209
 RooBinnedGenContext.cxx:210
 RooBinnedGenContext.cxx:211
 RooBinnedGenContext.cxx:212
 RooBinnedGenContext.cxx:213
 RooBinnedGenContext.cxx:214
 RooBinnedGenContext.cxx:215
 RooBinnedGenContext.cxx:216
 RooBinnedGenContext.cxx:217
 RooBinnedGenContext.cxx:218
 RooBinnedGenContext.cxx:219
 RooBinnedGenContext.cxx:220
 RooBinnedGenContext.cxx:221
 RooBinnedGenContext.cxx:222
 RooBinnedGenContext.cxx:223
 RooBinnedGenContext.cxx:224
 RooBinnedGenContext.cxx:225
 RooBinnedGenContext.cxx:226
 RooBinnedGenContext.cxx:227
 RooBinnedGenContext.cxx:228
 RooBinnedGenContext.cxx:229
 RooBinnedGenContext.cxx:230
 RooBinnedGenContext.cxx:231
 RooBinnedGenContext.cxx:232
 RooBinnedGenContext.cxx:233
 RooBinnedGenContext.cxx:234
 RooBinnedGenContext.cxx:235
 RooBinnedGenContext.cxx:236
 RooBinnedGenContext.cxx:237
 RooBinnedGenContext.cxx:238
 RooBinnedGenContext.cxx:239
 RooBinnedGenContext.cxx:240
 RooBinnedGenContext.cxx:241
 RooBinnedGenContext.cxx:242
 RooBinnedGenContext.cxx:243
 RooBinnedGenContext.cxx:244
 RooBinnedGenContext.cxx:245
 RooBinnedGenContext.cxx:246
 RooBinnedGenContext.cxx:247
 RooBinnedGenContext.cxx:248
 RooBinnedGenContext.cxx:249
 RooBinnedGenContext.cxx:250
 RooBinnedGenContext.cxx:251
 RooBinnedGenContext.cxx:252
 RooBinnedGenContext.cxx:253
 RooBinnedGenContext.cxx:254
 RooBinnedGenContext.cxx:255
 RooBinnedGenContext.cxx:256
 RooBinnedGenContext.cxx:257
 RooBinnedGenContext.cxx:258
 RooBinnedGenContext.cxx:259
 RooBinnedGenContext.cxx:260