#include "../../include/rcore.h"
#include "../../core/allocations.h"
#include "../../core/data.table.h"
#include "../../fitted/fitted.h"
#include "../../include/globals.h"
#include "../../include/sampling.h"
#include "../../minimal/data.frame.h"
#include "../../minimal/strings.h"
#include "cpq.common.h"

/* working memory freed on return; a failing predicate longjmps straight out
 * and leaks it. */
typedef struct {

  fitted_bn bn;        /* the fitted network to sample from. */
  fixed_node *fixed;   /* the evidence fixed during generation (all-unfixed). */
  int *colmap;         /* output column mapping (cpdist only, else NULL). */

} cpq_scratch;

static void free_scratch(cpq_scratch *sc) {

  if (sc->colmap)
    Free1D(sc->colmap);
  if (sc->fixed)
    FreeEvidence(sc->fixed, sc->bn.nnodes);
  FreeFittedBN(sc->bn);

}/*FREE_SCRATCH*/

/* -- logic sampling (query) ------------------------------------------------ */

/* the running sample counts. */
typedef struct {

  int total;           /* count of samples matching the evidence. */
  int matching;        /* count of samples matching event and evidence. */

} counts;

/* process one batch of m observations: generate them into the reused buffer,
 * evaluate the evidence and the event, and accumulate the counts. */
static void ls_batch(const cpq_query *q, cpq_scratch *sc, SEXP buffer, int m,
    counts *acc, int debug) {

int i = 0, nprot = 0, *ev = NULL, *evt = NULL;
int batch_total = 0, batch_matching = 0;
int event_true = isTRUE(q->target), evidence_true = isTRUE(q->evidence);
tabular dt = { 0 };
SEXP env = R_NilValue, rsxp = R_NilValue;

  /* refill the reused buffer in place (c_rbn_master brackets the RNG). */
  dt = tabular_from_SEXP(buffer, 0, 0, R_NilValue);
  c_rbn_master(sc->bn, dt, sc->fixed, FALSE);
  FreeTAB(dt);

  if (debug)
    Rprintf("* generated %d samples from the bayesian network.\n", m);

  /* build the evaluation environment once for both predicates. */
  if (!evidence_true || !event_true) {

    env = make_data_env(buffer, q->eval_env);
    nprot++;

  }/*THEN*/

  /* evaluate the evidence. */
  if (!evidence_true) {

    rsxp = eval(q->evidence, env);
    PROTECT(rsxp);
    nprot++;
    if (TYPEOF(rsxp) != LGLSXP)
      error("evidence must evaluate to a logical vector.");
    if (length(rsxp) != m)
      error("logical vector for evidence is of length %d instead of %d.",
        length(rsxp), m);
    ev = LOGICAL(rsxp);

  }/*THEN*/

  /* evaluate the event. */
  if (!event_true) {

    rsxp = eval(q->target, env);
    PROTECT(rsxp);
    nprot++;
    if (TYPEOF(rsxp) != LGLSXP)
      error("event must evaluate to a logical vector.");
    if (length(rsxp) != m)
      error("logical vector for event is of length %d instead of %d.",
        length(rsxp), m);
    evt = LOGICAL(rsxp);

  }/*THEN*/

  /* count the samples matching the evidence (and not missing on the event) and
   * those matching both, mirroring the R filtering. */
  for (i = 0; i < m; i++) {

    int fev = evidence_true || (ev[i] != NA_LOGICAL && ev[i]);
    int evt_notna = event_true || (evt[i] != NA_LOGICAL);
    int evt_true = event_true || (evt[i] != NA_LOGICAL && evt[i]);

    if (fev && evt_notna)
      batch_total++;
    if (fev && evt_true)
      batch_matching++;

  }/*FOR*/

  acc->total += batch_total;
  acc->matching += batch_matching;

  if (nprot > 0)
    UNPROTECT(nprot);

  if (debug) {

    double ratio =
      (batch_total == 0) ? 0 : (double)batch_matching / batch_total;

    if (!evidence_true)
      Rprintf("  > evidence matches %d samples out of %d (p = %.7g).\n",
        batch_total, m, (double)batch_total / m);
    else
      Rprintf("  > evidence matches %d samples out of %d (p = 1).\n", m, m);

    if (!event_true)
      Rprintf("  > event matches %d samples out of %d (p = %.7g).\n",
        batch_matching, batch_total, ratio);
    else
      Rprintf("  > event matches %d samples out of %d (p = 1).\n",
        batch_total, batch_total);

  }/*THEN*/

}/*LS_BATCH*/

