#include "../../include/rcore.h"
#include "../../core/allocations.h"
#include "../../core/correlation.h"
#include "../../core/covariance.matrix.h"
#include "../../core/data.table.h"
#include "../../core/moments.h"
#include "../../core/sets.h"
#include "../tests.h"

/* parametric tests for discrete variables. */
double ct_discrete(tabular dtx, tabular dty, tabular dtz, double *pvalue, double *df,
    test_e test) {

int i = 0, llx = 0, lly = dty.nlvl[0], llz = 0;
int *xptr = NULL, *yptr = dty.dcol[0], *zptr = NULL;
double statistic = 0;

  /* if y or any conditioning variable is latent, no candidate has complete
   * observations: they are all independent. */
  if (dty.m.flag[0].latent || tabular_has_latent(dtz)) {

    for (i = 0; i < dtx.m.ncols; i++)
      pvalue[i] = 1;
    return statistic;

  }/*THEN*/

  zptr = Calloc1D(dtz.m.nobs, sizeof(int));
  c_fast_config(dtz.dcol, dtz.m.nobs, dtz.m.ncols, dtz.nlvl, zptr, &llz, 1);

  for (i = 0; i < dtx.m.ncols; i++) {

    /* a latent x is independent of y. */
    if (dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    xptr = dtx.dcol[i];
    llx = dtx.nlvl[i];

    if (test == MI || test == MI_ADF || test == X2 || test == X2_ADF) {

      /* mutual information and Pearson's X^2 asymptotic tests. */
      statistic = c_cchisqtest(xptr, llx, yptr, lly, zptr, llz, dtz.m.nobs,
                    df, test, (test == MI) || (test == MI_ADF));
      pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

    }/*THEN*/
    else if (test == MI_SH) {

      /* shrinkage mutual information test. */
      statistic = c_shcmi(xptr, llx, yptr, lly, zptr, llz, dtz.m.nobs, df, TRUE);
      pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

    }/*THEN*/
    else if (test == JT) {

      /* Jonckheere-Terpstra test. */
      statistic = c_cjt(xptr, llx, yptr, lly, zptr, llz, dtz.m.nobs);
      pvalue[i] = 2 * pnorm(fabs(statistic), 0, 1, FALSE, FALSE);

    }/*THEN*/

  }/*FOR*/

  Free1D(zptr);

  return statistic;

}/*CT_DISCRETE*/

/* parametric tests for Gaussian variables. */
double ct_gaustests(tabular dtx, tabular dt, double *pvalue, double *df,
    test_e test) {

int i = 0, n = 0, ntests = dtx.m.ncols, ncomplete = 0;
double transform = 0, statistic = 0, lambda = 0, *mean = NULL;
bool complete = tabular_is_complete(dtx) && tabular_is_complete(dt);
bool *missing_yz = NULL, *missing_all = NULL;
covariance cov = { 0 }, basecov = { 0 };

  /* if y or any conditioning variable is latent, no candidate has complete
   * observations: they are all independent. (The leading slot reserved for the
   * variable under test is never latent.) */
  if (tabular_has_latent(dt)) {

    for (i = 0; i < ntests; i++)
      pvalue[i] = 1;
    return statistic;

  }/*THEN*/

  if (complete) {

    /* the degrees of freedom are the same for all the tests. */
    *df = gaussian_cdf(test, dt.m.nobs, dt.m.ncols - 2);

    if (*df < 1) {

      /* if there are not enough degrees of freedom, return independence. */
      warning("trying to do a conditional independence test with zero degrees of freedom.");

      *df = 0;
      statistic = 0;
      for (i = 0; i < ntests; i++)
        pvalue[i] = 1;

      return statistic;

    }/*THEN*/

    /* cache the means and compute the covariance matrix of the conditioning
     * variables once; it is updated with each variable under test. */
    tabular_cache_means(&dt, 1);
    cov = new_covariance(dt.m.ncols, TRUE);
    basecov = new_covariance(dt.m.ncols, TRUE);
    c_covmat(dt.ccol, dt.mean, dt.m.nobs, dt.m.ncols, cov, 1);
    if (ntests > 1)
      copy_covariance(&cov, &basecov);
    mean = dt.mean;

  }/*THEN*/
  else {

    /* the covariance matrix is recomputed from the complete cases of each
     * variable under test. */
    mean = Calloc1D(dt.m.ncols, sizeof(double));
    cov = new_covariance(dt.m.ncols, TRUE);
    missing_yz = Calloc1D(dt.m.nobs, sizeof(bool));
    missing_all = Calloc1D(dt.m.nobs, sizeof(bool));
    tabular_incomplete_cases(&dt, missing_yz, 0, 1);

  }/*ELSE*/

  /* compute the partial correlation and the test statistic. */
  for (i = 0; i < ntests; i++) {

    /* a latent x is independent of y. */
    if (dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    /* extract and plug in the i-th variable. */
    dt.ccol[0] = dtx.ccol[i];

    if (complete) {

      dt.mean[0] = c_mean(dt.ccol[0], dt.m.nobs);
      if (ntests > 1)
        copy_covariance(&basecov, &cov);
      c_update_covmat(dt.ccol, dt.mean, 0, dt.m.nobs, dt.m.ncols, cov.mat);
      n = dt.m.nobs;

    }/*THEN*/
    else {

      c_covmat_with_missing(dt.ccol, dt.m.nobs, dt.m.ncols, missing_yz,
        missing_all, mean, cov.mat, &ncomplete);
      *df = gaussian_cdf(test, ncomplete, dt.m.ncols - 2);

      if ((ncomplete == 0) || (*df < 1)) {

        /* if there are not enough degrees of freedom, return independence. */
        warning("trying to do a conditional independence test with zero degrees of freedom.");

        *df = 0;
        statistic = 0;
        pvalue[i] = 1;
        continue;

      }/*THEN*/

      n = ncomplete;

    }/*ELSE*/

    if (test == COR) {

      statistic = c_fast_pcor(cov, 0, 1, NULL, TRUE);
      transform = cor_t_trans(statistic, *df);
      pvalue[i] = 2 * pt(fabs(transform), *df, FALSE, FALSE);

    }/*THEN*/
    else if (test == MI_G) {

      statistic = c_fast_pcor(cov, 0, 1, NULL, TRUE);
      statistic = 2 * n * cor_mi_trans(statistic);
      pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

    }/*THEN*/
    else if (test == MI_G_SH) {

      lambda = covmat_lambda(dt.ccol, mean, cov, dt.m.nobs,
                 complete ? NULL : missing_all, n);
      covmat_shrink(cov, lambda);
      statistic = c_fast_pcor(cov, 0, 1, NULL, TRUE);
      statistic = 2 * n * cor_mi_trans(statistic);
      pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

    }/*THEN*/
    else if (test == ZF) {

      statistic = c_fast_pcor(cov, 0, 1, NULL, TRUE);
      statistic = cor_zf_trans(statistic, *df);
      pvalue[i] = 2 * pnorm(fabs(statistic), 0, 1, FALSE, FALSE);

    }/*THEN*/

  }/*FOR*/

  if (complete) {

    /* the means were cached on the local copy of the table, free them here. */
    Free1D(dt.mean);
    FreeCOV(basecov);
    FreeCOV(cov);

  }/*THEN*/
  else {

    Free1D(mean);
    Free1D(missing_yz);
    Free1D(missing_all);
    FreeCOV(cov);

  }/*ELSE*/

  return statistic;

}/*CT_GAUSTESTS*/

/* conditional linear Gaussian mutual information test. */
double c_micg_chisq(tabular dtx, tabular dty, tabular dtz, int *zptr, int llz,
    double *df, bool copy) {

double statistic = 0;

  if (dty.m.flag[0].discrete && dtx.m.flag[0].discrete) {

    if (dtz.nccols - 1 > 0) {

      /* need to reverse conditioning to actually compute the test. */
      statistic = 2 * dtz.m.nobs *
                    c_cmicg_unroll(dtx.dcol[0], dtx.nlvl[0], dty.dcol[0],
                      dty.nlvl[0], zptr, llz, dtz.ccol + 1, dtz.nccols - 1, df,
                      dtz.m.nobs);

    }/*THEN*/
    else {

      /* the test reverts back to a discrete mutual information test. */
      statistic = c_cchisqtest(dtx.dcol[0], dtx.nlvl[0],
                    dty.dcol[0], dty.nlvl[0], zptr, llz, dtz.m.nobs, df, MI, TRUE);

    }/*ELSE*/

  }/*THEN*/
  else if (dty.m.flag[0].continuous && dtx.m.flag[0].continuous) {

    if (copy)
      memcpy(dtz.ccol[0], dtx.ccol[0], dtz.m.nobs * sizeof(double));
    else
     dtz.ccol[0] = dtx.ccol[0];

    statistic = 2 * dtz.m.nobs * c_cmicg(dty.ccol[0], dtz.ccol, dtz.nccols, NULL,
                      0, zptr, llz, dtz.nlvl, dtz.m.nobs, df);

  }/*THEN*/
  else if (dty.m.flag[0].discrete && dtx.m.flag[0].continuous) {

    if (copy)
      memcpy(dtz.dcol[0], dty.dcol[0], dtz.m.nobs * sizeof(int));
    else
      dtz.dcol[0] = dty.dcol[0];

    dtz.nlvl[0] = dty.nlvl[0];
    statistic = 2 * dtz.m.nobs * c_cmicg(dtx.ccol[0], dtz.ccol + 1,
                      dtz.nccols - 1, dtz.dcol, dtz.ndcols, zptr, llz, dtz.nlvl,
                      dtz.m.nobs, df);

  }/*THEN*/
  else if (dty.m.flag[0].continuous && dtx.m.flag[0].discrete) {

    if (copy)
      memcpy(dtz.dcol[0], dtx.dcol[0], dtz.m.nobs * sizeof(int));
    else
      dtz.dcol[0] = dtx.dcol[0];

    dtz.nlvl[0] = dtx.nlvl[0];
    statistic = 2 * dtz.m.nobs * c_cmicg(dty.ccol[0], dtz.ccol + 1, dtz.nccols - 1,
                      dtz.dcol, dtz.ndcols, zptr, llz, dtz.nlvl,
                      dtz.m.nobs, df);

  }/*ELSE*/

  return statistic;

}/*C_MICG_CHISQ*/

double ct_micg(tabular dtx, tabular dty, tabular dtz, double *pvalue,
    double *df) {

int i = 0, j = 0, *zptr = 0, llz = 0;
double statistic = 0;
bool complete = tabular_is_complete(dtx) && tabular_is_complete(dty) &&
                  tabular_is_complete(dtz);
bool *missing_x = NULL, *missing_yz = NULL, *missing_all = NULL;
tabular dtx_cur = { 0 }, dt_complete = { 0 }, dtx_complete = { 0 };
tabular dty_complete = { 0 };

  /* if y or any conditioning variable is latent, no candidate has complete
   * observations: they are all independent. */
  if (dty.m.flag[0].latent || tabular_has_latent(dtz)) {

    for (i = 0; i < dtx.m.ncols; i++)
      pvalue[i] = 1;
    return statistic;

  }/*THEN*/

  /* the variable to test is plugged into this one-column table in turn. */
  dtx_cur = empty_tabular(dtx.m.nobs, 1, 1);

  if (complete) {

    /* the discrete parents' configurations do not change across tests. */
    if (dtz.ndcols - 1 > 0) {

      zptr = Calloc1D(dtz.m.nobs, sizeof(int));
      c_fast_config(dtz.dcol + 1, dtz.m.nobs, dtz.ndcols - 1, dtz.nlvl + 1,
        zptr, &llz, 1);

    }/*THEN*/

  }/*THEN*/
  else {

    /* allocate the tables for the complete observations and the missing values
     * indicators, recomputed for each variable. */
    dt_complete = new_tabular(dtz.m.nobs, dtz.ndcols, dtz.nccols);
    dty_complete = new_tabular(dty.m.nobs, dty.ndcols, dty.nccols);
    missing_x = Calloc1D(dtz.m.nobs, sizeof(bool));
    missing_yz = Calloc1D(dtz.m.nobs, sizeof(bool));
    missing_all = Calloc1D(dtz.m.nobs, sizeof(bool));

    tabular_incomplete_cases(&dty, missing_yz, 0, 0);
    tabular_incomplete_cases(&dtz, missing_yz, 1, 1);

  }/*ELSE*/

  for (i = 0; i < dtx.m.ncols; i++) {

    /* a latent x is independent of y. */
    if (dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    /* extract the variable to test. */
    tabular_subset_columns(&dtx, &dtx_cur, &i, 1);

    if (complete) {

      statistic = c_micg_chisq(dtx_cur, dty, dtz, zptr, llz, df, FALSE);

    }/*THEN*/
    else {

      /* the complete-case x table is sized for this one variable; subset the
       * observations that are complete for it. */
      dtx_complete = new_tabular(dtx_cur.m.nobs, dtx_cur.ndcols, dtx_cur.nccols);
      tabular_incomplete_cases(&dtx_cur, missing_x, 0, 0);
      for (j = 0; j < dtz.m.nobs; j++)
        missing_all[j] = missing_x[j] || missing_yz[j];
      tabular_subsample_by_logical(&dtx_cur, &dtx_complete, missing_all, 0, 0);
      tabular_subsample_by_logical(&dtz, &dt_complete, missing_all, 1, 1);
      tabular_subsample_by_logical(&dty, &dty_complete, missing_all, 0, 0);

      /* assume independence if there are no complete observations. */
      if (dt_complete.m.nobs == 0) {

        statistic = 0;
        pvalue[i] = 1;
        FreeTAB(dtx_complete);
        continue;

      }/*THEN*/

      if (dtz.ndcols - 1 > 0) {

        zptr = Calloc1D(dtz.m.nobs, sizeof(int));
        c_fast_config(dt_complete.dcol + 1, dt_complete.m.nobs,
          dt_complete.ndcols - 1, dt_complete.nlvl + 1, zptr, &llz, 1);

      }/*THEN*/

      statistic = c_micg_chisq(dtx_complete, dty_complete, dt_complete,
                    zptr, llz, df, TRUE);
      Free1D(zptr);
      FreeTAB(dtx_complete);

    }/*ELSE*/

    pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

  }/*FOR*/

  FreeTAB(dtx_cur);

  if (complete) {

    Free1D(zptr);

  }/*THEN*/
  else {

    Free1D(missing_x);
    Free1D(missing_yz);
    Free1D(missing_all);
    FreeTAB(dt_complete);
    FreeTAB(dty_complete);

  }/*ELSE*/

  return statistic;

}/*CT_MICG*/

/* discrete permutation tests. */
double ct_dperm(tabular dtx, tabular dty, tabular dtz, double *pvalue, double *df,
    test_e type, int B, double a) {

int i = 0, *zptr = NULL, llz = 0;
double statistic = 0;

  /* if y or any conditioning variable is latent, no candidate has complete
   * observations: they are all independent (and no random numbers are
   * consumed). */
  if (dty.m.flag[0].latent || tabular_has_latent(dtz)) {

    for (i = 0; i < dtx.m.ncols; i++)
      pvalue[i] = 1;
    return statistic;

  }/*THEN*/

  zptr = Calloc1D(dtz.m.nobs, sizeof(int));
  c_fast_config(dtz.dcol, dtz.m.nobs, dtz.m.ncols, dtz.nlvl, zptr, &llz, 1);

  for (i = 0; i < dtx.m.ncols; i++) {

    /* a latent x is independent of y. */
    if (dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    c_cmcarlo(dtx.dcol[i], dtx.nlvl[i], dty.dcol[0], dty.nlvl[0], zptr, llz,
      dtx.m.nobs, B, &statistic, pvalue + i, a, type, df);

  }/*FOR*/

  Free1D(zptr);

  return statistic;

}/*CT_DPERM*/

/* continuous permutation tests. */
double ct_gperm(tabular dtx, tabular dt, double *pvalue, double *df, test_e type,
    int B, double a, bool complete) {

int i = 0, j = 0, k = 0, nc = 0;
double *yptr = dt.ccol[1], **complete_column = NULL, statistic = 0;
bool *missing_yz = NULL;
tabular dt_complete = { 0 };

  /* if y or any conditioning variable is latent, no candidate has complete
   * observations: they are all independent (and no random numbers are
   * consumed). The leading slot reserved for the variable under test is never
   * latent; a latent x is handled per candidate below. */
  if (tabular_has_latent(dt)) {

    for (i = 0; i < dtx.m.ncols; i++)
      pvalue[i] = 1;
    return statistic;

  }/*THEN*/

  if (!complete) {

    /* allocate and initialize the missing values indicators. */
    missing_yz = Calloc1D(dt.m.nobs, sizeof(bool));
    tabular_incomplete_cases(&dt, missing_yz, 0, 1);

    dt_complete = new_tabular(dt.m.nobs, 0, dt.m.ncols);
    complete_column = dt_complete.ccol;

  }/*THEN*/
  else {

    complete_column = dt.ccol;

  }/*ELSE*/

  for (i = 0; i < dtx.m.ncols; i++) {

    /* a latent x is independent of y. */
    if (dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    /* swap the first column and restore the second, which is that undergoing
     * permutation (backward compatibility from set random seed). */
    dt.ccol[0] = dtx.ccol[i];
    dt.ccol[1] = yptr;

    if (!complete) {

      for (k = 0, nc = 0; k < dt.m.nobs; k++) {

        if (missing_yz[k] || ISNAN(dt.ccol[0][k]))
          continue;

        for (j = 0; j < dt.m.ncols; j++)
          complete_column[j][nc] = dt.ccol[j][k];
        nc++;

      }/*FOR*/

    }/*THEN*/
    else {

      nc = dt.m.nobs;

    }/*ELSE*/

    c_gauss_cmcarlo(complete_column, dt.m.ncols, nc, 0, 1, B, &statistic,
      pvalue + i, a, type);

  }/*FOR*/

  if (!complete) {

    Free1D(missing_yz);
    FreeTAB(dt_complete);

  }/*THEN*/

  return statistic;

}/*CT_GPERM*/

