// @(#)root/unuran:$Id$
// Authors: L. Moneta, J. Leydold Wed Feb 28 2007

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

// Implementation file for class TUnuranSampler
#include "TUnuranSampler.h"

#include "TUnuranContDist.h"
#include "TUnuranDiscrDist.h"
#include "TUnuranMultiContDist.h"
#include "TUnuran.h"
#include "Math/OneDimFunctionAdapter.h"
#include "Math/DistSamplerOptions.h"
#include "Fit/DataRange.h"
//#include "Math/WrappedTF1.h"

#include "TRandom.h"
#include "TError.h"

#include "TF1.h"
#include <cassert>
#include <cmath>

ClassImp(TUnuranSampler)

TUnuranSampler::TUnuranSampler() : ROOT::Math::DistSampler(),
   fOneDim(false),
   fDiscrete(false),
   fHasMode(false), fHasArea(false),
   fMode(0), fArea(0),
   fFunc1D(0),
   fUnuran(new TUnuran()  )
{
   fLevel = ROOT::Math::DistSamplerOptions::DefaultPrintLevel();
}

TUnuranSampler::~TUnuranSampler() {
   assert(fUnuran != 0);
   delete fUnuran;
}

bool TUnuranSampler::Init(const char * algo) {
   // initialize unuran classes using the given algorithm
   assert (fUnuran != 0 );
   if (NDim() == 0)  {
      Error("TUnuranSampler::Init","Distribution function has not been set ! Need to call SetFunction first.");
      return false;
   }

   if (fLevel < 0) fLevel =  ROOT::Math::DistSamplerOptions::DefaultPrintLevel();

   TString method(algo);
   if (method.IsNull() ) {
      if (NDim() == 1) method = ROOT::Math::DistSamplerOptions::DefaultAlgorithm1D();
      else  method = ROOT::Math::DistSamplerOptions::DefaultAlgorithmND();
   }
   method.ToUpper();

   bool ret = false;
   if (NDim() == 1) {
       // check if distribution is discrete by
      // using first string in the method name is "D"
      if (method.First("D") == 0) {
         if (fLevel>1) Info("TUnuranSampler::Init","Initialize one-dim discrete distribution with method %s",method.Data());
         ret =  DoInitDiscrete1D(method);
      }
      else {
         if (fLevel>1) Info("TUnuranSampler::Init","Initialize one-dim continuous distribution with method %s",method.Data());
         ret =  DoInit1D(method);
      }
   }
   else {
      if (fLevel>1) Info("TUnuranSampler::Init","Initialize multi-dim continuous distribution with method %s",method.Data());
      ret = DoInitND(method);
   }
   // set print level in UNURAN (must be done after having initialized) -
   if (fLevel>0) {
      //fUnuran->SetLogLevel(fLevel); ( seems not to work  disable for the time being)
      if (ret) Info("TUnuranSampler::Init","Successfully initailized Unuran with method %s",method.Data() );
      else Error("TUnuranSampler::Init","Failed to  initailize Unuran with method %s",method.Data() );
      // seems not to work in UNURAN (cll only when level > 0 )
   }
   return ret;
}


bool TUnuranSampler::Init(const ROOT::Math::DistSamplerOptions & opt ) {
   // default initialization with algorithm name
   SetPrintLevel(opt.PrintLevel() );
   return Init(opt.Algorithm().c_str() );
}


bool TUnuranSampler::DoInit1D(const char * method) {
   // initilize for 1D sampling
   // need to create 1D interface from Multidim one
   // (to do: use directly 1D functions ??)
   fOneDim = true;
   TUnuranContDist * dist = 0;
   if (fFunc1D == 0) {
      ROOT::Math::OneDimMultiFunctionAdapter<> function(ParentPdf() );
      dist = new TUnuranContDist(function,0,false,true);
   }
   else {
      dist = new TUnuranContDist(*fFunc1D); // no need to copy the function
   }
   // set range in distribution (support only one range)
   const ROOT::Fit::DataRange & range = PdfRange();
   if (range.Size(0) > 0) {
      double xmin, xmax;
      range.GetRange(0,xmin,xmax);
      dist->SetDomain(xmin,xmax);
   }
   if (fHasMode) dist->SetMode(fMode);
   if (fHasArea) dist->SetPdfArea(fArea);

   bool ret = false;
   if (method) ret =  fUnuran->Init(*dist, method);
   else ret =  fUnuran->Init(*dist);
   delete dist;
   return ret;
}

