#include "../../include/rcore.h"
#include "../../core/allocations.h"
#include "../../core/sets.h"
#include "../../scores/scores.h"
#include "../../tests/tests.h"

/* unconditional mutual information, to be used in C code; when missing is set
 * incomplete cases (a missing value in either variable) are skipped and the
 * number of complete cases is returned in ncomplete. */
double c_micg(double *yy, int *xx, int llx, int num, double *df, int *ncomplete,
    bool missing) {

int i = 0, nc = 0, *ni = NULL;
double lognum = 0, logden = 0, res = 0, ym = 0, ysd = 0, *mu = NULL, *sd = NULL;
long double sum = 0, sse = 0;

  mu = Calloc1D(llx, sizeof(double));
  sd = Calloc1D(llx, sizeof(double));
  ni = Calloc1D(llx, sizeof(int));

  /* marginal and per-level means and counts; the marginal mean accumulates in
   * long double to match c_mean() on complete data. */
  for (i = 0; i < num; i++) {

    if (missing && (ISNAN(yy[i]) || (xx[i] == NA_INTEGER)))
      continue;

    sum += yy[i];
    mu[xx[i] - 1] += yy[i];
    ni[xx[i] - 1]++;
    nc++;

  }/*FOR*/

  /* with no complete data points, or just one, return independence. */
  if (nc <= 1)
    goto free_and_return;

  ym = (double)(sum / nc);
  for (i = 0; i < llx; i++)
    mu[i] /= ni[i];

  /* marginal and per-level variances (matching c_sse() on complete data). */
  for (i = 0; i < num; i++) {

    if (missing && (ISNAN(yy[i]) || (xx[i] == NA_INTEGER)))
      continue;

    sse += (yy[i] - ym) * (yy[i] - ym);
    sd[xx[i] - 1] += (yy[i] - mu[xx[i] - 1]) * (yy[i] - mu[xx[i] - 1]);

  }/*FOR*/

  ysd = sqrt((double)sse / (nc - 1));
  for (i = 0; i < llx; i++) {

    if (ni[i] == 0)
      sd[i] = NA_REAL;
    else if (ni[i] == 1)
      sd[i] = 0;
    else
      sd[i] = sqrt(sd[i] / (ni[i] - 1));

  }/*FOR*/

  /* the denominator (null) and numerator (alternative) log-likelihoods. */
  for (i = 0; i < num; i++) {

    if (missing && (ISNAN(yy[i]) || (xx[i] == NA_INTEGER)))
      continue;

    logden += dnorm(yy[i], ym, ysd, TRUE);
    lognum += dnorm(yy[i], mu[xx[i] - 1], sd[xx[i] - 1], TRUE);

  }/*FOR*/

  /* a level with a single observation gives a zero variance and an infinite
   * density, so guard against a singular alternative model as in c_cmicg(). */
  res = (R_FINITE(logden) && R_FINITE(lognum)) ? (lognum - logden) / nc : 0;

free_and_return:
  if (ncomplete)
    *ncomplete = nc;

  /* the discrete variable is a parent of the continuous variable, then the
   * number of parameters is equal to one mean and one variance for each
   * level of the discrete variable; and the degrees of freedom of the test
   * are smaller by the one mean and variance under the null. */
  if (df)
    *df = 2 * (llx - 1);

  Free1D(mu);
  Free1D(sd);
  Free1D(ni);

  return res;

}/*C_MICG*/