static double ls_loop(const cpq_query *q, cpq_scratch *sc, int n, int batch,
    int debug) {

int b = 0, nbatches = n / batch, last = n % batch, nprot = 0;
double result = 0;
counts acc = { 0, 0 };
SEXP buffer_full = R_NilValue, buffer_last = R_NilValue;

  if (nbatches > 0) {

    PROTECT(buffer_full = fit2df(q->fitted, batch));
    nprot++;

    for (b = 0; b < nbatches; b++)
      ls_batch(q, sc, buffer_full, batch, &acc, debug);

  }/*THEN*/

  if (last > 0) {

    PROTECT(buffer_last = fit2df(q->fitted, last));
    nprot++;
    ls_batch(q, sc, buffer_last, last, &acc, debug);

  }/*THEN*/

  /* prevent divide-by-zero errors. */
  result = (acc.total == 0) ? 0 : (double)acc.matching / acc.total;

  if (debug && (nbatches > 1)) {

    Rprintf("* generated a grand total of %d samples.\n", n);
    Rprintf("  > event matches %d samples out of %d (p = %.7g).\n",
      acc.matching, acc.total, result);

  }/*THEN*/

  UNPROTECT(nprot);

  return result;

}/*LS_LOOP*/

/* conditional probability query with forward/logic sampling. */
SEXP logic_sampling(SEXP fitted, SEXP event, SEXP evidence, SEXP n, SEXP batch,
    SEXP eval_env, SEXP debug) {

cpq_query q;
cpq_scratch sc = { 0 };
double result;

  /* the query to answer. */
  q.fitted = fitted;
  q.target = event;
  q.evidence = evidence;
  q.eval_env = eval_env;

  /* the working memory; generation is unconditional, so fixed and colmap are
   * left NULL. */
  sc.bn = fitted_network_from_SEXP(fitted);

  result = ls_loop(&q, &sc, INT(n), INT(batch), isTRUE(debug));
  free_scratch(&sc);

  return ScalarReal(result);

}/*LOGIC_SAMPLING*/

/* -- logic sampling (distribution) ----------------------------------------- */

/* the growing result: target columns, over-allocated, trimmed at end. */
typedef struct {

  SEXP frame;          /* over-allocated result (target columns, n rows). */
  int ncols;           /* number of target columns. */
  int kept;            /* number of matching rows kept so far. */

} builder;

/* process one batch of m observations: generate them into the reused buffer,
 * evaluate the evidence, and copy the matching rows into the accumulator. */
static void ld_batch(const cpq_query *q, cpq_scratch *sc, SEXP buffer, int m,
    builder *acc, int debug) {

int i = 0, oj = 0, len = 0, nprot = 0, *ev = NULL;
int filter_count = 0;
int evidence_true = isTRUE(q->evidence);
tabular dt = { 0 };
SEXP rsxp = R_NilValue;

  /* refill the reused buffer in place (c_rbn_master brackets the RNG). */
  dt = tabular_from_SEXP(buffer, 0, 0, R_NilValue);
  c_rbn_master(sc->bn, dt, sc->fixed, FALSE);
  FreeTAB(dt);

  if (debug)
    Rprintf("* generated %d samples from the bayesian network.\n", m);

  /* evaluate the evidence. */
  if (!evidence_true) {

    SEXP env = make_data_env(buffer, q->eval_env);
    nprot++;
    PROTECT(rsxp = eval(q->evidence, env));
    nprot++;
    if (TYPEOF(rsxp) != LGLSXP)
      error("evidence must evaluate to a logical vector.");
    len = length(rsxp);
    if ((len != 1) && (len != m))
      error("logical vector for evidence is of length %d instead of %d.",
        len, m);
    ev = LOGICAL(rsxp);

  }/*THEN*/

  /* copy the matching rows into the accumulator, recycling a length-one
   * evidence vector as R does when subsetting. */
  for (i = 0; i < m; i++) {

    int keep = evidence_true;

    if (!evidence_true) {

      int ri = ev[(len == 1) ? 0 : i];
      keep = (ri != NA_LOGICAL) && ri;

    }/*THEN*/

    if (keep) {

      for (oj = 0; oj < acc->ncols; oj++) {

        SEXP bcol = VECTOR_ELT(buffer, sc->colmap[oj]);
        SEXP ocol = VECTOR_ELT(acc->frame, oj);

        if (TYPEOF(ocol) == REALSXP)
          REAL(ocol)[acc->kept] = REAL(bcol)[i];
        else
          INTEGER(ocol)[acc->kept] = INTEGER(bcol)[i];

      }/*FOR*/

      acc->kept++;

    }/*THEN*/

  }/*FOR*/

  if (nprot > 0)
    UNPROTECT(nprot);

  if (debug) {

    if (!evidence_true) {

      /* how.many() counts over the (possibly length-one) evidence vector. */
      for (i = 0; i < len; i++)
        if ((ev[i] != NA_LOGICAL) && ev[i])
          filter_count++;

      Rprintf("  > evidence matches %d samples out of %d (p = %.7g).\n",
        filter_count, m, (double)filter_count / m);

    }/*THEN*/
    else
      Rprintf("  > evidence matches %d samples out of %d (p = 1).\n", m, m);

  }/*THEN*/

}/*LD_BATCH*/

