// @(#)root/unuran:$Id$
// Authors: L. Moneta, J. Leydold Tue Sep 26 16:25:09 2006

/**********************************************************************
 *                                                                    *
 * Copyright (c) 2006  LCG ROOT Math Team, CERN/PH-SFT                *
 *                                                                    *
 *                                                                    *
 **********************************************************************/

// Implementation file for class TUnuran

#include "TUnuran.h"

#include "TUnuranContDist.h"
#include "TUnuranMultiContDist.h"
#include "TUnuranDiscrDist.h"
#include "TUnuranEmpDist.h"

#include "UnuranRng.h"
#include "UnuranDistrAdapter.h"

#include "TRandom.h"
#include "TSystem.h"

#include "TH1.h"

#include <cassert>


#include <unuran.h>

#include "TError.h"


TUnuran::TUnuran(TRandom * r, unsigned int debugLevel) : 
   fGen(0),
   fUdistr(0),
   fUrng(0),
   fRng(r) 
{
   // constructor implementation with a ROOT random generator
   // if no generator is given the ROOT default is used
   if (fRng == 0) fRng = gRandom; 
   // set debug level at global level 
   // (should be in a static  initialization function of the library ? )
   if ( debugLevel > 1) 
      unur_set_default_debug(UNUR_DEBUG_ALL);
   else if (debugLevel == 1) 
      unur_set_default_debug(UNUR_DEBUG_INIT);
   else
      unur_set_default_debug(UNUR_DEBUG_OFF);
      
}


TUnuran::~TUnuran() 
{
   // Destructor implementation
   if (fGen != 0) unur_free(fGen); 
   if (fUrng != 0) unur_urng_free(fUrng);
  // we can delete now the distribution object
   if (fUdistr != 0) unur_distr_free(fUdistr);
}

//private (no impl.)
TUnuran::TUnuran(const TUnuran &) 
{
   // Implementation of copy constructor.
}

TUnuran & TUnuran::operator = (const TUnuran &rhs) 
{
   // Implementation of assignment operator.
   if (this == &rhs) return *this;  // time saving self-test
   return *this;
}

bool  TUnuran::Init(const std::string & dist, const std::string & method) 
{
   // initialize with a string
   std::string s = dist + " & " + method; 
   fGen = unur_str2gen(s.c_str() ); 
   if (fGen == 0) { 
      Error("Init","Cannot create generator object"); 
      return false; 
   } 
   if (! SetRandomGenerator() ) return false; 

   return true; 
}

bool TUnuran::Init(const TUnuranContDist & distr, const std::string  & method)  
{ 
   // initialization with a distribution and and generator
   // the distribution object is copied in and managed by this class
   // use auto_ptr to manage previously existing distribution objects
   TUnuranContDist * distNew = distr.Clone(); 
   fDist = std::auto_ptr< TUnuranBaseDist>(distNew);

   fMethod = method; 
   if (! SetContDistribution(*distNew) ) return false;
   if (! SetMethodAndInit() ) return false;
   if (! SetRandomGenerator() ) return false; 
   return true;
}

    
bool TUnuran::Init(const TUnuranMultiContDist & distr, const std::string  & method)  
{ 
   //  initialization with a distribution and method
   // the distribution object is copied in and managed by this class
   // use auto_ptr to manage previously existing distribution objects
   TUnuranMultiContDist * distNew = distr.Clone(); 
   fDist = std::auto_ptr< TUnuranBaseDist>(distNew);

   fMethod = method;
   if (! SetMultiDistribution(*distNew) ) return false;
   if (! SetMethodAndInit() ) return false;
   if (! SetRandomGenerator() ) return false; 
   return true; 
}


bool TUnuran::Init(const TUnuranDiscrDist & distr, const std::string & method ) {
   //   initialization with a distribution and and generator
   // the distribution object is copied in and managed by this class
   // use auto_ptr to manage previously existing distribution objects
   TUnuranDiscrDist * distNew = distr.Clone(); 
   fDist = std::auto_ptr< TUnuranBaseDist>(distNew);

   fMethod = method; 
   if (! SetDiscreteDistribution(*distNew) ) return false;
   if (! SetMethodAndInit() ) return false;
   if (! SetRandomGenerator() ) return false; 
   return true;
}

