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

/* -- particle weights ------------------------------------------------------ */

/* build the tabular passed to c_lw_weights() from the SEXP data. */
tabular lw_data_from_SEXP(SEXP data, SEXP fitted, SEXP keep) {

tabular dt = { 0 };
SEXP keep2;

  /* the columns to keep are the fixed ones; the names and completeness flags
   * are read from the data's metadata by tabular_from_SEXP(). */
  PROTECT(keep2 = match(keep, getAttrib(fitted, R_NamesSymbol), 0));
  dt = tabular_from_SEXP(data, 0, 0, keep2);
  UNPROTECT(1);

  return dt;

}/*LW_DATA_FROM_SEXP*/

/* compute the weights of the likelihood weighting particles. */
SEXP lw_weights(SEXP fitted, SEXP data, SEXP keep, SEXP debug) {

int n = length(VECTOR_ELT(data, 0));
fitted_bn bn = fitted_network_from_SEXP(fitted);
tabular dt = lw_data_from_SEXP(data, fitted, keep);
SEXP weights;

  PROTECT(weights = allocVector(REALSXP, n));

  c_lw_weights(bn, dt, REAL(weights), isTRUE(debug));

  FreeTAB(dt);
  FreeFittedBN(bn);

  UNPROTECT(1);

  return weights;

}/*LW_WEIGHTS*/

/* -- likelihood weighting (query) ------------------------------------------ */

/* 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. */
  double *w;           /* likelihood weights, one per generated observation. */

} cpq_scratch;

static void free_scratch(cpq_scratch *sc) {

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

}/*FREE_SCRATCH*/

/* the running probability masses. */
typedef struct {

  double total;        /* accumulated probability mass of the evidence. */
  double matching;     /* accumulated mass of the event and the evidence. */

} masses;

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

int i = 0, *r = NULL;
int event_true = isTRUE(q->target), evidence_true = isTRUE(q->evidence);
double *w = sc->w;
long double batch_total = 0, batch_matching = 0;
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 event, unless it is unconditionally TRUE. */
  if (!event_true) {

    SEXP env = make_data_env(buffer, q->eval_env);
    PROTECT(rsxp = eval(q->target, env));
    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);
    r = LOGICAL(rsxp);

  }/*THEN*/

  /* the weights are all one for an unconditional query, the likelihood
   * weights otherwise. */
  if (evidence_true) {

    for (i = 0; i < m; i++)
      w[i] = 1;

  }/*THEN*/
  else {

    tabular wt = { 0 };

    add_simulation_metadata(buffer, m);
    wt = lw_data_from_SEXP(buffer, q->fitted,
           getAttrib(q->evidence, R_NamesSymbol));
    c_lw_weights(sc->bn, wt, w, debug);
    FreeTAB(wt);

  }/*ELSE*/

  /* accumulate the masses: total over non-missing events, matching over hits.
   * sum each batch in long double and round to double before accumulating, to
   * match R's cpe + sum(w[...]) exactly. */
  for (i = 0; i < m; i++) {

    int notna = event_true || (r[i] != NA_LOGICAL);
    int match = event_true || (r[i] != NA_LOGICAL && r[i]);

    if (notna)
      batch_total += w[i];
    if (match)
      batch_matching += w[i];

  }/*FOR*/

  acc->total += (double)batch_total;
  acc->matching += (double)batch_matching;

  if (!event_true)
    UNPROTECT(2);

  if (debug)
    Rprintf("  > event has a probability mass of %.7g out of %.7g.\n",
      acc->matching, acc->total);

}/*WS_BATCH*/

/* the batch loop; a failing predicate longjmps straight out and leaks the C
 * scratch (see free_scratch()). */
static double ws_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;
masses acc = { 0, 0 };
SEXP buffer_full = R_NilValue, buffer_last = R_NilValue;

  /* reuse one buffer for all the full-size batches ... */
  if (nbatches > 0) {

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

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

  }/*THEN*/

  /* ... and a smaller one for the final partial batch, if any, so that each
   * batch draws exactly its number of rows from the RNG. */
  if (last > 0) {

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

  }/*THEN*/

  result = acc.matching / acc.total;

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

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

  }/*THEN*/

  UNPROTECT(nprot);

  return result;

}/*WS_LOOP*/

/* conditional probability query with likelihood weighting. */
SEXP weighting_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: network, fixed evidence, and the weights buffer. */
  sc.bn = fitted_network_from_SEXP(fitted);
  sc.fixed = evidence_from_SEXP(evidence, sc.bn);
  sc.w = Calloc1D(INT(batch), sizeof(double));

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

  return ScalarReal(result);

}/*WEIGHTING_SAMPLING*/

/* -- likelihood weighting (distribution) ----------------------------------- */

/* sample observations with the associated likelihood weights. */
SEXP weighting_distribution(SEXP fitted, SEXP nodes, SEXP n, SEXP fix,
    SEXP debug) {

int nsims = INT(n), max_id = 0;
double *weights = NULL;
fitted_bn bn = fitted_network_from_SEXP(fitted);
fixed_node *fixed = evidence_from_SEXP(fix, bn);
tabular dt = { 0 };
SEXP result, simulation, wgt, from;

  /* allocate the scratch space for the simulation. */
  PROTECT(simulation = fit2df(fitted, nsims));
  dt = tabular_from_SEXP(simulation, 0, 0, R_NilValue);
  /* generate the random observations. */
  c_rbn_master(bn, dt, fixed, FALSE);
  FreeEvidence(fixed, bn.nnodes);
  /* attach the metadata that c_lw_weights() relies on. */
  add_simulation_metadata(simulation, nsims);
  FreeTAB(dt);

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

  /* compute the weights. */
  PROTECT(wgt = allocVector(REALSXP, nsims));
  weights = REAL(wgt);
  PROTECT(from = getAttrib(fix, R_NamesSymbol));
  dt = lw_data_from_SEXP(simulation, fitted, from);
  c_lw_weights(bn, dt, weights, FALSE);
  FreeTAB(dt);
  FreeFittedBN(bn);

  /* if all weights are zero or NA, the evidence is making it impossible to
   * generate a set of valid random observations. */
  max_id = d_which_max(weights, nsims);

  if (max_id == NA_INTEGER)
    error("all weights are NA, the probability of the evidence is impossible "
          "to compute.");
  if (weights[d_which_max(weights, nsims) - 1] == 0)
    error("all weights are zero, the evidence has probability zero.");

  /* prepare the return value with all the attributes. */
  PROTECT(result = c_dataframe_column(simulation, nodes, FALSE, TRUE));
  minimal_data_frame(result, nsims);
  setAttrib(result, BN_WeightsSymbol, wgt);
  setAttrib(result, BN_MethodSymbol, mkString("lw"));
  setAttrib(result, R_ClassSymbol, mkStringVec(2, "bn.cpdist", "data.frame"));

  UNPROTECT(4);

  return result;

}/*WEIGHTING_DISTRIBUTION*/
