/*
 Binomial GLM posterior kernel for CMCMC.

 Expects:
   - DIM is the coefficient dimension (compile-time).
   - Xi[0] = n (number of observations)
   - Xi[1] = link code:
       1 = logit
       2 = probit
   - Xf layout:
       X   : n*DIM floats, row-major (i from 0..n-1, j from 0..DIM-1)
       y   : n floats, values in {0,1}
       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 ~ Bernoulli(g^{-1}(X_i theta)), and theta_j ~ N(0, var).
 */

#include <math.h>
#include <float.h>

#define BINOM_LINK_LOGIT   1
#define BINOM_LINK_PROBIT  2

static __device__ __forceinline__ float clamp_prob(float p) {
  if (p <= 0.0f) return FLT_MIN;
  if (p >= 1.0f) return 1.0f - FLT_EPSILON;
  return p;
}

static __device__ __forceinline__ float log_sigmoid(float x) {
  if (x >= 0.0f) return -log1pf(expf(-x));
  return x - log1pf(expf(x));
}

static __device__ __forceinline__ float log1m_sigmoid(float x) {
  if (x >= 0.0f) return -x - log1pf(expf(-x));
  return -log1pf(expf(x));
}

static __device__ __forceinline__ int binomial_log_terms(
    float eta, int link_code, float *logp, float *log1mp) {
  if (link_code == BINOM_LINK_LOGIT) {
    *logp = log_sigmoid(eta);
    *log1mp = log1m_sigmoid(eta);
    return 1;
  }

  if (link_code == BINOM_LINK_PROBIT) {
    const float p = clamp_prob(0.5f * erfcf(-eta * 0.70710678118654752440f));
    *logp = logf(p);
    *log1mp = log1pf(-p);
    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 logp = 0.0f;
    float log1mp = 0.0f;
    if (!binomial_log_terms(eta, link_code, &logp, &log1mp)) {
      return -INFINITY;
    }
    ll += (y[i] > 0.5f) ? logp : log1mp;
  }

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