bool TUnuran::Init(const TUnuranEmpDist & distr, const std::string & method ) {
   //   initialization with a distribution and and generator
   // the distribution object is copied in and managed by this class
   // use auto_ptr to manage previously existing distribution objects
   TUnuranEmpDist * distNew = distr.Clone(); 
   fDist = std::auto_ptr< TUnuranBaseDist>(distNew);
   
   fMethod = method; 
   if (distr.IsBinned()) fMethod = "hist";
   else if (distr.NDim() > 1) fMethod = "vempk";
   if (! SetEmpiricalDistribution(*distNew) ) return false;
   if (! SetMethodAndInit() ) return false;
   if (! SetRandomGenerator() ) return false; 
   return true;
}


bool  TUnuran::SetRandomGenerator()
{
   // set an external random generator
   if (fRng == 0) return false; 
   if (fGen == 0) return false; 

   fUrng = unur_urng_new(&UnuranRng<TRandom>::Rndm, fRng );
   if (fUrng == 0) return false; 
   unsigned int ret = 0; 
   ret |= unur_urng_set_delete(fUrng, &UnuranRng<TRandom>::Delete); 
   ret |= unur_urng_set_seed(fUrng, &UnuranRng<TRandom>::Seed);
   if (ret != 0) return false; 

   unur_chg_urng( fGen, fUrng); 
   return true;    
}

bool  TUnuran::SetContDistribution(const TUnuranContDist & dist )
{
   // internal method to set in unuran the function pointer for a continuous univariate distribution 
   if (fUdistr != 0)  unur_distr_free(fUdistr);
   fUdistr = unur_distr_cont_new(); 
   if (fUdistr == 0) return false; 
   unsigned int ret = 0; 
   ret = unur_distr_set_extobj(fUdistr, &dist);  
   if ( ! dist.IsLogPdf() ) { 
      ret |= unur_distr_cont_set_pdf(fUdistr, &ContDist::Pdf);  
      ret |= unur_distr_cont_set_dpdf(fUdistr, &ContDist::Dpdf);  
      if (dist.HasCdf() ) ret |= unur_distr_cont_set_cdf(fUdistr, &ContDist::Cdf);  
   }
   else { 
      // case user provides log of pdf 
      ret |= unur_distr_cont_set_logpdf(fUdistr, &ContDist::Pdf);  
      ret |= unur_distr_cont_set_dlogpdf(fUdistr, &ContDist::Dpdf);  
   }

   double xmin, xmax = 0; 
   if (dist.GetDomain(xmin,xmax) ) { 
      ret = unur_distr_cont_set_domain(fUdistr,xmin,xmax);  
      if (ret != 0)  { 
         Error("SetContDistribution","invalid domain xmin = %g xmax = %g ",xmin,xmax);
         return false; 
      }
   }
   if (dist.HasMode() ) { 
      ret = unur_distr_cont_set_mode(fUdistr, dist.Mode());  
      if (ret != 0)  { 
         Error("SetContDistribution","invalid mode given,  mode = %g ",dist.Mode());
         return false; 
      }
   }
   if (dist.HasPdfArea() ) { 
      ret = unur_distr_cont_set_pdfarea(fUdistr, dist.PdfArea());  
      if (ret != 0)  { 
         Error("SetContDistribution","invalid area given,  area = %g ",dist.PdfArea());
         return false; 
      }
   }

   return (ret ==0) ? true : false; 
}


