#include "../../include/rcore.h"
#include "../../core/allocations.h"
#include "../../include/globals.h"
#include "../../minimal/common.h"
#include "../../minimal/strings.h"
#include "../../scores/scores.h"

/* build the parents (transposed) sparse adjacency matrix of the DAG obtained by
 * applying an arc operation to the base network, so that the perturbed scores
 * can read the updated parents from it. */
static sparse_amat score_delta_sam(SEXP network, SEXP arc, const char *op) {

sparse_amat sam = { 0 };
int base_narcs = 0, nnodes = 0, narcs = 0, cap = 0, fi = 0, ti = 0;
int *from = NULL, *to = NULL;
bool remove_arc = (strcmp(op, "drop") == 0) || (strcmp(op, "reverse") == 0);
SEXP arcs, nodes, arc_ids, ends;

  arcs = getListElement(network, "arcs");
  PROTECT(nodes = getAttrib(getListElement(network, "nodes"), R_NamesSymbol));
  nnodes = length(nodes);
  base_narcs = length(arcs) / 2;

  /* map the base and the operated arc endpoints to node indexes. */
  PROTECT(arc_ids = match(nodes, arcs, 0));
  PROTECT(ends = match(nodes, arc, 0));
  fi = INTEGER(ends)[0] - 1;
  ti = INTEGER(ends)[1] - 1;

  /* set and reverse add at most one arc to the base structure. */
  cap = base_narcs + 1;
  from = Calloc1D(cap, sizeof(int));
  to = Calloc1D(cap, sizeof(int));

  /* copy the base arcs, dropping the operated arc for drop and reverse. */
  for (int k = 0; k < base_narcs; k++) {

    int f = INTEGER(arc_ids)[k] - 1, t = INTEGER(arc_ids)[k + base_narcs] - 1;

    if (remove_arc && (f == fi) && (t == ti))
      continue;

    from[narcs] = f;
    to[narcs++] = t;

  }/*FOR*/

  /* add the arc (set) or its reverse (reverse). */
  if (strcmp(op, "set") == 0) {

    from[narcs] = fi;
    to[narcs++] = ti;

  }/*THEN*/
  else if (strcmp(op, "reverse") == 0) {

    from[narcs] = ti;
    to[narcs++] = fi;

  }/*THEN*/

  /* transposing swaps the endpoints, turning children into parents. */
  sam = new_sparse_amat(to, from, narcs, nnodes);
  sam.labels = Calloc1D(nnodes, sizeof(char *));
  for (int k = 0; k < nnodes; k++)
    sam.labels[k] = (char *) CHAR(STRING_ELT(nodes, k));

  Free1D(from);
  Free1D(to);
  UNPROTECT(3);

  return sam;

}/*SCORE_DELTA_SAM*/

/* like score_delta_sam() but reading the base arcs off a base parents CSR
 * snapshot instead of the SEXP network, so the perturbed scores can be computed
 * without the network SEXP. the result shares the base labels. */
static sparse_amat score_delta_sam_from_base(sparse_amat base, int fi, int ti,
    const char *op) {

sparse_amat sam = { 0 };
int nnodes = base.dim, narcs = 0, cap = base.narcs + 1;
int *from = NULL, *to = NULL;
bool remove_arc = (strcmp(op, "drop") == 0) || (strcmp(op, "reverse") == 0);

  from = Calloc1D(cap, sizeof(int));
  to = Calloc1D(cap, sizeof(int));

  /* read the base arcs off the parents CSR (row = child, colidx = parents),
   * dropping the operated arc for drop and reverse. */
  for (int r = 0; r < nnodes; r++) {

    for (int p = base.rowptr[r]; p < base.rowptr[r + 1]; p++) {

      int f = base.colidx[p], t = r;

      if (remove_arc && (f == fi) && (t == ti))
        continue;

      from[narcs] = f;
      to[narcs++] = t;

    }/*FOR*/

  }/*FOR*/

  /* add the arc (set) or its reverse (reverse). */
  if (strcmp(op, "set") == 0) {

    from[narcs] = fi;
    to[narcs++] = ti;

  }/*THEN*/
  else if (strcmp(op, "reverse") == 0) {

    from[narcs] = ti;
    to[narcs++] = fi;

  }/*THEN*/

  /* transposing swaps the endpoints, turning children into parents. */
  sam = new_sparse_amat(to, from, narcs, nnodes);
  sam.labels = Calloc1D(nnodes, sizeof(char *));
  for (int k = 0; k < nnodes; k++)
    sam.labels[k] = base.labels[k];

  Free1D(from);
  Free1D(to);

  return sam;

}/*SCORE_DELTA_SAM_FROM_BASE*/

