#include "../include/rcore.h"
#include "../core/allocations.h"
#include "../core/sort.h"
#include "../minimal/common.h"
#include "sparse.amat.h"

/* build a CSR adjacency matrix from an arc set. */
sparse_amat new_sparse_amat(int *from, int *to, int narcs, int nnodes) {

sparse_amat sam = { 0 };
int *pos = NULL;

  sam.dim = nnodes;
  sam.narcs = narcs;

  /* count the children of each node to lay out the per-node offsets. */
  sam.rowptr = Calloc1D(nnodes + 1, sizeof(int));
  for (int k = 0; k < narcs; k++)
    sam.rowptr[from[k] + 1]++;
  for (int i = 0; i < nnodes; i++)
    sam.rowptr[i + 1] += sam.rowptr[i];

  /* place each child into its parent's slice of the children array. */
  sam.colidx = Calloc1D(narcs, sizeof(int));
  pos = Calloc1D(nnodes, sizeof(int));
  for (int i = 0; i < nnodes; i++)
    pos[i] = sam.rowptr[i];
  for (int k = 0; k < narcs; k++)
    sam.colidx[pos[from[k]]++] = to[k];
  Free1D(pos);

  /* sort each node's children by index, to match the order in which a dense
   * adjacency matrix would have been scanned. */
  for (int i = 0; i < nnodes; i++)
    i_sort(sam.colidx + sam.rowptr[i], NULL, sam.rowptr[i + 1] - sam.rowptr[i]);

  return sam;

}/*NEW_SPARSE_AMAT*/

/* build a CSR adjacency matrix from a "bn" object: match its arcs against its own
 * node labels and call new_sparse_amat(), optionally transposing (so that the
 * children matrix becomes the parents matrix), then attach the node labels. */
sparse_amat sparse_amat_from_SEXP(SEXP bn, bool transpose) {

sparse_amat sam = { 0 };
int nnodes = 0, narcs = 0;
int *from = NULL, *to = NULL;
SEXP arcs, nodes, arc_ids;

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

  /* map the arc endpoints to node indexes. */
  PROTECT(arc_ids = match(nodes, arcs, 0));
  from = Calloc1D(narcs, sizeof(int));
  to = Calloc1D(narcs, sizeof(int));
  for (int k = 0; k < narcs; k++) {

    from[k] = INTEGER(arc_ids)[k] - 1;
    to[k] = INTEGER(arc_ids)[k + narcs] - 1;

  }/*FOR*/

  /* transposing swaps the arc endpoints, turning children into parents. */
  if (transpose)
    sam = new_sparse_amat(to, from, narcs, nnodes);
  else
    sam = new_sparse_amat(from, to, narcs, nnodes);

  Free1D(from);
  Free1D(to);

  /* attach the row/column (node) labels. */
  sam.labels = Calloc1D(nnodes, sizeof(char *));
  for (int k = 0; k < nnodes; k++)
    sam.labels[k] = (char *) CHAR(STRING_ELT(nodes, k));

  UNPROTECT(2);

  return sam;

}/*SPARSE_AMAT_FROM_SEXP*/

/* build a CSR children matrix from an arc set, matching the arcs against the
 * given node labels (which fixes the node order used for the coordinates). */
sparse_amat sparse_amat_from_arcs(SEXP arcs, SEXP nodes) {

sparse_amat sam = { 0 };
int narcs = length(arcs) / 2, nnodes = length(nodes);
int *from = NULL, *to = NULL;
SEXP arc_ids;

  /* map the arc endpoints to node indexes. */
  PROTECT(arc_ids = match(nodes, arcs, 0));
  from = Calloc1D(narcs, sizeof(int));
  to = Calloc1D(narcs, sizeof(int));
  for (int k = 0; k < narcs; k++) {

    from[k] = INTEGER(arc_ids)[k] - 1;
    to[k] = INTEGER(arc_ids)[k + narcs] - 1;

  }/*FOR*/

  sam = new_sparse_amat(from, to, narcs, nnodes);

  Free1D(from);
  Free1D(to);
  UNPROTECT(1);

  return sam;

}/*SPARSE_AMAT_FROM_ARCS*/

/* build a CSR adjacency matrix from a sorted set of flat CMC() arc coordinates
 * (from + to * nnodes), matching the labels against `nodes`, optionally
 * transposing (children -> parents). This is the inverse of sparse_amat_coords(). */