/* conditional mutual information, to be used in C code. */
double c_cmicg(double *yy, double **xx, int nx, int **zz, int nz, int *z0,
    int nz0, int *nlvls, int num, double *df) {

double logden = 0, lognum = 0;
int *z1 = NULL, nz1 = 0;

  /* CONTRACT (zz != NULL branch): the alternative uses config(zz) against the
   * null's config(z0), so zz MUST contain the base z0 (config(zz) nests over
   * config(z0)); zz[0] is the newly added discrete variable (used for the df).
   * Passing non-nested configurations makes the statistic wrong and asymmetric. */

  if (!zz) {

    /* compute the denominator (model under the null). */
    logden = c_fast_ccgloglik(yy, xx + 1, nx - 1, num, z0, nz0);
    /* compute the numerator (model under the alternative). */
    lognum = c_fast_ccgloglik(yy, xx, nx, num, z0, nz0);

    /* one regression coefficient for each conditioning level is added
     * (and there are no new standard errors since there are no new discrete
     * parents); if all conditioning variables are continuous that's just
     * one global regression coefficient. */
    if (df)
      *df = (nz0 == 0) ? 1 : nz0;

  }/*THEN*/
  else {

    /* compute the denominator (model under the null). */
    logden = c_fast_ccgloglik(yy, xx, nx, num, z0, nz0);
    /* compute the numerator (model under the alternative). */
    z1 = Calloc1D(num, sizeof(int));
    c_fast_config(zz, num, nz, nlvls, z1, &nz1, 1);
    lognum = c_fast_ccgloglik(yy, xx, nx, num, z1, nz1);
    Free1D(z1);

    /* for each additional configuration of the discrete conditioning
     * variables plus the discrete yptr, one whole set of regression
     * coefficients (plus the intercept and the standard error) is added.
     * under H0: llz * (ngp + 2) vs under H1: (lly * llz) * (ngp + 2). */
    if (df)
      *df = (nlvls[0] - 1) * ((nz0 == 0) ? 1 : nz0) * (nx + 2);

  }/*ELSE*/

  /* if the null model is singular, the alternative model is even more singular
   * so it should always be rejected. */
  return (R_FINITE(logden) && R_FINITE(lognum)) ? (lognum - logden) / num : 0;

}/*C_CMICG*/

/* conditional mutual information between two discrete variables, conditional
 * on both discrete and continuous variables. */
double c_cmicg_unroll(int *xx, int llx, int *yy, int lly, int *zz, int llz,
    double **gp, int ngp, double *df, int num) {

int i = 0, llz2 = 0, tlvls[2] = {0, 0};
int *tt[2] = {NULL, NULL}, *zz2 = NULL;
int *dcols[2] = {NULL, NULL}, dlvls[2] = {0, 0};
double logden = 0, lognum = 0;

  if (!zz) {

    /* no discrete conditioning variables, reuse yy the denominator. */
    zz2 = yy;
    llz2 = lly;

    /* remainder term: the mutual information test of xx and yy. */
    lognum = c_chisqtest(xx, llx, yy, lly, num, df, MI, FALSE);

  }/*THEN*/
  else {

    /* combine zz and yy to get the configurations for the denominator. */
    zz2 = Calloc1D(num, sizeof(int));
    tt[0] = yy;
    tt[1] = zz;
    tlvls[0] = lly;
    tlvls[1] = llz;
    c_fast_config(tt, num, 2, tlvls, zz2, &llz2, 1);

    /* remainder term: the mutual information test of xx and yy given zz. */
    lognum = c_cchisqtest(xx, llx, yy, lly, zz, llz, num, df, MI, FALSE);

  }/*ELSE*/

  /* iterate over the continuous variables using the chain rule. adding xx to
   * the conditioning set is a properly nested comparison: the numerator adds
   * it on top of {yy, zz} giving I(xx; gp | yy, zz), the denominator on top of
   * {zz} giving I(xx; gp | zz). the two conditioning sets must therefore
   * include zz2 and zz, not just xx (otherwise the models are not nested and
   * the statistic is not symmetric in xx and yy). */
  dcols[0] = xx;
  dcols[1] = zz2;
  dlvls[0] = llx;
  dlvls[1] = llz2;
  for (i = 0; i < ngp; i++)
    lognum += c_cmicg(gp[i], gp + i + 1, ngp - i - 1, dcols, 2, zz2, llz2,
                dlvls, num, NULL);

  dcols[1] = zz;
  dlvls[1] = llz;
  for (i = 0; i < ngp; i++)
    logden += c_cmicg(gp[i], gp + i + 1, ngp - i - 1, dcols, zz ? 2 : 1, zz,
                zz ? llz : 0, dlvls, num, NULL);

  if (zz)
    Free1D(zz2);

  /* set the degrees of freedom: for each continuous variable added by the
   * chain rule, the interaction between the two discrete variables, which is
   * symmetric in their numbers of levels. */
  if (df)
    *df += ((llz < 1) ? 1 : llz) * (ngp * (ngp + 3) / 2) * (llx - 1) * (lly - 1);

  /* lognum and logden are already per-observation quantities (unlike c_cmicg(),
   * which sums log-likelihoods), so return the mutual information directly; the
   * caller scales it by 2 * nobs like the other branches. */
  return (R_FINITE(logden) && R_FINITE(lognum)) ? (lognum - logden) : 0;

}/*C_CMICG_UNROLL*/