bool TUnuranSampler::DoInitDiscrete1D(const char * method) {
   // initilize for 1D sampling of discrete distributions
   fOneDim = true;
   fDiscrete = true;
   TUnuranDiscrDist * dist = 0;
   if (fFunc1D == 0) {
      // need to copy the passed function pointer in this case
      ROOT::Math::OneDimMultiFunctionAdapter<> function(ParentPdf() );
      dist = new TUnuranDiscrDist(function,true);
   }
   else {
      // no need to copy the function since fFunc1D is managed outside
      dist = new TUnuranDiscrDist(*fFunc1D, false);
   }
   // set range in distribution (support only one range)
   // otherwise 0, inf is assumed
   const ROOT::Fit::DataRange & range = PdfRange();
   if (range.Size(0) > 0) {
      double xmin, xmax;
      range.GetRange(0,xmin,xmax);
      if (xmin < 0) {
         Warning("DoInitDiscrete1D","range starts from negative values - set minimum to zero");
         xmin = 0;
      }
      dist->SetDomain(int(xmin+0.1),int(xmax+0.1));
   }
   if (fHasMode) dist->SetMode(int(fMode+0.1));
   if (fHasArea) dist->SetProbSum(fArea);

   bool ret =  fUnuran->Init(*dist, method);
   delete dist;
   return ret;
}


bool TUnuranSampler::DoInitND(const char * method) {
   // initilize for 1D sampling
   TUnuranMultiContDist dist(ParentPdf());
   // set range in distribution (support only one range)
   const ROOT::Fit::DataRange & range = PdfRange();
   if (range.IsSet()) {
      std::vector<double> xmin(range.NDim() );
      std::vector<double> xmax(range.NDim() );
      range.GetRange(&xmin[0],&xmax[0]);
      dist.SetDomain(&xmin.front(),&xmax.front());
//       std::cout << " range is min = ";
//       for (int j = 0; j < NDim(); ++j) std::cout << xmin[j] << "   ";
//       std::cout << " max = ";
//       for (int j = 0; j < NDim(); ++j) std::cout << xmax[j] << "   ";
//       std::cout << std::endl;
   }
   fOneDim = false;
   if (method) return fUnuran->Init(dist, method);
   return fUnuran->Init(dist);
}

void TUnuranSampler::SetFunction(TF1 * pdf) {
   // set function from a TF1 pointer
   SetFunction<TF1>(*pdf, pdf->GetNdim());
}

void TUnuranSampler::SetRandom(TRandom * r) {
   // set random generator (must be called before Init to have effect)
   fUnuran->SetRandom(r);
}

void TUnuranSampler::SetSeed(unsigned int seed) {
   // set random generator seed (must be called before Init to have effect)
   fUnuran->SetSeed(seed);
}

TRandom * TUnuranSampler::GetRandom() {
   // get random generator used
   return  fUnuran->GetRandom();
}

double TUnuranSampler::Sample1D() {
   // sample 1D distributions
   return (fDiscrete) ? (double) fUnuran->SampleDiscr() : fUnuran->Sample();
}

bool TUnuranSampler::Sample(double * x) {
   // sample multi-dim distributions
   if (!fOneDim) return fUnuran->SampleMulti(x);
   x[0] = Sample1D();
   return true;
}


bool TUnuranSampler::SampleBin(double prob, double & value, double *error) {
   // sample a bin according to Poisson statistics

   TRandom * r = fUnuran->GetRandom();
   if (!r) return false;
   value = r->Poisson(prob);
   if (error) *error = std::sqrt(value);
   return true;
}