#include "../include/rcore.h"
#include "../core/allocations.h"
#include "../include/globals.h"
#include "./hypergeometric.h"
#include "./hyperpoisson.h"
#include "linear.algebra.h"
#include "reweighted.least.squares.h"

/* clamp for the linear predictor before exponentiating / inverse-logit. */
#define ETA_MAX 30.0
/* bounds and controls for the negative-binomial dispersion Newton update. */
#define THETA_MIN 1e-4
#define THETA_MAX 1e6
#define THETA_MAX_ITER 25
#define THETA_TOL 1e-8
/* bounds for the hyper-Poisson dispersion. */
#define DISP_MIN 1e-6
#define DISP_MAX 1e6
/* ridge added to the diagonal of X'WX for numerical stability. */
#define IRLS_RIDGE 1e-8

/* numerically stable inverse-logit, clamped strictly inside (0, 1). */
double irls_inv_logit(double eta) {

double p = (eta >= 0) ? 1 / (1 + exp(-eta)) : exp(eta) / (1 + exp(eta));

  if (p < MACHINE_TOL)
    return MACHINE_TOL;
  else if (p > 1 - MACHINE_TOL)
    return 1 - MACHINE_TOL;
  else
    return p;

}/*IRLS_INV_LOGIT*/

/* linear predictor X_i . coef with the predictor clamped to [-ETA_MAX, ETA_MAX]. */
double irls_linpred(const double *X, int n, int p, const double *coef, int i) {

double eta = 0;

  for (int k = 0; k < p; k++)
    eta += X[i + k * n] * coef[k];

  if (eta > ETA_MAX)
    return ETA_MAX;
  else if (eta < -ETA_MAX)
    return -ETA_MAX;
  else
    return eta;

}/*IRLS_LINPRED*/

/* solve (X'WX) beta = X'Wz via a Cholesky factorisation. */
int c_weighted_least_squares(const double *X, int n, int p, const double *w,
    const double *z, double *beta, double *Xs, double *A, double *b) {

int info = 0;

  /* right-hand side b = X'Wz (a length-p reduction, kept as a tight loop). */
  for (int k = 0; k < p; k++) {

    double bk = 0;
    for (int i = 0; i < n; i++)
      bk += X[i + k * n] * w[i] * z[i];
    b[k] = bk;

  }/*FOR*/

  /* row-scale Xs = sqrt(w) X so that X'WX = Xs'Xs (tight loop). */
  for (int i = 0; i < n; i++) {

    double sw = sqrt(w[i]);
    for (int k = 0; k < p; k++)
      Xs[i + k * n] = sw * X[i + k * n];

  }/*FOR*/

  /* upper triangle of A = Xs'Xs via a symmetric rank-k BLAS update. */
  c_crossprod(A, Xs, n, p);

  /* a small ridge keeps the system positive definite under separation. */
  for (int k = 0; k < p; k++)
    A[k + k * p] += IRLS_RIDGE;

  /* Cholesky solve; on success the solution overwrites b, then beta. */
  info = c_chol_solve(A, b, p, 1);

  if (info == 0)
    memcpy(beta, b, p * sizeof(double));

  return info;

}/*C_WEIGHTED_LEAST_SQUARES*/

/* weighted logistic regression by IRLS. */
int c_irls_logistic(const double *X, int n, int p, const double *y,
    const double *prior_w, double *beta, int max_iter, double tol) {

double *w = Calloc1D(n, sizeof(double));
double *z = Calloc1D(n, sizeof(double));
double *beta_new = Calloc1D(p, sizeof(double));
double *Xs = Calloc1D(n * p, sizeof(double));
double *A = Calloc1D(p * p, sizeof(double));
double *b = Calloc1D(p, sizeof(double));
int iter = 0;

  for (iter = 0; iter < max_iter; iter++) {

    /* form the binomial working weights and responses at the current beta. */
    for (int i = 0; i < n; i++) {

      double eta = irls_linpred(X, n, p, beta, i);
      double pi = irls_inv_logit(eta);
      double vw = pi * (1 - pi);
      double pw = prior_w ? prior_w[i] : 1.0;

      w[i] = pw * vw;
      z[i] = eta + (y[i] - pi) / vw;

    }/*FOR*/

    /* a singular weighted system leaves the last valid beta untouched. */
    if (c_weighted_least_squares(X, n, p, w, z, beta_new, Xs, A, b) != 0)
      break;

    double delta = 0;
    for (int k = 0; k < p; k++) {

      double d = fabs(beta_new[k] - beta[k]);
      if (d > delta)
        delta = d;
      beta[k] = beta_new[k];

    }/*FOR*/

    if (delta < tol) {

      iter++;
      break;

    }/*THEN*/

  }/*FOR*/

  Free1D(w);
  Free1D(z);
  Free1D(beta_new);
  Free1D(Xs);
  Free1D(A);
  Free1D(b);

  return iter;

}/*C_IRLS_LOGISTIC*/

/* safeguarded Newton update of the negative-binomial dispersion theta on the
 * (prior-weighted) profile score, using the MASS::theta.ml score and expected
 * information; theta is recomputed at the supplied beta. */