double robust_score_difference(double old1, double old2, double new1,
    double new2) {

double delta = 0, new_sum = 0, old_sum = 0;

  /* sum old and new score components. */
  new_sum = new1 + new2;
  old_sum = old1 + old2;

  /* compare the network scores, minus numeric tolerance for better score
   * equivalence detection. */
  if (fabs(new_sum - old_sum) < MACHINE_TOL)
    delta = 0;
  else
    delta = new_sum - old_sum;

  /* catch: the difference between two -Inf scores must be -Inf and not NaN,
   * so that we can safely test it and reject the change; and if the old score
   * is -Inf just look at the new score. */
  if (old_sum == R_NegInf)
    delta = (new_sum != R_NegInf) ? new_sum : R_NegInf;

  return delta;

}/*ROBUST_SCORE_DIFFERENCE*/

SEXP score_delta_decomposable(SEXP arc, SEXP network, SEXP data, SEXP score,
    SEXP score_delta, SEXP reference_score, SEXP op, SEXP extra) {

int *t = NULL;
const char *o = CHAR(STRING_ELT(op, 0));
double diff = 0, *new = NULL, *old = NULL;
sparse_amat sam = { 0 };
SEXP delta, new_score, try, to_update, reference_names;

  /* the affected nodes are the arc head, plus the tail when reversing. */
  if (strcmp(o, "reverse") == 0) {

    PROTECT(to_update = allocVector(STRSXP, 2));
    SET_STRING_ELT(to_update, 0, STRING_ELT(arc, 0));
    SET_STRING_ELT(to_update, 1, STRING_ELT(arc, 1));

  }/*THEN*/
  else {

    PROTECT(to_update = allocVector(STRSXP, 1));
    SET_STRING_ELT(to_update, 0, STRING_ELT(arc, 1));

  }/*ELSE*/

  /* compute the updated scores with the perturbed structure. */
  PROTECT(new_score = allocVector(REALSXP, length(to_update)));
  new = REAL(new_score);
  memset(new, '\0', length(new_score) * sizeof(double));
  sam = score_delta_sam(network, arc, o);
  c_per_node_score(sam, data, score, to_update, extra, FALSE, new);
  FreeSparseAMAT(sam);
  /* update the test counter. */
  test_counter += length(new_score);

  /* get the corresponding components from the old score. */
  PROTECT(reference_names = getAttrib(reference_score, R_NamesSymbol));
  PROTECT(try = match(reference_names, to_update, 0));
  t = INTEGER(try);
  old = REAL(reference_score);

  if (length(new_score) == 1)
    diff = robust_score_difference(old[t[0] - 1], 0, new[0], 0);
  else
    diff = robust_score_difference(old[t[0] - 1], old[t[1] - 1], new[0], new[1]);

  /* build the return value. */
  PROTECT(delta = allocVector(VECSXP, 3));
  SET_VECTOR_ELT(delta, 0, ScalarLogical(diff > 0));
  SET_VECTOR_ELT(delta, 1, ScalarReal(diff));
  SET_VECTOR_ELT(delta, 2, new_score);
  setAttrib(delta, R_NamesSymbol, mkStringVec(3, "bool", "delta", "updates"));

  UNPROTECT(5);

  return delta;

}/*SCORE_DELTA_DECOMPOSABLE*/

/* coord-native Castelo-Siebes/marginal-uniform score delta: compute the score
 * difference for applying operation `o` to arc fi -> ti (label pair `arc`),
 * reading the current parents off the base (transposed) CSR `base`. the two
 * recomputed score components are written into `new` (length 2); the base is
 * owned by the caller and is not freed here. */