static SEXP ld_loop(const cpq_query *q, cpq_scratch *sc, int n, int batch,
    int debug) {

int b = 0, i = 0, oj = 0, nbatches = n / batch, last = n % batch, nprot = 0;
builder acc;
SEXP buffer_full = R_NilValue, buffer_last = R_NilValue;
SEXP out_full, names, result;

  /* the growing result: target columns, none kept yet. */
  acc.frame = R_NilValue;
  acc.ncols = length(q->target);
  acc.kept = 0;

  /* over-allocated accumulator: the target columns at n rows. */
  PROTECT(out_full = fit2df(q->fitted, n));
  nprot++;
  PROTECT(acc.frame = c_dataframe_column(out_full, q->target, FALSE, TRUE));
  nprot++;

  /* map each accumulator column to its column in the generated data. */
  names = getAttrib(out_full, R_NamesSymbol);
  for (oj = 0; oj < acc.ncols; oj++)
    for (i = 0; i < length(names); i++)
      if (strcmp(CHAR(STRING_ELT(names, i)),
            CHAR(STRING_ELT(q->target, oj))) == 0) {

        sc->colmap[oj] = i;
        break;

      }/*THEN*/

  /* reuse one buffer for the full-size batches, and a smaller one for the
   * final partial batch, so each batch draws exactly its number of rows. */
  if (nbatches > 0) {

    PROTECT(buffer_full = fit2df(q->fitted, batch));
    nprot++;

    for (b = 0; b < nbatches; b++)
      ld_batch(q, sc, buffer_full, batch, &acc, debug);

  }/*THEN*/

  if (last > 0) {

    PROTECT(buffer_last = fit2df(q->fitted, last));
    nprot++;
    ld_batch(q, sc, buffer_last, last, &acc, debug);

  }/*THEN*/

  /* trim the accumulator to the rows actually kept. */
  PROTECT(result = allocVector(VECSXP, acc.ncols));
  nprot++;

  for (oj = 0; oj < acc.ncols; oj++) {

    SEXP ocol = VECTOR_ELT(acc.frame, oj), tcol;

    PROTECT(tcol = allocVector(TYPEOF(ocol), acc.kept));
    if (TYPEOF(ocol) == REALSXP)
      for (i = 0; i < acc.kept; i++)
        REAL(tcol)[i] = REAL(ocol)[i];
    else
      for (i = 0; i < acc.kept; i++)
        INTEGER(tcol)[i] = INTEGER(ocol)[i];
    copyMostAttrib(ocol, tcol);
    SET_VECTOR_ELT(result, oj, tcol);
    UNPROTECT(1);

  }/*FOR*/

  setAttrib(result, R_NamesSymbol, q->target);
  minimal_data_frame(result, acc.kept);
  setAttrib(result, BN_MethodSymbol, mkString("ls"));
  setAttrib(result, R_ClassSymbol, mkStringVec(2, "bn.cpdist", "data.frame"));

  if (debug && (nbatches > 1))
    Rprintf("* generated a grand total of %d samples.\n", n);

  UNPROTECT(nprot);

  return result;

}/*LD_LOOP*/

/* sample from a conditional distribution with forward/logic sampling. */
SEXP logic_distribution(SEXP fitted, SEXP nodes, SEXP evidence, SEXP n,
    SEXP batch, SEXP eval_env, SEXP debug) {

cpq_query q;
cpq_scratch sc = { 0 };
SEXP result;

  /* the query to answer. */
  q.fitted = fitted;
  q.target = nodes;
  q.evidence = evidence;
  q.eval_env = eval_env;

  /* the working memory: network and column map; generation is unconditional,
   * so fixed is left NULL. */
  sc.bn = fitted_network_from_SEXP(fitted);
  sc.colmap = Calloc1D(length(nodes), sizeof(int));

  result = ld_loop(&q, &sc, INT(n), INT(batch), isTRUE(debug));
  free_scratch(&sc);

  return result;

}/*LOGIC_DISTRIBUTION*/
