#include "../../include/rcore.h"
#include "../../core/allocations.h"
#include "../../core/sort.h"
#include "../../include/globals.h"
#include "../../include/graph.h"
#include "../../include/learning.h"
#include "../../math/linear.algebra.h"
#include "../../math/sparse.amat.h"
#include "../../minimal/strings.h"
#include "search.state.h"

void tabu_add(double *cache_value, sparse_amat sam, int *b, int nbl, double *np,
    double *mp, SEXP bestop, SEXP nodes, int *nnodes, int *from, int *to,
    double *max, search_state *st, int cur, int *path, int *scratch,
    bool debugging);
void tabu_del(double *cache_value, int *w, int nwl, sparse_amat sam, SEXP bestop,
    SEXP nodes, int *nnodes, int *from, int *to, double *max, search_state *st,
    int cur, bool debugging);
void tabu_rev(double *cache_value, int *b, int nbl, sparse_amat sam, SEXP bestop,
    SEXP nodes, int *nnodes, int *from, int *to, double *max, int *update,
    search_state *st, int cur, double *mp, double *np, int *path, int *scratch,
    bool debugging);

/* can the arc i -> j be added to the network? (it must not already be present in
 * either direction, its target must be below the parent limit, and it must not
 * be blacklisted). This is the single source of truth for arc addability, shared
 * with the batch enumerator hc_to_be_added(). */
bool is_addable(sparse_amat sam, int i, int j, double *np, double *mp,
    int *b, int nbl, int nnodes) {

  /* test the O(1) gates (self-loop, parent limit) before the has_arc() binary
   * searches, so a saturated head short-circuits without them. */
  return (i != j) && (np[j] < mp[j]) && !has_arc(sam, i, j) &&
         !has_arc(sam, j, i) && !coord_set_member(b, nbl, CMC(i, j, nnodes));

}/*IS_ADDABLE*/

/* a single step of tabu search (one arc addition/removal/reversal, plus
 * lookup in the tabu list). */
SEXP tabu_step(SEXP nodes, SEXP reference,
    SEXP whitelist, SEXP blacklist, SEXP state, SEXP current, SEXP baseline,
    SEXP nparents, SEXP maxp, SEXP debug) {

int nnodes = length(nodes), narcs = 0, i = 0, j = 0, p = 0, cur = 0;
int nwl = length(whitelist), nbl = length(blacklist);
int *w = NULL, *b = NULL, *coords = NULL;
int counter = 0, update = 1, from = 0, to = 0;
int *path = NULL, *scratch = NULL;
double *cache_value = NULL, max = NUM(baseline);
double *mp = REAL(maxp), *np = REAL(nparents);
bool debugging = isTRUE(debug);
search_state *st = search_state_from_SEXP(state);
sparse_amat sam = { 0 };
SEXP bestop;

  /* allocate and initialize the return value (use FALSE as a canary value); the
   * two trailing slots carry the score deltas of the chosen move, so the caller
   * can test them without reading the cache. */
  PROTECT(bestop = allocVector(VECSXP, 5));
  setAttrib(bestop, R_NamesSymbol,
    mkStringVec(5, "op", "from", "to", "to.delta", "from.delta"));

  /* allocate and initialize a dummy FALSE object. */
  SET_VECTOR_ELT(bestop, 0, ScalarLogical(FALSE));

  /* allocate the depth-first search buffers for the cycle checks. */
  path = Calloc1D(nnodes, sizeof(int));
  scratch = Calloc1D(nnodes, sizeof(int));

  /* the current network and its sorted coordinate set come from the state;
   * candidate networks differ from it by one coordinate, which the tabu-list
   * lookups compare against without a dense matrix. */
  narcs = st->narcs;
  coords = st->coords;
  sam = sparse_amat_from_coords(coords, narcs, nodes, FALSE);

  /* save pointers to the numeric/integer matrices. */
  cache_value = st->cache;
  w = INTEGER(whitelist);
  b = INTEGER(blacklist);
  cur = asInteger(current);

  if (debugging) {

     /* count how may arcs are to be tested: the arcs that can be added. */
     for (i = 0, counter = 0; i < nnodes; i++)
       for (j = 0; j < nnodes; j++)
         counter += is_addable(sam, i, j, np, mp, b, nbl, nnodes);

     Rprintf("----------------------------------------------------------------\n");
     Rprintf("* trying to add one of %d arcs.\n", counter);

  }/*THEN*/

  /* test neighbours by arc addition. */
  tabu_add(cache_value, sam, b, nbl, np, mp, bestop, nodes, &nnodes, &from, &to,
    &max, st, cur, path, scratch, debugging);

  if (debugging) {

     /* count how may arcs are to be tested: the arcs that are not whitelisted. */
     for (i = 0, counter = 0; i < narcs; i++)
       counter += (1 - coord_set_member(w, nwl, coords[i]));

     Rprintf("----------------------------------------------------------------\n");
     Rprintf("* trying to remove one of %d arcs.\n", counter);

  }/*THEN*/

  /* test neighbours by arc deletion. */
  tabu_del(cache_value, w, nwl, sam, bestop, nodes, &nnodes, &from, &to, &max,
    st, cur, debugging);

  if (debugging) {

     /* count how may arcs are to be tested: the arcs whose reverse is not
      * blacklisted. */
     for (i = 0, counter = 0; i < nnodes; i++)
       for (p = sam.rowptr[i]; p < sam.rowptr[i + 1]; p++) {

         j = sam.colidx[p];
         counter += (1 - coord_set_member(b, nbl, CMC(j, i, nnodes)));

       }/*FOR*/

     Rprintf("----------------------------------------------------------------\n");
     Rprintf("* trying to reverse one of %d arcs.\n", counter);

  }/*THEN*/

  /* test neighbours by arc reversal. */
  tabu_rev(cache_value, b, nbl, sam, bestop, nodes, &nnodes, &from, &to, &max,
    &update, st, cur, mp, np, path, scratch, debugging);

  /* update the reference scores. */
  REAL(reference)[to] += cache_value[CMC(from, to, nnodes)];
  if (update == 2)
    REAL(reference)[from] += cache_value[CMC(to, from, nnodes)];

  /* return the chosen move's score deltas for the caller's singular-node reset;
   * the reverse-only "from" delta is left as NA otherwise. */
  SET_VECTOR_ELT(bestop, 3, ScalarReal(cache_value[CMC(from, to, nnodes)]));
  SET_VECTOR_ELT(bestop, 4,
    ScalarReal((update == 2) ? cache_value[CMC(to, from, nnodes)] : NA_REAL));

  FreeSparseAMAT(sam);
  Free1D(path);
  Free1D(scratch);

  UNPROTECT(1);

  return bestop;

}/*TABU_STEP*/