static double nb_update_theta(const double *X, int n, int p, const double *y,
    const double *prior_w, const double *beta, double th) {

  for (int it = 0; it < THETA_MAX_ITER; it++) {

    double score = 0, info = 0;

    for (int i = 0; i < n; i++) {

      double pw = prior_w ? prior_w[i] : 1.0;
      if (pw <= 0)
        continue;

      double mu = exp(irls_linpred(X, n, p, beta, i));
      if (mu < MACHINE_TOL)
        mu = MACHINE_TOL;
      double tpm = th + mu;

      score += pw * (digamma(th + y[i]) - digamma(th) + log(th) + 1
                     - log(tpm) - (y[i] + th) / tpm);
      info  += pw * (- trigamma(th + y[i]) + trigamma(th) - 1.0 / th
                     + 2.0 / tpm - (y[i] + th) / (tpm * tpm));

    }/*FOR*/

    /* a non-positive or non-finite information makes the Newton step
     * unreliable, so stop and keep the current theta. */
    if (!R_FINITE(score) || !R_FINITE(info) || info < MACHINE_TOL)
      break;

    double th_new = th + score / info;
    if (th_new < THETA_MIN)
      th_new = THETA_MIN;
    else if (th_new > THETA_MAX)
      th_new = THETA_MAX;

    if (fabs(th_new - th) < THETA_TOL * th) {

      th = th_new;
      break;

    }/*THEN*/

    th = th_new;

  }/*FOR*/

  return th;

}/*NB_UPDATE_THETA*/

/* weighted profile score d/d(dispersion) of the hyper-Poisson count log-
 * likelihood: sum_i pw_i [ -(psi(b + y_i) - psi(b)) - dlog1F1_db(b, lambda_i) ]. */
static double hp_disp_score(const double *X, int n, int p, const double *y,
    const double *prior_w, const double *beta, double disp) {

double score = 0;

  for (int i = 0; i < n; i++) {

    double w = prior_w ? prior_w[i] : 1.0;
    if (w <= 0)
      continue;
    double lambda = exp(irls_linpred(X, n, p, beta, i));
    if (lambda < MACHINE_TOL)
      lambda = MACHINE_TOL;
    score += w * (-(digamma(disp + y[i]) - digamma(disp)) -
                  dlog1F1_db(disp, lambda));

  }/*FOR*/

  return score;

}/*HP_DISP_SCORE*/

/* one safeguarded Newton step for the hyper-Poisson dispersion on the weighted
 * profile score, using a numerical second derivative (no closed form for
 * d2/db2 log 1F1); the dispersion is kept within [DISP_MIN, DISP_MAX]. */
static double hp_update_dispersion(const double *X, int n, int p,
    const double *y, const double *prior_w, const double *beta, double disp) {

double s0 = hp_disp_score(X, n, p, y, prior_w, beta, disp);
double db = disp * 1e-4 + 1e-8;
double s1 = hp_disp_score(X, n, p, y, prior_w, beta, disp + db);
double deriv = (s1 - s0) / db;

  if (!R_FINITE(s0) || !R_FINITE(deriv) || fabs(deriv) < MACHINE_TOL)
    return disp;

  double disp_new = disp - s0 / deriv;
  if (disp_new < DISP_MIN)
    return DISP_MIN;
  else if (disp_new > DISP_MAX)
    return DISP_MAX;
  else
    return disp_new;

}/*HP_UPDATE_DISPERSION*/

/* weighted count regression by Fisher-scoring IRLS, jointly updating the
 * dispersion, for the family selected by `family`. */
int c_irls_count(glm_family_e family, const double *X, int n, int p,
    const double *y, const double *prior_w, double *beta, double *dispersion,
    int max_iter, double tol) {

double *w = Calloc1D(n, sizeof(double));
double *z = Calloc1D(n, sizeof(double));
double *beta_new = Calloc1D(p, sizeof(double));
double *Xs = Calloc1D(n * p, sizeof(double));
double *A = Calloc1D(p * p, sizeof(double));
double *b = Calloc1D(p, sizeof(double));
int iter = 0;
double disp = *dispersion;

  for (iter = 0; iter < max_iter; iter++) {

    /* form the family-specific working weights and responses (log link). */
    for (int i = 0; i < n; i++) {

      double eta = irls_linpred(X, n, p, beta, i);
      double pw = prior_w ? prior_w[i] : 1.0;
      double mu = 0, weight = 0, deriv = 0;

      if (family == GLM_NEGBIN) {

        /* mean exp(eta); dmu/deta = mu; weight = mu * theta / (theta + mu). */
        mu = exp(eta);
        if (mu < MACHINE_TOL)
          mu = MACHINE_TOL;
        weight = mu * disp / (disp + mu);
        deriv = mu;

      }/*THEN*/
      else {

        /* hyper-Poisson: dmu/deta = variance (one-parameter exponential family
         * in eta = log(lambda)), so the weight equals the variance. */
        double lambda = exp(eta), var = 0;
        if (lambda < MACHINE_TOL)
          lambda = MACHINE_TOL;
        hypois_moments(lambda, disp, &mu, &var);
        if (var < MACHINE_TOL)
          var = MACHINE_TOL;
        weight = var;
        deriv = var;

      }/*ELSE*/

      w[i] = pw * weight;
      z[i] = eta + (y[i] - mu) / deriv;

    }/*FOR*/

    if (c_weighted_least_squares(X, n, p, w, z, beta_new, Xs, A, b) != 0)
      break;

    double delta = 0;
    for (int k = 0; k < p; k++) {

      double d = fabs(beta_new[k] - beta[k]);
      if (d > delta)
        delta = d;
      beta[k] = beta_new[k];

    }/*FOR*/

    /* update the dispersion at the new coefficients. */
    if (family == GLM_NEGBIN)
      disp = nb_update_theta(X, n, p, y, prior_w, beta, disp);
    else
      disp = hp_update_dispersion(X, n, p, y, prior_w, beta, disp);

    if (delta < tol) {

      iter++;
      break;

    }/*THEN*/

  }/*FOR*/

  *dispersion = disp;

  Free1D(w);
  Free1D(z);
  Free1D(beta_new);
  Free1D(Xs);
  Free1D(A);
  Free1D(b);

  return iter;

}/*C_IRLS_COUNT*/
