/*
 Gamma GLM posterior kernel for CMCMC.

 Expects:
   - DIM is the full parameter dimension (compile-time).
     theta[0:(DIM-2)] are coefficients and theta[DIM-1] is the
     Gamma shape parameter alpha.
   - Xi[0] = n (number of observations)
   - Xi[1] = link code, or encoded as link + 10 * prior code:
       link 1 = log
       prior 1 = Uniform(alpha_lower, alpha_upper)
       prior 2 = Uniform(cv_lower, cv_upper), CV=1/sqrt(alpha)
       prior 3 = Exponential(rate) on CV
       prior 4 = Half-Cauchy(scale) on CV
       prior 5 = Exponential(rate) on alpha
   - Xf layout:
       X           : n*(DIM-1) floats, row-major
       y           : n floats, positive responses
       var         : 1 float, prior variance for each coefficient
       prior_param1: 1 float, interpretation depends on prior code
       prior_param2: 1 float, interpretation depends on prior code

 logdens(theta) = sum_i log p(y_i | X_i, beta, alpha) + log p(beta)
 with E[y_i] = g^{-1}(X_i beta), Var[y_i] = mu_i^2 / alpha,
 beta_j ~ N(0, var), and p(alpha) determined by the prior code.
 */

#include <math.h>

#define GAMMA_LINK_LOG      1

#define GAMMA_PRIOR_UNIFORM_ALPHA  1
#define GAMMA_PRIOR_UNIFORM_CV     2
#define GAMMA_PRIOR_EXP_CV         3
#define GAMMA_PRIOR_HALF_CAUCHY_CV 4
#define GAMMA_PRIOR_EXP_ALPHA      5

static __device__ __forceinline__ int gamma_mean(
    float eta, int link_code, float *mu) {
  if (link_code == GAMMA_LINK_LOG) {
    const float m = expf(eta);
    if (!isfinite(m) || m <= 0.0f) return 0;
    *mu = m;
    return 1;
  }

  return 0;
}

__device__ float logdens(const float *theta, const float *Xf, const int *Xi) {
  const int n = (Xi ? Xi[0] : 0);
  const int raw_link_code = (Xi ? Xi[1] : 0);
  const int prior_code = (raw_link_code >= 10 ? raw_link_code / 10 : GAMMA_PRIOR_UNIFORM_ALPHA);
  const int link_code = (raw_link_code >= 10 ? raw_link_code - 10 * prior_code : raw_link_code);
  if (n <= 0 || DIM <= 1) return -INFINITY;
  const int beta_dim = DIM - 1;

  const float *X = Xf;
  const float *y = Xf + (size_t)n * (size_t)beta_dim;
  const size_t scalar_offset = (size_t)n * (size_t)beta_dim + (size_t)n;
  const float prior_var_in = Xf[scalar_offset];
  const float prior_param1 = Xf[scalar_offset + 1u];
  const float prior_param2 = Xf[scalar_offset + 2u];
  const float alpha = theta[beta_dim];
  if (!isfinite(alpha) || alpha <= 0.0f) return -INFINITY;

  const float prior_var = (prior_var_in > 0.0f ? prior_var_in : 100.0f);
  const float inv_prior_var = 1.0f / prior_var;
  const float log_alpha = logf(alpha);
  const float log_gamma_alpha = lgammaf(alpha);

  float aux_lp = 0.0f;
  if (prior_code == GAMMA_PRIOR_UNIFORM_ALPHA) {
    const float alpha_lower = prior_param1;
    const float alpha_upper = prior_param2;
    if (!isfinite(alpha_lower) || !isfinite(alpha_upper) ||
        alpha_lower < 0.0f || alpha_upper <= alpha_lower) {
      return -INFINITY;
    }
    if (alpha <= alpha_lower || alpha >= alpha_upper) {
      return -INFINITY;
    }
  } else if (prior_code == GAMMA_PRIOR_UNIFORM_CV) {
    const float cv_lower = prior_param1;
    const float cv_upper = prior_param2;
    if (!isfinite(cv_lower) || !isfinite(cv_upper) ||
        cv_lower < 0.0f || cv_upper <= cv_lower) {
      return -INFINITY;
    }
    const float cv = 1.0f / sqrtf(alpha);
    if (cv <= cv_lower || cv >= cv_upper) {
      return -INFINITY;
    }
    aux_lp += -1.5f * log_alpha;
  } else if (prior_code == GAMMA_PRIOR_EXP_CV) {
    const float rate = prior_param1;
    if (!isfinite(rate) || rate <= 0.0f) return -INFINITY;
    aux_lp += -rate / sqrtf(alpha) - 1.5f * log_alpha;
  } else if (prior_code == GAMMA_PRIOR_HALF_CAUCHY_CV) {
    const float scale = prior_param1;
    if (!isfinite(scale) || scale <= 0.0f) return -INFINITY;
    const float cv2 = 1.0f / alpha;
    aux_lp += -log1pf(cv2 / (scale * scale)) - 1.5f * log_alpha;
  } else if (prior_code == GAMMA_PRIOR_EXP_ALPHA) {
    const float rate = prior_param1;
    if (!isfinite(rate) || rate <= 0.0f) return -INFINITY;
    aux_lp += -rate * alpha;
  } else {
    return -INFINITY;
  }

  float ll = 0.0f;
  for (int i = 0; i < n; i++) {
    if (!isfinite(y[i]) || y[i] <= 0.0f) return -INFINITY;

    float eta = 0.0f;
    const float *Xi_row = X + (size_t)i * (size_t)beta_dim;
    for (int j = 0; j < beta_dim; j++) {
      eta += Xi_row[j] * theta[j];
    }

    float mu = 0.0f;
    if (!gamma_mean(eta, link_code, &mu)) {
      return -INFINITY;
    }
    ll += alpha * log_alpha - alpha * logf(mu) - log_gamma_alpha +
          (alpha - 1.0f) * logf(y[i]) - alpha * y[i] / mu;
  }

  float lp = 0.0f;
  for (int j = 0; j < beta_dim; j++) {
    const float tj = theta[j];
    lp += -0.5f * tj * tj * inv_prior_var;
  }

  return ll + lp + aux_lp;
}