/* one iteration of the search loop: sync the current network into the state,
 * (optionally) store its tabu hash, fill the score cache and take the best step,
 * all in a single .Call(). This just orchestrates the existing pieces in the same
 * order the R loop used to call them in, so their behaviour and debug output are
 * unchanged. */
SEXP tabu_search_step(SEXP state, SEXP arcs, SEXP nodes, SEXP current, SEXP hash,
    SEXP data, SEXP score, SEXP extra, SEXP reference, SEXP equivalence,
    SEXP decomposability, SEXP updated, SEXP whitelist, SEXP blacklist,
    SEXP nparents, SEXP maxp, SEXP baseline, SEXP debug) {

SEXP bestop;

  /* sync the current network (as a coordinate set) into the search state. */
  search_set_network(state, arcs, nodes);

  /* store the hash of the current network into the tabu ring (tabu search
   * only). */
  if (isTRUE(hash))
    search_hash(state, current);

  /* fill the score-delta cache held in the state (modified in place). */
  score_cache_fill(nodes, data, score, extra, reference, equivalence,
    decomposability, updated, state, blacklist, debug);

  /* find and return the best move (the reference score is updated in place). */
  bestop = tabu_step(nodes, reference, whitelist, blacklist, state, current,
             baseline, nparents, maxp, debug);

  return bestop;

}/*TABU_SEARCH_STEP*/

/* try to add an arc to the current network, minding the tabu list. */
void tabu_add(double *cache_value, sparse_amat sam, int *b, int nbl, double *np,
    double *mp, SEXP bestop, SEXP nodes, int *nnodes, int *from, int *to,
    double *max, search_state *st, int cur, int *path, int *scratch,
    bool debugging) {

int i = 0, j = 0, idx = 0;
double temp = 0, tol = MACHINE_TOL;

  for (i = 0; i < *nnodes; i++) {

    for (j = 0; j < *nnodes; j++) {

      /* only consider arcs that can actually be added. */
      if (!is_addable(sam, i, j, np, mp, b, nbl, *nnodes))
        continue;

      /* retrieve the score delta from the cache. */
      temp = cache_value[CMC(i, j, *nnodes)];

      if (debugging) {

        Rprintf("  > trying to add %s -> %s.\n", NODE(i), NODE(j));
        Rprintf("    > delta between scores for nodes %s %s is %lf.\n",
          NODE(i), NODE(j), temp);

      }/*THEN*/

      /* this score delta is the best one at the moment, so add the arc if it
       * does not introduce cycles in the graph and does not match any of the
       * networks stored in the tabu list. */
      if (temp - *max > tol) {

        if (c_has_path_sparse(sam, j, i, FALSE, path, scratch)) {

          if (debugging)
            Rprintf("    > not adding, introduces cycles in the graph.\n");

          continue;

        }/*THEN*/

        /* look up the candidate network (this one plus the new arc). */
        idx = tabu_match(st, cur, st->current_hash ^ coord_hash(CMC(i, j, *nnodes)));

        if (idx > 0) {

          if (debugging)
            Rprintf("    > not adding, network matches element %d in the tabu list.\n", idx);

          continue;

        }/*THEN*/

        if (debugging)
          Rprintf("    @ adding %s -> %s.\n", NODE(i), NODE(j));

        /* update the return value. */
        bestop_update(bestop, "set", NODE(i), NODE(j));
        /* store the node indices to update the reference scores. */
        *from = i;
        *to = j;

        /* update the threshold score delta. */
        *max = temp;

      }/*THEN*/

    }/*FOR*/

  }/*FOR*/

}/*TABU_ADD*/

