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

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

int i = 0, ntests = dtx.m.ncols, nobs = dtx.m.nobs;
int llx = 0, lly = dty.nlvl[0], *xptr = NULL, *yptr = dty.dcol[0];
double statistic = 0;

  for (i = 0; i < ntests; i++) {

    /* a latent x or y leaves the contingency table empty: independence. */
    if (dty.m.flag[0].latent || 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_chisqtest(xptr, llx, yptr, lly, 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_shmi(xptr, llx, yptr, lly, nobs, TRUE);
      *df = ((double)(llx - 1) * (double)(lly - 1));
      pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

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

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

    }/*THEN*/

  }/*FOR*/

  return statistic;

}/*UT_DISCRETE*/

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

int i = 0, n = 0, ntests = dtx.m.ncols, nobs = dtx.m.nobs, ncomplete = 0;
double transform = 0, *xptr = NULL, *yptr = dty.ccol[0];
double xm = 0, ym = 0, xsd = 0, ysd = 0, statistic = 0;
bool complete = tabular_is_complete(dtx) && tabular_is_complete(dty);

  if (complete) {

    /* the degrees of freedom and the mean and variance of y are the same for
     * all the tests. */
    *df = gaussian_df(test, nobs);

    if (*df < 1) {

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

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

      return statistic;

    }/*THEN*/

    ym = c_mean(yptr, nobs);
    ysd = c_sse(yptr, ym, nobs);

  }/*THEN*/

  for (i = 0; i < ntests; i++) {

    /* a latent x or y has no complete cases: independence, no test. */
    if (dty.m.flag[0].latent || dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    xptr = dtx.ccol[i];

    if (complete) {

      xm = c_mean(xptr, nobs);
      xsd = c_sse(xptr, xm, nobs);
      statistic = c_fast_cor(xptr, yptr, nobs, xm, ym, xsd, ysd);
      n = nobs;

    }/*THEN*/
    else {

      /* the correlation is computed from the complete cases of each variable. */
      statistic = c_cor_with_missing(xptr, yptr, nobs, &xm, &ym, &xsd, &ysd,
                    &ncomplete);
      *df = gaussian_df(test, ncomplete);

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

        /* if there are not enough degrees of freedom, including when there are
         * no complete observations, return perfect independence. */
        warning("trying to do an independence test with zero degrees of freedom.");

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

      }/*THEN*/

      n = ncomplete;

    }/*ELSE*/

    if (test == COR) {

      transform = cor_t_trans(statistic, *df);
      pvalue[i] = 2 * pt(fabs(transform), *df, FALSE, FALSE);

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

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

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

      statistic *= 1 - cor_lambda(xptr, yptr, nobs, n, xm, ym, xsd, ysd,
                         statistic);
      statistic = 2 * n * cor_mi_trans(statistic);
      pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

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

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

    }/*THEN*/

  }/*FOR*/

  return statistic;

}/*UT_GAUSTESTS*/

