#include "../../include/rcore.h"
#include "../../core/data.table.h"
#include "../../include/globals.h"
#include "../../minimal/common.h"
#include "../../minimal/data.frame.h"
#include "../patterns.h"
#include "../tests.h"

/* user-provided test function; kept at the R/C boundary because it builds the
 * response SEXP and calls back into R for each node. */
static double ut_custom(SEXP x, SEXP y, SEXP data, SEXP custom_fn,
    SEXP custom_args, double *pvalue) {

double statistic = 0;
int i = 0;
SEXP xi;

  PROTECT(xi = allocVector(STRSXP, 1));

  for (i = 0; i < length(x); i++) {

    /* iterate over the nodes, one at a time... */
    SET_STRING_ELT(xi, 0 , STRING_ELT(x, i));
    /* ... compute the test statistic, and fill the p-value array. */
    statistic = custom_test_function(xi, y, R_NilValue, data, custom_fn,
                  custom_args, pvalue + i);

  }/*FOR*/

  UNPROTECT(1);

  return statistic;

}/*UT_CUSTOM*/

/* unconditional independence tests. */
SEXP utest(SEXP x, SEXP y, SEXP data, SEXP test, SEXP alpha, SEXP extra_args,
    SEXP learning, SEXP debug) {

int ntests = length(x);
bool debugging = isTRUE(debug);
double a = NUM(alpha);
double *pvalue = NULL, statistic = 0, df = NA_REAL;
const char *t = CHAR(STRING_ELT(test, 0));
test_e test_type = test_to_enum(t);
SEXP xx, yy, result;

  /* allocate the return value, which has the same length as x. */
  PROTECT(result = allocVector(REALSXP, ntests));
  setAttrib(result, R_NamesSymbol, x);
  pvalue = REAL(result);
  /* set all elements to zero. */
  memset(pvalue, '\0', ntests * sizeof(double));

  /* extract the variables from the data, keeping the column names so the
   * completeness flags can be matched onto them. */
  PROTECT(xx = c_dataframe_column(data, x, FALSE, TRUE));
  PROTECT(yy = c_dataframe_column(data, y, 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);

    statistic = ut_discrete(dtx, dty, pvalue, &df, test_type);

    FreeTAB(dtx);
    FreeTAB(dty);

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

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

    statistic = ut_gaustests(dtx, dty, pvalue, &df, test_type);

    FreeTAB(dtx);
    FreeTAB(dty);

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

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

    statistic = ut_micg(dtx, dty, pvalue, &df);

    FreeTAB(dtx);
    FreeTAB(dty);

  }/*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);
    int B = INT(getListElement(extra_args, "B"));

    statistic = ut_dperm(dtx, dty, pvalue, &df, test_type, B,
                  IS_SMC(test_type) ? NUM(alpha) : 1);

    FreeTAB(dtx);
    FreeTAB(dty);

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

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

    statistic = ut_gperm(dtx, dty, pvalue, test_type, B,
                  IS_SMC(test_type) ? NUM(alpha) : 1,
                  tabular_is_complete(dtx) && tabular_is_complete(dty));

    FreeTAB(dtx);
    FreeTAB(dty);

  }/*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");

    statistic = ut_custom(x, y, data, custom_fn, custom_args, pvalue);

  }/*THEN*/

  /* print the outcome of each test as soon as it has been computed, so that the
   * trace is available even if a later test errors out. */
  if (debugging && (test_type != ENOTEST))
    for (int i = 0; i < ntests; i++)
      debug_test_outcome(CHAR(STRING_ELT(x, i)), CHAR(STRING_ELT(y, 0)),
        NULL, 0, pvalue[i], a);

  UNPROTECT(3);

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

  /* increase the test counter. */
  test_counter += ntests;

  if (isTRUE(learning))
    return result;
  else
    return c_create_htest(statistic, test, pvalue[ntests - 1], df, extra_args);

}/*UTEST*/

