// -*- mode: c++; c-basic-offset: 4 -*-
// (C) Nicholas Polson, James Scott, Jesse Windle, 2012-2019

// This file is part of BayesLogit.

// BayesLogit is 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.

// BayesLogit 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.

// You should have received a copy of the GNU General Public License along with
// BayesLogit.  If not, see <https://www.gnu.org/licenses/>.


#include "PolyaGammaApproxAlt.h"
#include <stdexcept>


// double rtinvchi2(double h, double trunc, RNG& r)
// {
//   double h2 = h * h;
//   double R = trunc / h2;
//   double X = 0.0;
//   // I need to consider using a different truncated normal sampler.
//   double E1 = r.expon_rate(1.0); double E2 = r.expon_rate(1.0);
//   while ( (E1*E1) > (2 * E2 / R)) {
//     // printf("E %g %g %g %g\n", E1, E2, E1*E1, 2*E2/R);
//     E1 = r.expon_rate(1.0); E2 = r.expon_rate(1.0);
//   }
//   // printf("E %g %g \n", E1, E2);
//   X = 1 + E1 * R;
//   X = R / (X * X);
//   X = h2 * X;
//   return X;
// }


double PolyaGammaApproxAlt::a_coef(int n, double x, double h)
{
  double d_n = 2.0 * (double) n + h;
  double log_out = h * log(2.0) - lgamma(h) + lgamma(n+h) \
    - lgamma(n+1) + log(d_n)					\
    - 0.5 * log(2.0 * __PI * x * x * x) - 0.5 * d_n * d_n / x;
  double out = exp(log_out);
  return out;
}

double PolyaGammaApproxAlt::a_coef_recursive(double n, double x, double h, double coef_h, double& gnh_over_gn1_gh)
{
  double d_n = 2.0 * (double) n + h;
  // gamma_nh_over_n *= (n + h - 1) / n;  // Can speed up further by separate function for a0 and an, n > 0.
  if (n != 0)
    gnh_over_gn1_gh *= (n + h - 1) / n;
  else
    gnh_over_gn1_gh  = 1.0;
  double coef       = coef_h * gnh_over_gn1_gh;
  double log_kernel = - 0.5 * (log(x * x * x) + d_n * d_n / x) + log(d_n);
  return coef * exp(log_kernel);
}

double PolyaGammaApproxAlt::pigauss(double x, double z, double lambda)
{
  // z = 1 / mean
  double b = sqrt(lambda / x) * (x * z - 1);
  double a = sqrt(lambda / x) * (x * z + 1) * -1.0;
  double y = p_norm(b, false) + exp(2 * lambda * z) * p_norm(a, false);
  return y;
}

double PolyaGammaApproxAlt::w_left(double trunc, double h, double z)
{
  double out = 0;
  if (z != 0) 
    out = exp(h * (log(2.0) - z)) * pigauss(trunc, z/h, h*h);
  else
    out = exp(h * log(2.0)) * (1.0 - p_gamma_rate(1/trunc, 0.5, 0.5*h*h, false));
  return out;
}

double PolyaGammaApproxAlt::w_right(double trunc, double h, double z)
{
  double lambda_z = PISQ * 0.125 + 0.5 * z * z;
  double p = exp(h * log(HALFPI / lambda_z)) * (1.0-p_gamma_rate(trunc, h, lambda_z, false));
  return p;
}

double PolyaGammaApproxAlt::rtigauss(double h, double z, double trunc)
{
  z = fabs(z);
  double mu = h/z;
  double X = trunc + 1.0;
  if (mu > trunc) { // mu > t
    double alpha = 0.0;
    while (unif() > alpha) {
      X = rtinvchi2(h, trunc);
      alpha = exp(-0.5 * z*z * X);
    }
    // printf("rtigauss, part i: %g\n", X);
  }
  else {
    while (X > trunc) {
      X = igauss(mu, h*h);
    }
    // printf("rtigauss, part ii: %g\n", X);
  }
  return X;
}

double PolyaGammaApproxAlt::g_tilde(double x, double h, double trunc)
{
  double out = 0;
  if (x > trunc) 
    out = exp(h * log(0.5 * __PI) + (h-1) * log(x) - PISQ * 0.125 * x - lgamma(h));
  else 
    out = h * exp( h * log(2.0) - 0.5 * log(2.0 * __PI * x * x * x) - 0.5 * h * h / x);
    // out = h * pow(2, h) * pow(2 * __PI * pow(x,3), -0.5) * exp(-0.5 * pow(h,2) / x);
  return out;
}

