#include "../../include/rcore.h"
#include "../../core/allocations.h"
#include "../../core/covariance.matrix.h"
#include "../../core/data.table.h"
#include "../../core/math.functions.h"
#include "../../core/sets.h"
#include "../../include/globals.h"
#include "../../minimal/common.h"
#include "../../minimal/data.frame.h"
#include "../../minimal/strings.h"
#include "../patterns.h"
#include "../tests.h"

/* marshal the pure-C search result into the SEXP the callers expect: a named
 * vector with the p-value range, carrying the separating set as an attribute. */
static SEXP ast_result_to_SEXP(ast_result res) {

SEXP retval, dsep_set;

  PROTECT(retval = mkRealVec(3, res.pvalue, res.min_pvalue, res.max_pvalue));
  setAttrib(retval, R_NamesSymbol,
    mkStringVec(3, "p.value", "min.p.value", "max.p.value"));

  if (res.separated) {

    PROTECT(dsep_set = allocVector(STRSXP, res.ndsep));
    for (int i = 0; i < res.ndsep; i++)
      SET_STRING_ELT(dsep_set, i, mkChar(res.dsep[i]));
    setAttrib(retval, BN_DsepsetSymbol, dsep_set);
    UNPROTECT(1);

  }/*THEN*/
  else {

    setAttrib(retval, BN_DsepsetSymbol, R_NilValue);

  }/*ELSE*/

  UNPROTECT(1);

  return retval;

}/*AST_RESULT_TO_SEXP*/

/* user-provided test function; kept at the R/C boundary because it builds the
 * conditioning-set SEXP and calls back into R for each subset. */
static ast_result ast_custom(SEXP x, SEXP y, SEXP sx, SEXP fixed, SEXP data,
    int minsize, int maxsize, double a, SEXP custom_fn, SEXP custom_args,
    bool debugging) {

int i = 0, j = 0, cursize = 0, nf = length(fixed), *subset = NULL;
double pvalue = 0, min_pvalue = 1, max_pvalue = 0;
ast_result retval;
SEXP sxi;

  for (cursize = imax(1, minsize); cursize <= maxsize; cursize++) {

    /* allocate and initialize the subset. */
    subset = Calloc1D(cursize + nf, sizeof(int));
    /* initialize the first subset. */
    first_subset(subset + nf, cursize, nf);
    for (i = 0; i < nf; i++)
      subset[i] = i;

    /* iterate over subsets. */
    do {

      /* prepare the current subset. */
      PROTECT(sxi = allocVector(STRSXP, cursize + nf));
       for (j = 0; j < cursize + nf; j++)
         SET_STRING_ELT(sxi, j, STRING_ELT(sx, subset[j]));

      /* compute the test statistic and update the p-values. */
      custom_test_function(x, y, sxi, data, custom_fn, custom_args, &pvalue);
      update_pvalue_range(pvalue, &min_pvalue, &max_pvalue);
      /* increment the test counter. */
      test_counter++;

      if (debugging) {

        int nsx = length(sxi);
        const char **cond = Calloc1D(nsx, sizeof(char *));

        for (j = 0; j < nsx; j++)
          cond[j] = CHAR(STRING_ELT(sxi, j));
        debug_test_outcome(CHAR(STRING_ELT(x, 0)), CHAR(STRING_ELT(y, 0)), cond,
          nsx, pvalue, a);

        Free1D(cond);

      }/*THEN*/

      /* if the p-value is not significant... */
      if (pvalue > a) {

        const char **subnames = NULL;

        /* ... collect the labels of the nodes in the separating set... */
        subnames = Calloc1D(length(sxi), sizeof(char*));
        for (j = 0; j < length(sxi); j++)
          subnames[j] = CHAR(STRING_ELT(sxi, j));

        /*... and return. */
        retval = ast_result_make(pvalue, min_pvalue, max_pvalue, TRUE,
                   subnames, length(sxi));

        Free1D(subset);
        Free1D(subnames);

        UNPROTECT(1);
        return retval;

      }/*THEN*/

      UNPROTECT(1);

    } while (next_subset(subset + nf, cursize, length(sx) - nf, nf));

    Free1D(subset);

  }/*FOR*/

  return ast_result_make(pvalue, min_pvalue, max_pvalue, FALSE, NULL, 0);

}/*AST_CUSTOM*/