/* conditional linear Gaussian mutual information test. */
double ut_micg(tabular dtx, tabular dty, double *pvalue, double *df) {

int i = 0, n = 0, nc = 0, ntests = dtx.m.ncols, nobs = dtx.m.nobs;
int lly = 0, disc_lvls = 0, *disc_ptr = NULL;
double xm = 0, xsd = 0, ym = 0, ysd = 0, statistic = 0, *cont_ptr = NULL;
void *xptr = NULL, *yptr = NULL;
bool complete = tabular_is_complete(dtx) && tabular_is_complete(dty);
bool ydiscrete = dty.m.flag[0].discrete, xdiscrete = FALSE;
tabular dtx_cur = { 0 };

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

  if (ydiscrete) {

    lly = dty.nlvl[0];
    yptr = dty.dcol[0];

  }/*THEN*/
  else {

    yptr = dty.ccol[0];
    /* cache the mean and variance for the complete-data fast path. */
    if (complete) {

      ym = c_mean(yptr, nobs);
      ysd = c_sse(yptr, ym, nobs);

    }/*THEN*/

  }/*ELSE*/

  for (i = 0; i < ntests; i++) {

    /* a latent x or y has no complete cases: independence, no test. */
    if (dty.m.flag[0].latent || dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    /* extract the variable under test, which may be discrete or continuous. */
    tabular_subset_columns(&dtx, &dtx_cur, &i, 1);
    xdiscrete = dtx_cur.m.flag[0].discrete;

    if (ydiscrete && xdiscrete) {

      /* if both nodes are discrete, the test reverts back to a discrete mutual
       * information test, which handles missing data internally. */
      statistic = c_chisqtest(dtx_cur.dcol[0], dtx_cur.nlvl[0], yptr, lly, nobs,
                    df, MI, TRUE);
      pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

    }/*THEN*/
    else if (!ydiscrete && !xdiscrete) {

      /* if both nodes are continuous, the test reverts back to a Gaussian
       * mutual information test. */
      if (complete) {

        xptr = dtx_cur.ccol[0];
        xm = c_mean(xptr, nobs);
        xsd = c_sse(xptr, xm, nobs);
        statistic = c_fast_cor(xptr, yptr, nobs, xm, ym, xsd, ysd);
        n = nobs;

      }/*THEN*/
      else {

        statistic = c_cor_with_missing(dtx_cur.ccol[0], yptr, nobs, NULL, NULL,
                      NULL, NULL, &nc);
        n = nc;

      }/*ELSE*/

      *df = 1;
      statistic = 2 * n * cor_mi_trans(statistic);
      pvalue[i] = pchisq(statistic, *df, FALSE, FALSE);

    }/*THEN*/
    else {

      /* one node is discrete and the other continuous. */
      if (xdiscrete) {

        disc_ptr = dtx_cur.dcol[0];
        disc_lvls = dtx_cur.nlvl[0];
        cont_ptr = yptr;

      }/*THEN*/
      else {

        disc_ptr = yptr;
        disc_lvls = lly;
        cont_ptr = dtx_cur.ccol[0];

      }/*ELSE*/

      if (complete) {

        statistic = 2 * nobs * c_micg(cont_ptr, disc_ptr, disc_lvls, nobs, df,
                                  NULL, FALSE);

      }/*THEN*/
      else {

        statistic = c_micg(cont_ptr, disc_ptr, disc_lvls, nobs, df, &nc, TRUE);
        statistic *= 2 * nc;

      }/*ELSE*/

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

    }/*ELSE*/

  }/*FOR*/

  FreeTAB(dtx_cur);

  return statistic;

}/*UT_MICG*/

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

int i = 0, ntests = dtx.m.ncols, nobs = dtx.m.nobs;
int llx = 0, lly = dty.nlvl[0], *xptr = NULL, *yptr = dty.dcol[0];
double statistic = 0;

  for (i = 0; i < ntests; i++) {

    /* a latent x or y leaves the contingency table empty: independence (and no
     * random numbers are consumed by the permutation kernel). */
    if (dty.m.flag[0].latent || dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    xptr = dtx.dcol[i];
    llx = dtx.nlvl[i];
    c_mcarlo(xptr, llx, yptr, lly, nobs, B, &statistic, pvalue + i,
      a, type, df);

  }/*FOR*/

  return statistic;

}/*UT_DPERM*/

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

int i = 0, k = 0, ntests = dtx.m.ncols, nobs = dtx.m.nobs, nc = nobs;
double *yptr = dty.ccol[0], *xptr = NULL;
double *yptr_complete = NULL, *xptr_complete = NULL, statistic = 0;

  if (!complete) {

    xptr_complete = Calloc1D(nobs, sizeof(double));
    yptr_complete = Calloc1D(nobs, sizeof(double));

  }/*THEN*/
  else {

    yptr_complete = yptr;

  }/*ELSE*/

  for (i = 0; i < ntests; i++) {

    /* a latent x or y has no complete cases: independence (and no random
     * numbers are consumed by the permutation kernel). */
    if (dty.m.flag[0].latent || dtx.m.flag[i].latent) {

      pvalue[i] = 1;
      continue;

    }/*THEN*/

    xptr = dtx.ccol[i];

    if (!complete) {

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

        if (ISNAN(xptr[k]) || ISNAN(yptr[k]))
          continue;

        xptr_complete[nc] = xptr[k];
        yptr_complete[nc] = yptr[k];
        nc++;

      }/*FOR*/

    }/*THEN*/
    else {

      xptr_complete = xptr;

    }/*ELSE*/

    c_gauss_mcarlo(xptr_complete, yptr_complete, nc, B, pvalue + i,
      a, type, &statistic);

  }/*FOR*/

  if (!complete) {

    Free1D(xptr_complete);
    Free1D(yptr_complete);

  }/*THEN*/

  return statistic;

}/*UT_GPERM*/
