/*
 Poisson GLM posterior kernel for CMCMC.

 Expects:
   - DIM is the coefficient dimension (compile-time).
   - Xi[0] = n (number of observations)
   - Xi[1] = link code:
       1 = log
   - Xf layout:
       X   : n*DIM floats, row-major (i from 0..n-1, j from 0..DIM-1)
       y   : n floats, non-negative counts
       var : 1 float, prior variance for each coefficient

 logdens(theta) = sum_i log p(y_i | X_i, theta) + log p(theta)
 with y_i ~ Poisson(g^{-1}(X_i theta)), and theta_j ~ N(0, var).
 Terms constant in theta, such as -log(y_i!), are omitted.
 */

#include <math.h>

#define POIS_LINK_LOG      1

static __device__ __forceinline__ int poisson_log_mean(
    float eta, int link_code, float *mu, float *log_mu) {
  if (link_code == POIS_LINK_LOG) {
    const float m = expf(eta);
    if (!isfinite(m) || m <= 0.0f) return 0;
    *mu = m;
    *log_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 link_code = (Xi ? Xi[1] : 0);
  if (n <= 0) return -INFINITY;

  const float *X = Xf;
  const float *y = Xf + (size_t)n * (size_t)DIM;
  const float prior_var_in = Xf[(size_t)n * (size_t)DIM + (size_t)n];
  const float prior_var = (prior_var_in > 0.0f ? prior_var_in : 100.0f);
  const float inv_prior_var = 1.0f / prior_var;

  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)DIM;
    for (int j = 0; j < DIM; j++) {
      eta += Xi_row[j] * theta[j];
    }

    float mu = 0.0f;
    float log_mu = 0.0f;
    if (!poisson_log_mean(eta, link_code, &mu, &log_mu)) {
      return -INFINITY;
    }
    ll += y[i] * log_mu - mu;
  }

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

  return ll + lp;
}
