/*
 Gaussian 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 sigma2.
   - Xi[0] = n (number of observations)
   - Xi[1] = link code, or encoded as link + 10 * prior code:
       link 1 = log
       link 2 = inverse
       link 3 = identity
       prior 1 = Uniform(sigma2_lower, sigma2_upper)
       prior 2 = Uniform(sigma_lower, sigma_upper)
       prior 3 = Exponential(rate) on sigma
       prior 4 = Half-Cauchy(scale) on sigma
   - Xf layout:
       X           : n*(DIM-1) floats, row-major
       y           : n floats, finite 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, sigma2) + log p(beta)
 + log p(sigma), with y_i ~ N(g^{-1}(X_i beta), sigma2) and
 beta_j ~ N(0, var).
 */

#include <math.h>

#define GAUSS_LINK_LOG      1
#define GAUSS_LINK_INVERSE  2
#define GAUSS_LINK_IDENTITY 3

#define GAUSS_PRIOR_UNIFORM_SIGMA2   1
#define GAUSS_PRIOR_UNIFORM_SIGMA    2
#define GAUSS_PRIOR_EXPONENTIAL_SIGMA 3
#define GAUSS_PRIOR_HALF_CAUCHY_SIGMA 4

static __device__ __forceinline__ int gaussian_mean(
    float eta, int link_code, float *mu) {
  if (link_code == GAUSS_LINK_LOG) {
    const float m = expf(eta);
    if (!isfinite(m)) return 0;
    *mu = m;
    return 1;
  }

  if (link_code == GAUSS_LINK_INVERSE) {
    if (!isfinite(eta) || eta == 0.0f) return 0;
    *mu = 1.0f / eta;
    return isfinite(*mu);
  }

  if (link_code == GAUSS_LINK_IDENTITY) {
    if (!isfinite(eta)) return 0;
    *mu = eta;
    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 : GAUSS_PRIOR_UNIFORM_SIGMA2);
  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 sigma2 = theta[beta_dim];

  const float prior_var = (prior_var_in > 0.0f ? prior_var_in : 100.0f);
  const float inv_prior_var = 1.0f / prior_var;
  if (!isfinite(sigma2) || sigma2 <= 0.0f) return -INFINITY;
  const float inv_sigma2 = 1.0f / sigma2;
  const float log_sigma2 = logf(sigma2);

  float aux_lp = 0.0f;
  if (prior_code == GAUSS_PRIOR_UNIFORM_SIGMA2) {
    const float sigma2_lower = prior_param1;
    const float sigma2_upper = prior_param2;
    if (!isfinite(sigma2_lower) || !isfinite(sigma2_upper) ||
        sigma2_lower < 0.0f || sigma2_upper <= sigma2_lower) {
      return -INFINITY;
    }
    if (sigma2 <= sigma2_lower || sigma2 >= sigma2_upper) {
      return -INFINITY;
    }
  } else if (prior_code == GAUSS_PRIOR_UNIFORM_SIGMA) {
    const float sigma_lower = prior_param1;
    const float sigma_upper = prior_param2;
    if (!isfinite(sigma_lower) || !isfinite(sigma_upper) ||
        sigma_lower < 0.0f || sigma_upper <= sigma_lower) {
      return -INFINITY;
    }
    const float sigma = sqrtf(sigma2);
    if (sigma <= sigma_lower || sigma >= sigma_upper) {
      return -INFINITY;
    }
    aux_lp += -0.5f * log_sigma2;
  } else if (prior_code == GAUSS_PRIOR_EXPONENTIAL_SIGMA) {
    const float rate = prior_param1;
    if (!isfinite(rate) || rate <= 0.0f) return -INFINITY;
    aux_lp += -rate * sqrtf(sigma2) - 0.5f * log_sigma2;
  } else if (prior_code == GAUSS_PRIOR_HALF_CAUCHY_SIGMA) {
    const float scale = prior_param1;
    if (!isfinite(scale) || scale <= 0.0f) return -INFINITY;
    aux_lp += -log1pf(sigma2 / (scale * scale)) - 0.5f * log_sigma2;
  } else {
    return -INFINITY;
  }

  float ll = 0.0f;
  for (int i = 0; i < n; i++) {
    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 (!gaussian_mean(eta, link_code, &mu)) {
      return -INFINITY;
    }
    const float r = y[i] - mu;
    ll += -0.5f * log_sigma2 - 0.5f * r * r * inv_sigma2;
  }

  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;
}
