/***************************************************************************** 
  * Project: RooFit                                                           * 
  *                                                                           * 
  * Binomial PDF                                                              *
  * author: Marco Bomben <marco.bomben@cern.ch>                               *
  * Inspired by RooPoisson by Kyle Cranmer <cranmer@cern.ch>                  * 
  *
  * Wed Oct 19 18:53:51 CEST 2022 
  *****************************************************************************/ 

#ifndef ROOBINOMIAL
#define ROOBINOMIAL

#include "RooAbsPdf.h"
#include "RooRealProxy.h"
#include "RooCategoryProxy.h"
#include "RooAbsReal.h"
#include "RooAbsCategory.h"
#include "RooTrace.h"
#include "RooBatchCompute.h"
#include "RooNaNPacker.h"
//#include "BatchModeHelpers.h"

class RooBinomial : public RooAbsPdf {
public:
  RooBinomial() { _noRounding = kFALSE ;   } ;
  RooBinomial(const char *name, const char *title, RooAbsReal& _k, RooAbsReal& _N, RooAbsReal& _p,  Bool_t noRounding=kFALSE);
  RooBinomial(const RooBinomial& other, const char* name=0) ;
  virtual TObject* clone(const char* newname) const override { return new RooBinomial(*this,newname); }
  inline virtual ~RooBinomial() {  }

  Int_t getAnalyticalIntegral(RooArgSet& allVars, RooArgSet& analVars, const char* rangeName=0) const override;
  Double_t analyticalIntegral(Int_t code, const char* rangeName=0) const override;

  Int_t getGenerator(const RooArgSet& directVars, RooArgSet &generateVars, Bool_t staticInitOK=kTRUE) const override;
  void generateEvent(Int_t code) override;
  
  void setNoRounding(bool flag = kTRUE) {_noRounding = flag;}
  void protectNegativeMean(bool flag = kTRUE) {_protectNegative = flag;}  // is this needed?

  /// Get the x variable.
  RooAbsReal const& getK() const { return k.arg(); }

  /// Get the probability parameter.
  RooAbsReal const& getP() const { return p.arg(); }

protected:

  RooRealProxy k ;
  RooRealProxy N ;
  RooRealProxy p ;
  Bool_t  _noRounding ;
  Bool_t  _protectNegative{true};
  
  Double_t evaluate() const override;
  void computeBatch(cudaStream_t*, double* output, size_t nEvents, RooBatchCompute::DataMap&) const override;
  inline bool canComputeBatchWithCuda() const override { return true; }
  // RooSpan<double> evaluateBatch(std::size_t begin, std::size_t batchSize) const override; // old implementation

  ClassDefOverride(RooBinomial,3) // A Binomial PDF
};
 
#endif