/*
 * Copyright (C) 1996-2011 Daniel Waggoner
 *
 * This free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * It is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * If you did not received a copy of the GNU General Public License
 * with this software, see <http://www.gnu.org/licenses/>.
 */

#include <boost/math/distributions/chi_squared.hpp>
#include <cmath>
#include <cstdio>
#include <random>
#include <sstream>

#include "dw_rand.h"

extern "C"
{
/*******************************************************************************/
/*************************** Uniform Random Numbers ****************************/
/*******************************************************************************/

static std::mt19937 gen;

/*
   Initializes seed value for uniform random number generator.  The seed value 
   can be any integer.  A value of 0 will initialize the seed from a (hopefully)
   non-deterministic random number generator.
*/
void dw_initialize_generator(int init)
{
  using seed_t = decltype(gen)::result_type;
  if (init == 0)
    {
      std::random_device rd;
      std::uniform_int_distribution<seed_t> dist {std::numeric_limits<seed_t>::min()};
      gen.seed(dist(rd));
    }
  else
    gen.seed(static_cast<seed_t>(init));
}

void dw_print_generator_state(FILE *f)
{
  std::ostringstream s;
  s << gen << "\n";
  std::fprintf(f, "%s", s.str().c_str());
}


/*******************************************************************************/
/*******************************************************************************/
/*******************************************************************************/

/*
   Generates a uniform (0,1) deviate.
*/
PRECISION dw_uniform_rnd(void)
{
  std::uniform_real_distribution<PRECISION> dist;
  return dist(gen);
}


/*
   Returns a standard gaussian deviate.  The density function for the
   standard gaussian is

                          1
                     ----------- exp(-0.5*x^2)
                      sqrt(2*Pi)

*/
PRECISION dw_gaussian_rnd(void)
{
  std::normal_distribution<PRECISION> dist;
  return dist(gen);
}

/*
   Returns a standard truncated gaussian deviate.  Usually, a <= b, but if this 
   is not the case, then the values of a and b will be swapped.

   Based on ideas in Robert's "Simulation of Truncated Normal Variables,"
   Statistics and Computing, June 1995.  (There are be eariler cites).

   If a >= 0 or b <= 0, use rejection method with an exponential proposal.  The 
   exponential family is

           g(x) = alpha*exp(-alpha*(x-a))/(1-exp(-alpha*(b-a)))

   If a >= 0, alpha is chosen to be (a+sqrt(a^2+4))/2 if this number is less than
   or equal to b and (a+b)/2 otherwise.  Emperical evidence indicates that the 
   worst case is a=0, b=infinity with the probability acceptance equal to 0.7602.
   The case b <= 0 is similar. 

   If a < 0 < b, and b - a < sqrt(2*pi), use rejection method with uniform 
   proposal.  Worst case: a=0, b=sqrt(2*pi) with probability of acceptance equal
   to 0.4939.  If a=0, then the exponential proposal would be used, but if a
   is slightly less than 0, then the acceptance probability will be slightly more
   than 0.4939.

   If a < 0 < b, and b - a >= sqrt(2*pi), draw from gaussian until the constraint 
   is satisfied.  Worst case: a=0, b=sqrt(2*pi) with probability of acceptance
   equal to 0.4939.  If a=0, then the exponential proposal would be used, but if 
   a is slightly less than 0, then the acceptance probability will be slightly 
   more than 0.4939.
*/
PRECISION dw_truncated_gaussian_rnd(PRECISION a, PRECISION b)
{
  PRECISION alpha, c, x;

  if (a == b) return a;

  if (a > b)
    {
      x=a;
      a=b;
      b=x;
    }

  if (a >= 0)
    {
      alpha=0.5*(a+sqrt(a*a+4.0));
      if (alpha >= (x=0.5*(a+b))) alpha=x;
      c=1-exp(-alpha*(b-a));
      x=a-log(1-c*dw_uniform_rnd())/alpha;
      while (dw_uniform_rnd() > exp(-0.5*(x-alpha)*(x-alpha)))
	x=a-log(1-c*dw_uniform_rnd())/alpha;
      return x;
    }

  if (b <= 0)
    {
      alpha=0.5*(-b+sqrt(b*b+4));
      if (alpha >= (x=-0.5*(a+b))) alpha=x;
      c=1-exp(-alpha*(b-a));
      x=-b-log(1-c*dw_uniform_rnd())/alpha;
      while (dw_uniform_rnd() > exp(-0.5*(x-alpha)*(x-alpha)))
	x=-b-log(1-c*dw_uniform_rnd())/alpha;
      return -x;
    }

  if (b-a < 2.506628274631000)   // sqrt(2*pi)
    {
      x=a+(b-a)*dw_uniform_rnd();
      while (dw_uniform_rnd() > exp(-x*x/2))
	x=a+(b-a)*dw_uniform_rnd();
      return x;
    }

  x=dw_gaussian_rnd();
  while ((x < a) || (x > b))
    x=dw_gaussian_rnd();
  return x;
}


/*
   Returns a standard gamma deviate.  The density function for a standard gamma
   distribution is

                                           x^(a-1)*exp(-x)
                   gamma_density(x;a) =   ----------------
                                              gamma(a)

   for a > 0.  The function gamma(a) is the integral with from 0 to infinity of 
   exp(-t)*t^(a-1).

   A general gamma variate can be obtained as follows.  Let z=b*x.  Then,
   z is drawn from a general gamma distribution whose density is

                                        z^(a-1)*exp(-z/b)
                gamma_density(z;a,b) = ------------------
                                          gamma(a)*b^a

   Notes:
    Does not check if a > 0.
*/
PRECISION dw_gamma_rnd(PRECISION a)
{
  std::gamma_distribution<PRECISION> dist {a};
  return dist(gen);
}


/*
   Returns a lognormal deviate.  The mean and standard deviations of the 
   underlying normal distributions are passed.
*/
PRECISION dw_lognormal_rnd(PRECISION mean, PRECISION standard_deviation)
{
  std::lognormal_distribution<PRECISION> dist {mean, standard_deviation};
  return dist(gen);
}



PRECISION dw_chi_square_cdf(PRECISION x, int df)
{
  boost::math::chi_squared_distribution<PRECISION> dist {static_cast<PRECISION>(df)};
  return cdf(dist, x);
}

PRECISION dw_chi_square_invcdf(PRECISION p, int df)
{
  boost::math::chi_squared_distribution<PRECISION> dist {static_cast<PRECISION>(df)};
  return quantile(dist, p);
}


/*
   Returns the natural logrithm of the gamma function applied to x.  The gamma 
   function of x is the integral from 0 to infinity of t^(x-1)*exp(-t)dt.
*/
PRECISION dw_log_gamma(PRECISION x)
{
  return std::lgamma(x);
}

} // extern "C"