bool  TUnuran::SetMultiDistribution(const TUnuranMultiContDist & dist )
{
   // internal method to set in unuran the function pointer for a multivariate distribution 
   if (fUdistr != 0)  unur_distr_free(fUdistr);
   fUdistr = unur_distr_cvec_new(dist.NDim() ); 
   if (fUdistr == 0) return false; 
   unsigned int ret = 0; 
   ret |= unur_distr_set_extobj(fUdistr, &dist );  
   if ( ! dist.IsLogPdf() ) { 
      ret |= unur_distr_cvec_set_pdf(fUdistr, &MultiDist::Pdf);  
      ret |= unur_distr_cvec_set_dpdf(fUdistr, &MultiDist::Dpdf);  
      ret |= unur_distr_cvec_set_pdpdf(fUdistr, &MultiDist::Pdpdf);  
   }
   else { 
      ret |= unur_distr_cvec_set_logpdf(fUdistr, &MultiDist::Pdf);  
      ret |= unur_distr_cvec_set_dlogpdf(fUdistr, &MultiDist::Dpdf);  
      ret |= unur_distr_cvec_set_pdlogpdf(fUdistr, &MultiDist::Pdpdf);  
   }

   const double * xmin = dist.GetLowerDomain();
   const double * xmax = dist.GetUpperDomain();
   if ( xmin != 0 || xmax != 0 ) {
      ret = unur_distr_cvec_set_domain_rect(fUdistr,xmin,xmax);  
      if (ret != 0)  { 
         Error("SetMultiDistribution","invalid domain");
         return false; 
      }
#ifdef OLDVERS
      Error("SetMultiDistribution","domain setting not available in UNURAN 0.8.1");
#endif

   }

   const double * xmode = dist.GetMode(); 
   if (xmode != 0) { 
      ret = unur_distr_cvec_set_mode(fUdistr, xmode);  
      if (ret != 0)  { 
         Error("SetMultiDistribution","invalid mode");
         return false; 
      }
   }
   return (ret ==0) ? true : false; 
}

bool TUnuran::SetEmpiricalDistribution(const TUnuranEmpDist & dist) { 

   // internal method to set in unuran the function pointer for am empiral distribution (from histogram)
   if (fUdistr != 0)  unur_distr_free(fUdistr);
   if (dist.NDim() == 1) 
      fUdistr = unur_distr_cemp_new(); 
   else 
      fUdistr = unur_distr_cvemp_new(dist.NDim() ); 

   if (fUdistr == 0) return false; 
   unsigned int ret = 0; 


   // get info from histogram 
   if (dist.IsBinned() ) { 
      int nbins = dist.Data().size(); 
      double min = dist.LowerBin();
      double max = dist.UpperBin();
      const double * pv = &(dist.Data().front());
      ret |= unur_distr_cemp_set_hist(fUdistr, pv, nbins, min, max); 
#ifdef OLDVERS
      Error("SetEmpiricalDistribution","hist method not available in UNURAN 0.8.1");
#endif
   } 
   else { 
      const double * pv = &dist.Data().front();
      // n is number of points (size/ndim)
      int n = dist.Data().size()/dist.NDim();
      if (dist.NDim() == 1)   
         ret |= unur_distr_cemp_set_data(fUdistr, pv, n); 
      else  
         ret |= unur_distr_cvemp_set_data(fUdistr, pv, n); 
   }
   if (ret != 0) { 
      Error("SetEmpiricalDistribution","invalid distribution object");
      return false; 
   }
   return true; 
}


bool  TUnuran::SetDiscreteDistribution(const TUnuranDiscrDist & dist)
{
   // internal method to set in unuran the function pointer for a discrete univariate distribution 
   if (fUdistr != 0)  unur_distr_free(fUdistr);
   fUdistr = unur_distr_discr_new(); 
   if (fUdistr == 0) return false; 
   unsigned int ret = 0; 
   // if a probability mesh function is provided 
   if (dist.ProbVec().size() == 0) { 
      ret = unur_distr_set_extobj(fUdistr, &dist );  
      ret |= unur_distr_discr_set_pmf(fUdistr, &DiscrDist::Pmf);  
      if (dist.HasCdf() ) ret |= unur_distr_discr_set_cdf(fUdistr, &DiscrDist::Cdf);  
 
   }
   else { 
      // case user provides vector of probabilities 
      ret |= unur_distr_discr_set_pv(fUdistr, &dist.ProbVec().front(), dist.ProbVec().size() );
   }

   int xmin, xmax = 0; 
   if (dist.GetDomain(xmin,xmax) ) { 
      ret = unur_distr_discr_set_domain(fUdistr,xmin,xmax);  
      if (ret != 0)  { 
         Error("SetDiscrDistribution","invalid domain xmin = %d xmax = %d ",xmin,xmax);
         return false; 
      }
   }
   if (dist.HasMode() ) { 
      ret = unur_distr_discr_set_mode(fUdistr, dist.Mode());  
      if (ret != 0)  { 
         Error("SetContDistribution","invalid mode given,  mode = %d ",dist.Mode());
         return false; 
      }
   }
   if (dist.HasProbSum() ) {
      ret = unur_distr_discr_set_pmfsum(fUdistr, dist.ProbSum());  
      if (ret != 0)  { 
         Error("SetContDistribution","invalid sum given,  mode = %g ",dist.ProbSum());
         return false; 
      }
   }

   return (ret ==0) ? true : false; 
}