////////////////////////////////////////////////////////////////////////////////
				  // Sample //
////////////////////////////////////////////////////////////////////////////////

double PolyaGammaApproxAlt::draw_abridged(double h, double z, int max_inner)
{
  if (h < 1 || h > 4) {
    #ifndef USE_R
    fprintf(stderr, "PolyaGammaApproxAlt::draw h = %g must be in [1,4]\n", h);
    #else
    Rprintf("PolyaGammaApproxAlt::draw h = %g must be in [1,4]\n", h);
    #endif
    return 0;
  }

  // Change the parameter.
  z = fabs(z) * 0.5;
  
  int    idx   = (int) floor((h-1.0)*100.0);
  double trunc = trunc_schedule[idx];

  // Now sample 0.25 * J^*(1, z := z/2).
  double rate_z       = 0.125 * __PI*__PI + 0.5 * z*z;
  double weight_left  = w_left (trunc, h, z);
  double weight_right = w_right(trunc, h, z);
  double prob_right   = weight_right / (weight_right + weight_left);

  // printf("prob_right: %g\n", prob_right);
  
  double coef1_h = exp(h * log(2.0) - 0.5 * log(2.0 * __PI));
  // double gamma_nh_over_n = RNG::Gamma(h);
  double gnh_over_gn1_gh = 1.0; // Will fill in value on first call to a_coef_recursive.

  int num_trials = 0;
  int total_iter = 0;

  while (num_trials < 10000) {
    num_trials++;

    double X = 0.0;
    double Y = 0.0;

    // if (unif() < p/(p+q))
    double uu = unif();
    if ( uu < prob_right )
      X = ltgamma(h, rate_z, trunc);
    else
      X = rtigauss(h, z, trunc);

    // double S  = a_coef(0, X, h);
    double S = a_coef_recursive(0.0, X, h, coef1_h, gnh_over_gn1_gh);
    double a_n = S;
    // double a_n2 = S2;
    // printf("a_n=%g, a_n2=%g\n", a_n, a_n2);
    double gt =  g_tilde(X, h, trunc);
    Y = unif() * gt;

    // printf("test gt: %g\n", g_tilde(trunc * 0.1, h, trunc));
    // printf("X, Y, S, gt: %g, %g, %g, %g\n", X, Y, S, gt);

    bool decreasing = false;

    int  n  = 0;
    bool go = true;
    // Cap the number of iterations?
    while (go && n < max_inner) {
      total_iter++;

      // Break infinite loop.  Put first so it always checks n==0.
      #ifdef USE_R
      if (n % 1000 == 0) R_CheckUserInterrupt();
      #endif

      ++n;
      double prev = a_n;
      // a_n  = a_coef(n, X, h);
      a_n = a_coef_recursive((double)n, X, h, coef1_h, gnh_over_gn1_gh);
      // printf("a_n=%g, a_n2=%g\n", a_n, a_n2);
      decreasing = a_n <= prev;

      if (n%2==1) {
	S = S - a_n;
	if ( Y<=S && decreasing) return 0.25 * X;
      }
      else {
	S = S + a_n;
	if ( Y>S && decreasing) go = false;
      }

    }
    // Need Y <= S in event that Y = S, e.g. when X = 0.

  }
  
  // We should never get here.
  return -1.0;
} // draw

double PolyaGammaApproxAlt::draw(double h, double z, int max_inner)
{
  if (h < 1) {
    #ifndef USE_R
    fprintf(stderr, "PolyaGammaApproxAlt::draw h = %g must be >= 1\n", h);
    #else
    Rprintf("PolyaGammaApproxAlt::draw h = %g must be >= 1\n", h);
    #endif
    return 0;
  }

  double n = floor( (h-1.0) / 4.0 );
  double remain = h - 4.0 * n;

  double x = 0.0;

  for (int i = 0; i < (int)n; i++) 
    x += draw_abridged(4.0, z);
  if (remain > 4.0)
    x += draw_abridged(0.5 * remain, z) + draw_abridged(0.5 * remain, z);
  else
    x += draw_abridged(remain, z);

  return x;
}

////////////////////////////////////////////////////////////////////////////////
				 // APPENDIX //
////////////////////////////////////////////////////////////////////////////////

// We should only have to calculate Gamma(h) once.  We can then get Gamma(n+h)
// from the recursion Gamma(z+1) = z Gamma(z).  Not sure how that is in terms of
// stability, but that should save us quite a few computations.  This affects
// a_coef and g_tilde.