/* try to delete an arc from the current network, minding the tabu list. */
void tabu_del(double *cache_value, int *w, int nwl, sparse_amat sam, SEXP bestop,
    SEXP nodes, int *nnodes, int *from, int *to, double *max, search_state *st,
    int cur, bool debugging) {

int i = 0, j = 0, p = 0, idx = 0;
double temp = 0, tol = MACHINE_TOL;

  /* iterate over the arcs already in the network (the CSR children). */
  for (i = 0; i < *nnodes; i++) {

    for (p = sam.rowptr[i]; p < sam.rowptr[i + 1]; p++) {

      j = sam.colidx[p];

      /* whitelisted arcs are not to be removed, ever. */
      if (coord_set_member(w, nwl, CMC(i, j, *nnodes)))
        continue;

      /* retrieve the score delta from the cache. */
      temp = cache_value[CMC(i, j, *nnodes)];

      if (debugging) {

        Rprintf("  > trying to remove %s -> %s.\n", NODE(i), NODE(j));
        Rprintf("    > delta between scores for nodes %s %s is %lf.\n",
          NODE(i), NODE(j), temp);

      }/*THEN*/

      if (temp - *max > tol) {

        /* look up the candidate network (this one minus the arc). */
        idx = tabu_match(st, cur, st->current_hash ^ coord_hash(CMC(i, j, *nnodes)));

        if (idx > 0) {

          if (debugging > 0)
            Rprintf("    > not removing, network matches element %d in the tabu list.\n", idx);

          continue;

        }/*THEN*/

        if (debugging)
          Rprintf("    @ removing %s -> %s.\n", NODE(i), NODE(j));

        /* update the return value. */
        bestop_update(bestop, "drop", NODE(i), NODE(j));
        /* store the node indices to update the reference scores. */
        *from = i;
        *to = j;

        /* update the threshold score delta. */
        *max = temp;

      }/*THEN*/

    }/*FOR*/

  }/*FOR*/

}/*TABU_DEL*/

/* try to reverse an arc in the current network, minding the tabu list. */
void tabu_rev(double *cache_value, int *b, int nbl, sparse_amat sam, SEXP bestop,
    SEXP nodes, int *nnodes, int *from, int *to, double *max, int *update,
    search_state *st, int cur, double *mp, double *np, int *path, int *scratch,
    bool debugging) {

int i = 0, j = 0, p = 0, idx = 0;
double temp = 0, tol = MACHINE_TOL;

  /* iterate over the arcs already in the network (the CSR children). */
  for (i = 0; i < *nnodes; i++) {

    for (p = sam.rowptr[i]; p < sam.rowptr[i + 1]; p++) {

      j = sam.colidx[p];

      /* don't reverse an arc if the one in the opposite direction is
       * blacklisted, ever. */
      if (coord_set_member(b, nbl, CMC(j, i, *nnodes)))
        continue;

      /* do not reverse an arc if that means violating the limit on the
       * maximum number of parents of the node that gains one. */
      if (np[i] >= mp[i])
        continue;

      /* retrieve the score delta from the cache. */
      temp = cache_value[CMC(i, j, *nnodes)] + cache_value[CMC(j, i, *nnodes)];
      /* nuke small values and negative zeroes. */
      if (fabs(temp) < tol) temp = 0;

      if (debugging) {

        Rprintf("  > trying to reverse %s -> %s.\n", NODE(i), NODE(j));
        Rprintf("    > delta between scores for nodes %s %s is %lf.\n",
          NODE(i), NODE(j), temp);

      }/*THEN*/

      if (temp - *max > tol) {

        if (c_has_path_sparse(sam, i, j, TRUE, path, scratch)) {

          if (debugging)
            Rprintf("    > not reversing, introduces cycles in the graph.\n");

          continue;

        }/*THEN*/

        /* look up the candidate network (this one with the arc reversed). */
        idx = tabu_match(st, cur, st->current_hash ^
                coord_hash(CMC(i, j, *nnodes)) ^ coord_hash(CMC(j, i, *nnodes)));

        if (idx > 0) {

          if (debugging)
            Rprintf("    > not reversing, network matches element %d in the tabu list.\n", idx);

          continue;

        }/*THEN*/

        if (debugging)
          Rprintf("    @ reversing %s -> %s.\n", NODE(i), NODE(j));

        /* update the return value. */
        bestop_update(bestop, "reverse", NODE(i), NODE(j));
        /* store the node indices to update the reference scores. */
        *from = i;
        *to = j;
        /* both nodes' reference scores must be updated. */
        *update = 2;

        /* update the threshold score delta. */
        *max = temp;

      }/*THEN*/

    }/*FOR*/

  }/*FOR*/

}/*TABU_REV*/