SEXP allsubs_test(SEXP x, SEXP y, SEXP sx, SEXP fixed, SEXP data, SEXP test,
    SEXP alpha, SEXP extra_args, SEXP min, SEXP max, SEXP debug) {

int minsize = INT(min), maxsize = INT(max);
int i = 0, nf = length(fixed);
double pvalue = 0, min_pvalue = 1, max_pvalue = 0, a = NUM(alpha);
const char *t = CHAR(STRING_ELT(test, 0));
bool debugging = isTRUE(debug);
test_e test_type = test_to_enum(t);
SEXP xx, yy, zz, res = R_NilValue;
ast_result r = { 0 };

  /* call indep_test to deal with tests which have no moving parts, that is when
   * the conditioning set is empty of completely fixed. */
  if (minsize == 0) {

    pvalue = NUM(indep_test(x, y, fixed, data, test, alpha, extra_args,
                   TRUESEXP, FALSESEXP));
    update_pvalue_range(pvalue, &min_pvalue, &max_pvalue);

    /* indep_test() already counted this test through utest()/ctest(). */

    if (debugging) {

      const char **cond = Calloc1D(nf, sizeof(char *));

      for (i = 0; i < nf; i++)
        cond[i] = CHAR(STRING_ELT(fixed, i));
      debug_test_outcome(CHAR(STRING_ELT(x, 0)), CHAR(STRING_ELT(y, 0)), cond,
        nf, pvalue, a);

      Free1D(cond);

    }/*THEN*/

    if (pvalue > a) {

      /* the fixed variables separate x and y; report them as the d-sep set. */
      PROTECT(res = ast_result_to_SEXP(
                      ast_result_make(pvalue, min_pvalue, max_pvalue, FALSE,
                        NULL, 0)));
      setAttrib(res, BN_DsepsetSymbol, fixed);

      UNPROTECT(1);
      return res;

    }/*THEN*/
    else {

      /* return even if the variables are not found to be independent, since
       * there are no more tests left to do. */
      if (maxsize == 0)
        return ast_result_to_SEXP(
                 ast_result_make(pvalue, min_pvalue, max_pvalue, FALSE, NULL, 0));

    }/*ELSE*/

  }/*THEN*/

  /* extract the variables from the data. */
  PROTECT(xx = c_dataframe_column(data, x, FALSE, TRUE));
  PROTECT(yy = c_dataframe_column(data, y, FALSE, TRUE));
  PROTECT(zz = c_dataframe_column(data, sx, FALSE, TRUE));

  if (IS_DISCRETE_ASYMPTOTIC_TEST(test_type)) {

    /* parametric tests for discrete variables. */
    tabular dtx = tabular_from_SEXP(xx, 0, 0, R_NilValue);
    tabular dty = tabular_from_SEXP(yy, 0, 0, R_NilValue);
    tabular dtz = tabular_from_SEXP(zz, 0, 0, R_NilValue);

    r = ast_discrete(dtx, dty, dtz, nf, minsize, maxsize, test_type, a,
          debugging);

    FreeTAB(dtx);
    FreeTAB(dty);
    FreeTAB(dtz);

  }/*THEN*/
  else if ((test_type == COR) || (test_type == ZF) || (test_type == MI_G) ||
           (test_type == MI_G_SH)) {

    /* parametric tests for Gaussian variables; x, y and the conditioning set
     * are extracted together so x, y occupy the leading columns. */
    SEXP xyz;
    PROTECT(xyz = c_dataframe_column_set(data, x, y, sx));
    tabular dt = tabular_from_SEXP(xyz, 0, 0, R_NilValue);

    r = ast_gaustests(dt, nf, minsize, maxsize, a, debugging, test_type);

    FreeTAB(dt);
    UNPROTECT(1);

  }/*THEN*/
  else if (test_type == MI_CG) {

    /* conditional linear Gaussian test. */
    tabular dtx = tabular_from_SEXP(xx, 0, 0, R_NilValue);
    tabular dty = tabular_from_SEXP(yy, 0, 0, R_NilValue);
    tabular dtz = tabular_from_SEXP(zz, 1, 1, R_NilValue);

    r = ast_micg(dtx, dty, dtz, nf, minsize, maxsize, a, debugging);

    FreeTAB(dtx);
    FreeTAB(dty);
    FreeTAB(dtz);

  }/*THEN*/
  else if (IS_DISCRETE_PERMUTATION_TEST(test_type)) {

    /* discrete permutation tests. */
    tabular dtx = tabular_from_SEXP(xx, 0, 0, R_NilValue);
    tabular dty = tabular_from_SEXP(yy, 0, 0, R_NilValue);
    tabular dtz = tabular_from_SEXP(zz, 0, 0, R_NilValue);
    int B = INT(getListElement(extra_args, "B"));

    r = ast_dperm(dtx, dty, dtz, nf, minsize, maxsize, a, test_type, B,
          IS_SMC(test_type) ? a : 1, debugging);

    FreeTAB(dtx);
    FreeTAB(dty);
    FreeTAB(dtz);

  }/*THEN*/
  else if (IS_CONTINUOUS_PERMUTATION_TEST(test_type)) {

    /* continuous permutation tests; x, y and the conditioning set are extracted
     * together so x, y occupy the leading columns. */
    SEXP xyz;
    PROTECT(xyz = c_dataframe_column_set(data, x, y, sx));
    tabular dt = tabular_from_SEXP(xyz, 0, 0, R_NilValue);
    tabular_cache_means(&dt, 0);
    int B = INT(getListElement(extra_args, "B"));

    r = ast_gperm(dt, nf, minsize, maxsize, a, test_type, B,
          IS_SMC(test_type) ? a : 1, tabular_is_complete(dt), debugging);

    FreeTAB(dt);
    UNPROTECT(1);

  }/*THEN*/
  else if (test_type == CUSTOM_T) {

    /* user-provided test function. */
    SEXP custom_fn = getListElement(extra_args, "fun");
    SEXP custom_args = getListElement(extra_args, "args");

    r = ast_custom(x, y, sx, fixed, data, minsize, maxsize, a, custom_fn,
          custom_args, debugging);

  }/*THEN*/

  /* the p-values of the tests performed under minsize == 0 should be considered
   * when computing the maximum and minimum p-values. */
  if (minsize == 0) {

    r.min_pvalue = r.min_pvalue < min_pvalue ? r.min_pvalue : min_pvalue;
    r.max_pvalue = r.max_pvalue > max_pvalue ? r.max_pvalue : max_pvalue;

  }/*THEN*/

  /* build the return value while the data columns are still protected, since
   * the separating-set names borrow from them, and release the owned copy. */
  res = ast_result_to_SEXP(r);
  Free1D(r.dsep);

  UNPROTECT(3);

  /* catch-all for unknown tests (after deallocating memory.) */
  if (test_type == ENOTEST)
    error("unknown test statistic '%s'.", t);

  return res;

}/*ALLSUBS_TEST*/