sparse_amat sparse_amat_from_coords(int *coords, int narcs, SEXP nodes,
    bool transpose) {

sparse_amat sam = { 0 };
int nnodes = length(nodes);
int *from = NULL, *to = NULL;

  from = Calloc1D(narcs, sizeof(int));
  to = Calloc1D(narcs, sizeof(int));
  for (int k = 0; k < narcs; k++) {

    from[k] = coords[k] % nnodes;
    to[k] = coords[k] / nnodes;

  }/*FOR*/

  /* transposing swaps the arc endpoints, turning children into parents. */
  if (transpose)
    sam = new_sparse_amat(to, from, narcs, nnodes);
  else
    sam = new_sparse_amat(from, to, narcs, nnodes);

  Free1D(from);
  Free1D(to);

  /* attach the row/column (node) labels. */
  sam.labels = Calloc1D(nnodes, sizeof(char *));
  for (int k = 0; k < nnodes; k++)
    sam.labels[k] = (char *) CHAR(STRING_ELT(nodes, k));

  return sam;

}/*SPARSE_AMAT_FROM_COORDS*/

/* read off the arcs of a CSR matrix as the sorted set of their flat CMC()
 * coordinates (from + to * dim); `coords` must have room for narcs entries. */
void sparse_amat_coords(sparse_amat sam, int *coords) {

int i = 0, p = 0, k = 0;

  for (i = 0; i < sam.dim; i++)
    for (p = sam.rowptr[i]; p < sam.rowptr[i + 1]; p++)
      coords[k++] = i + sam.colidx[p] * sam.dim;

  i_sort(coords, NULL, sam.narcs);

}/*SPARSE_AMAT_COORDS*/

/* is there an arc from -> to in the CSR children matrix? (binary search in the
 * sorted children of `from`). */
bool has_arc(sparse_amat sam, int from, int to) {

int lo = sam.rowptr[from], hi = sam.rowptr[from + 1] - 1, mid = 0;

  while (lo <= hi) {

    mid = lo + (hi - lo) / 2;

    if (sam.colidx[mid] == to)
      return true;
    else if (sam.colidx[mid] < to)
      lo = mid + 1;
    else
      hi = mid - 1;

  }/*WHILE*/

  return false;

}/*HAS_ARC*/

/* write into `cand` the CSR obtained from `base` by adding (add = true) or
 * removing (add = false) the value `val` in row `row` -- i.e. one arc operation
 * on the DAG. `cand` must be a reusable scratch with rowptr of length dim + 1,
 * colidx of capacity at least base.narcs + 1, and its labels shared with `base`.
 * The base row is sorted, so the spliced row stays sorted without re-sorting. */
void sparse_amat_candidate(sparse_amat base, sparse_amat *cand, int row, int val,
    bool add) {

int i = 0, p = 0, k = 0, shift = add ? 1 : -1;
int rs = base.rowptr[row], re = base.rowptr[row + 1];

  (*cand).dim = base.dim;
  (*cand).labels = base.labels;

  /* copy the rows before `row` unchanged. */
  for (p = 0; p < rs; p++)
    (*cand).colidx[k++] = base.colidx[p];

  if (add) {

    /* splice `val` into the sorted row. */
    for (p = rs; (p < re) && (base.colidx[p] < val); p++)
      (*cand).colidx[k++] = base.colidx[p];
    (*cand).colidx[k++] = val;
    for (; p < re; p++)
      (*cand).colidx[k++] = base.colidx[p];

  }/*THEN*/
  else {

    /* drop `val` from the row. */
    for (p = rs; p < re; p++)
      if (base.colidx[p] != val)
        (*cand).colidx[k++] = base.colidx[p];

  }/*ELSE*/

  /* copy the rows after `row` unchanged. */
  for (p = re; p < base.narcs; p++)
    (*cand).colidx[k++] = base.colidx[p];

  (*cand).narcs = k;

  /* the row offsets are unchanged up to `row`, then shifted by one. */
  for (i = 0; i <= row; i++)
    (*cand).rowptr[i] = base.rowptr[i];
  for (i = row + 1; i <= base.dim; i++)
    (*cand).rowptr[i] = base.rowptr[i] + shift;

}/*SPARSE_AMAT_CANDIDATE*/

/* index of a node given its label (linear scan over the labels). */
int sparse_amat_node(sparse_amat sam, const char *name) {

  for (int i = 0; i < sam.dim; i++)
    if (strcmp(sam.labels[i], name) == 0)
      return i;

  return -1;

}/*SPARSE_AMAT_NODE*/

void FreeSparseAMAT(sparse_amat sam) {

  Free1D(sam.rowptr);
  Free1D(sam.colidx);
  Free1D(sam.labels);

}/*FREESPARSEAMAT*/