//bool TUnuran::SetMethodAndInit(const std::string & s) { 
bool TUnuran::SetMethodAndInit() { 

   // internal function to set a method from a distribution and 
   // initialize unuran with the given distribution.
   if (fUdistr == 0) return false; 

   struct unur_slist *mlist = NULL;
   
   UNUR_PAR * par = _unur_str2par(fUdistr, fMethod.c_str(), &mlist);
   if (par == 0) {
      Error("SetMethod","missing distribution information or syntax error");
      if (mlist != 0)  _unur_slist_free(mlist);
      return false;
   }


   // set unuran to not use a private copy of the distribution object
   unur_set_use_distr_privatecopy (par, false);

   // need to free fGen if already existing ? 
   if (fGen != 0 )  unur_free(fGen);
   fGen = unur_init(par); 
   _unur_slist_free(mlist);
   if (fGen == 0) { 
      Error("SetMethod","initializing Unuran: condition for method violated");
      return false; 
   }
   return true;
 }


int TUnuran::SampleDiscr()
{
   // sample one-dimensional distribution
   assert(fGen != 0); 
   return unur_sample_discr(fGen);
}

double TUnuran::Sample()
{
   // sample one-dimensional distribution
   assert(fGen != 0); 
   return unur_sample_cont(fGen);
}

bool TUnuran::SampleMulti(double * x)
{
   // sample multidimensional distribution
   if (fGen == 0) return false;  
   unur_sample_vec(fGen,x);
   return true; 
}

void TUnuran::SetSeed(unsigned int seed) { 
   return fRng->SetSeed(seed); 
}

bool  TUnuran::SetLogLevel(unsigned int debugLevel) 
{
   if (fGen == 0) return false; 
   int ret = 0; 
   if ( debugLevel > 1) 
      ret |= unur_chg_debug(fGen, UNUR_DEBUG_ALL);
   else if (debugLevel == 1) 
      ret |= unur_chg_debug(fGen, UNUR_DEBUG_ALL);
   else
      ret |= unur_chg_debug(fGen, UNUR_DEBUG_OFF);

   return (ret ==0) ? true : false; 

}

bool TUnuran::InitPoisson(double mu, const std::string & method) { 
   // initializaton for a Poisson
   double p[1];
   p[0] = mu; 

   fUdistr = unur_distr_poisson(p,1);

   fMethod = method;
   if (fUdistr == 0) return false; 
   if (! SetMethodAndInit() ) return false;
   if (! SetRandomGenerator() ) return false; 
   return true; 
}


bool TUnuran::InitBinomial(unsigned int ntot, double prob, const std::string & method ) { 
   // initializaton for a Binomial
   double par[2];
   par[0] = ntot; 
   par[1] = prob; 
   fUdistr = unur_distr_binomial(par,2);

   fMethod = method;
   if (fUdistr == 0) return false; 
   if (! SetMethodAndInit() ) return false;
   if (! SetRandomGenerator() ) return false; 
   return true; 
}


bool TUnuran::ReInitDiscrDist(unsigned int npar, double * par) { 
   // re-initialization of UNURAN without freeing and creating a new fGen object
   // works only for pre-defined distribution by changing their parameters 
   if (!fGen ) return false;
   if (!fUdistr) return false;
   unur_distr_discr_set_pmfparams(fUdistr,par,npar);
   int iret = unur_reinit(fGen);
   if (iret) Warning("ReInitDiscrDist","re-init failed - a full initizialization must be performed");
   return (!iret); 
}