double score_delta_cs_core(SEXP arc, int fi, int ti, sparse_amat base, SEXP data,
    SEXP score, SEXP reference_score, const char *o, SEXP extra, double *new) {

int *t = NULL;
double retval = 0, *old = NULL, new_prior = 0, old_prior = 0;
sparse_amat sam = { 0 };
SEXP prior, cs_prior, order, try, target, off_target;

  /* get the prior specification. */
  prior = getListElement(extra, "prior");
  cs_prior = getListElement(extra, "beta");
  order = getAttrib(cs_prior, BN_NodesSymbol);

  PROTECT(try = match(order, arc, 0));
  t = INTEGER(try);

  PROTECT(target = allocVector(STRSXP, 1));
  SET_STRING_ELT(target, 0, STRING_ELT(arc, 1));
  PROTECT(off_target = allocVector(STRSXP, 1));
  SET_STRING_ELT(off_target, 0, STRING_ELT(arc, 0));
  /* compute the one score component that needs recomputing. */
  memset(new, '\0', 2 * sizeof(double));
  sam = score_delta_sam_from_base(base, fi, ti, o);
  if (strcmp(o, "reverse") == 0)
    c_per_node_score(sam, data, score, arc, extra, FALSE, new);
  else
    c_per_node_score(sam, data, score, target, extra, FALSE, new + 1);
  /* update the test counter. */
  test_counter++;

  /* the graph prior reads the parents and children off the sparse matrices of
   * the perturbed (sam) and the base structures. */
  if ((t[0] < t[1]) || (strcmp(o, "reverse") == 0)) {

    new_prior = graph_prior_prob(prior, target, cs_prior, sam, FALSE);
    new_prior += graph_prior_prob(prior, off_target, cs_prior, sam, FALSE);
    old_prior = graph_prior_prob(prior, target, cs_prior, base, FALSE);
    old_prior += graph_prior_prob(prior, off_target, cs_prior, base, FALSE);

  }/*THEN*/

  /* when trying to add a new arc, make sure the prior from the reference
   * score does not cover the reverse arc (already present as target ->
   * off_target in the base structure). */
  if (strcmp(o, "set") == 0) {

    if (has_arc(base, fi, ti)) {

      sparse_amat rev_sam = { 0 };

      old_prior -= graph_prior_prob(prior, target, cs_prior, base, FALSE);
      old_prior -= graph_prior_prob(prior, off_target, cs_prior, base, FALSE);

      rev_sam = score_delta_sam_from_base(base, ti, fi, "drop");
      old_prior += graph_prior_prob(prior, target, cs_prior, rev_sam, FALSE);
      old_prior += graph_prior_prob(prior, off_target, cs_prior, rev_sam, FALSE);
      FreeSparseAMAT(rev_sam);

    }/*THEN*/

  }/*THEN*/

  FreeSparseAMAT(sam);

  old = REAL(reference_score);
  if (strcmp(o, "reverse") != 0)
     new[0] = old[t[0] - 1] - old_prior + new_prior;
  else
     new[0] -= - old_prior + new_prior;

  retval = robust_score_difference(old[t[0] - 1], old[t[1] - 1], new[0], new[1]);

  UNPROTECT(3);

  return retval;

}/*SCORE_DELTA_CS_CORE*/

SEXP score_delta_cs(SEXP arc, SEXP network, SEXP data, SEXP score,
    SEXP score_delta, SEXP reference_score, SEXP op, SEXP extra) {

const char *o = CHAR(STRING_ELT(op, 0));
int fi = 0, ti = 0;
double retval = 0, *new = NULL;
sparse_amat base = { 0 };
SEXP nodes, ends, delta, new_score;

  /* the base parents (transposed) CSR of the current network, plus the arc
   * endpoints in its node order. */
  base = sparse_amat_from_SEXP(network, TRUE);
  PROTECT(nodes = getAttrib(getListElement(network, "nodes"), R_NamesSymbol));
  PROTECT(ends = match(nodes, arc, 0));
  fi = INTEGER(ends)[0] - 1;
  ti = INTEGER(ends)[1] - 1;

  PROTECT(new_score = allocVector(REALSXP, 2));
  new = REAL(new_score);
  retval = score_delta_cs_core(arc, fi, ti, base, data, score, reference_score,
             o, extra, new);

  FreeSparseAMAT(base);

  /* build the return value. */
  PROTECT(delta = allocVector(VECSXP, 3));
  SET_VECTOR_ELT(delta, 0, ScalarLogical(retval > 0));
  SET_VECTOR_ELT(delta, 1, ScalarReal(retval));
  SET_VECTOR_ELT(delta, 2, new_score);
  setAttrib(delta, R_NamesSymbol, mkStringVec(3, "bool", "delta", "updates"));

  UNPROTECT(4);

  return delta;

}/*SCORE_DELTA_CS*/

SEXP score_delta_monolithic(SEXP arc, SEXP network, SEXP data, SEXP score,
    SEXP score_delta, SEXP reference_score, SEXP op, SEXP extra) {

  switch(score_to_enum(CHAR(STRING_ELT(score, 0)))) {

    case BDE:
    case BDS:
    case BDJ:
    case MBDE:
    case BGE:

      switch(gprior_to_enum(CHAR(STRING_ELT(getListElement(extra, "prior"), 0)))) {

        case CS:
        case MU:
          return score_delta_cs(arc, network, data, score, score_delta,
                   reference_score, op, extra);

        default:
          error("uknown prior in monolithic score function.");

      }/*SWITCH*/


    default:
      error("unknown monolithic score function.");

  }/*SWITCH*/

  return R_NilValue;

}/*SCORE_DELTA_MONOLITHIC*/

SEXP score_delta(SEXP arc, SEXP network, SEXP data, SEXP score,
    SEXP score_delta, SEXP reference_score, SEXP op, SEXP extra, SEXP decomposable) {

  if (isTRUE(decomposable)) {

    return score_delta_decomposable(arc, network, data, score, score_delta,
             reference_score, op, extra);

  }/*THEN*/
  else {

    return score_delta_monolithic(arc, network, data, score, score_delta,
             reference_score, op, extra);

  }/*ELSE*/

}/*SCORE_DELTA*/
