// @(#)root/roostats:$Id$
// Authors: Kevin Belasco        7/22/2009
// Authors: Kyle Cranmer         7/22/2009
/*************************************************************************
 * Copyright (C) 1995-2008, Rene Brun and Fons Rademakers.               *
 * All rights reserved.                                                  *
 *                                                                       *
 * For the licensing terms see $ROOTSYS/LICENSE.                         *
 * For the list of contributors see $ROOTSYS/README/CREDITS.             *
 *************************************************************************/

#ifndef RooStats_ProposalHelper
#define RooStats_ProposalHelper

#include "Rtypes.h"
#include "RooStats/ProposalFunction.h"
#include "RooStats/UniformProposal.h"
#include "RooStats/PdfProposal.h"
#include "RooArgSet.h"
#include "RooMsgService.h"
#include "RooRealVar.h"
#include "TMatrixDSym.h"
#include "TObject.h"



namespace RooStats {

   class ProposalHelper : public TObject {

   public:
      ProposalHelper();

      /// Set the PDF to be the proposal density function
      virtual void SetPdf(RooAbsPdf& pdf) { fPdf = &pdf; }
      /// Set the bank of clues to add to the current proposal density function
      virtual void SetClues(RooDataSet& clues) { fClues = &clues; }

      /// Get the ProposalFunction that we've been designing
      virtual ProposalFunction* GetProposalFunction();

      virtual void SetCacheSize(Int_t size)
      {
         if (size > 0) {
            fCacheSize = size;
         } else {
            coutE(Eval) << "Warning: Requested non-positive cache size: " << size << ". Cache size unchanged."
                        << std::endl;
         }
      }

      virtual void SetUpdateProposalParameters(bool updateParams)
      { fUseUpdates = updateParams; }

      virtual void SetVariables(RooArgList& vars)
      { fVars = &vars; }

      virtual void SetVariables(const RooArgList& vars)
      { fVars = new RooArgList(vars); fOwnsVars = true; }

      /// set what fraction of the proposal density function should come from
      /// a uniform proposal distribution
      virtual void SetUniformFraction(double uniFrac) { fUniFrac = uniFrac; }

      /// set what fraction of the proposal density function should come from
      /// the bank of clues
      virtual void SetCluesFraction(double cluesFrac) { fCluesFrac = cluesFrac; }

      /// set the covariance matrix to use for a multi-variate Gaussian proposal
      virtual void SetCovMatrix(const TMatrixDSym& covMatrix)
      { fCovMatrix = new TMatrixDSym(covMatrix); }

      /// set what divisor we will use when dividing the range of a variable to
      /// determine the width of the proposal function for each dimension
      /// e.g. divisor = 6 for sigma = 1/6th
      virtual void SetWidthRangeDivisor(double divisor)
      { if (divisor > 0.) fSigmaRangeDivisor = divisor; }

      /// set the option string to pass to the RooNDKeysPdf constructor
      /// if the bank of clues pdf is being automatically generated by this
      /// ProposalHelper
      virtual void SetCluesOptions(const Option_t* options)
      { if (options != nullptr) fCluesOptions = options; }

      virtual void SetVariables(RooArgSet& vars)
      {
         RooArgList* argList = new RooArgList(vars);
         SetVariables(*argList);
         fOwnsVars = true;
      }

      ~ProposalHelper() override
      {
         if (fOwnsPdfProp)      delete fPdfProp;
         if (fOwnsPdf)          delete fPdf;
         if (fOwnsCluesPdf)     delete fCluesPdf;
         if (fOwnsVars)         delete fVars;
         delete fCovMatrix;
         delete fUniformPdf;
      }

   protected:
      RooAbsPdf *fPdf = nullptr;         ///< the main proposal density function
      RooAbsPdf *fCluesPdf = nullptr;    ///< proposal dens. func. with clues for certain points
      RooAbsPdf *fUniformPdf = nullptr;  ///< uniform proposal dens. func.
      RooDataSet *fClues = nullptr;      ///< data set of clues
      TMatrixDSym *fCovMatrix = nullptr; ///< covariance matrix for multi var gaussian pdf
      PdfProposal* fPdfProp = nullptr;   ///< the PdfProposal we are (probably) going to return
      RooArgList *fVars = nullptr;   ///< the RooRealVars to generate proposals for
      Int_t fCacheSize = -1;         ///< for generating proposals from PDFs
      double fSigmaRangeDivisor;   ///< range divisor to get sigma for each variable
      double fUniFrac = -1;        ///< what fraction of the PDF integral is uniform
      double fCluesFrac = -1;      ///< what fraction of the PDF integral comes from clues
      bool fOwnsPdfProp = true;    ///< whether we own the PdfProposal; equivalent to:
                                   ///< !(whether we have returned it in GetProposalFunction)
      bool fOwnsPdf = false;       ///< whether we created (and own) the main pdf
      bool fOwnsCluesPdf = false;  ///< whether we created (and own) the clues pdf
      bool fOwnsVars = false;      ///< whether we own fVars
      bool fUseUpdates = false;    ///< whether to set updates for proposal params in PdfProposal
      const Option_t *fCluesOptions = nullptr; ///< option string for clues RooNDKeysPdf

      void CreatePdf();
      void CreateCluesPdf();
      void CreateUniformPdf();
      void CreateCovMatrix(RooArgList& xVec);

      ClassDefOverride(ProposalHelper,1)
   };
}
#endif
