
# common steps in constraint-based structure learning: infer arc orientation
# from the skeleton and the data.
learn.arc.directions = function(data, cluster = NULL, local.structure,
    whitelist, blacklist, test, alpha, extra.args = list(), max.sx = ncol(data),
    debug = FALSE) {

  nodes = names(data)

  # build a list of the undirected arcs in the graph.
  arcs = sets2arcs(local.structure, nodes)

  # drop blacklisted arcs from the arc set.
  to.keep = apply(arcs, 1, function(x){ !is.blacklisted(blacklist, x) })
  arcs = arcs[to.keep, , drop = FALSE]

  # 3. [Orient Edges]
  # 3.1 detect v-structures.
  vs = do.call("rbind",
         vstruct.detect(nodes = nodes, arcs = arcs, mb = local.structure,
           data = data, alpha = alpha, extra.args = extra.args, test = test,
           blacklist = blacklist, max.sx = max.sx, debug = debug))
  rownames(vs) = NULL

  if (!is.null(vs)) {

    # 3.2 sort them in p-value order.
    vs = vs[order(vs[,"max_a"], decreasing = FALSE),]
    # 3.3 apply.
    arcs = vstruct.apply(arcs = arcs, vs = vs, nodes = nodes, debug = debug)

  }#THEN

  # 4. propagate directions.
  pdag = list(learning = list(whitelist = whitelist, blacklist = blacklist),
              nodes = structure(rep(0, length(nodes)), names = nodes),
              arcs = arcs)

  return(cpdag.backend(pdag, fix = TRUE, debug = debug)$arcs)

}#LEARN.ARC.DIRECTIONS

# construct a fake markov blanket using all the nodes within distance 2.
fake.markov.blanket = function(learn, target) {

  mb = c(unlist(lapply(learn[[target]]$nbr,
         function(cur) learn[[cur]]$nbr)), learn[[target]]$nbr)
  mb = setdiff(unique(mb), target)

  return(mb)

}#FAKE.MARKOV.BLANKET

# build the neighbourhood of a node from the markov blanket.
neighbour = function(target, mb, data, alpha, extra.args = list(),
    whitelist, blacklist, test, empty.dsep = TRUE, markov = TRUE,
    max.sx = ncol(data), debug = FALSE) {

  # initialize the neighbourhood using the markov blanket.
  candidate.neighbours = mb[[target]]

  # if the markov blanket is empty there's nothing to do.
  if (length(candidate.neighbours) == 0) {

    if (debug) {

      cat("----------------------------------------------------------------\n")
      cat("* markov blanket of", target, "is empty; the neighbourhood as well.\n")

    }#THEN

    return(list(mb = character(0), nbr = character(0)))

  }#THEN

  # whitelisted nodes are included (arc orientation is irrelevant), and
  # blacklisted nodes are removed if both directed arcs are banned and both
  # are not in the whitelist.
  blacklisted = candidate.neighbours[sapply(candidate.neighbours,
        function(y) { is.blacklisted(blacklist, c(target, y), both = TRUE) })]
  whitelisted = candidate.neighbours[sapply(candidate.neighbours,
        function(y) { is.whitelisted(whitelist, c(target, y), either = TRUE) })]

  candidate.neighbours = setdiff(candidate.neighbours, blacklisted)
  candidate.neighbours = union(candidate.neighbours, whitelisted)

  if (debug) {

    cat("----------------------------------------------------------------\n")
    cat("* learning neighbourhood of", target, ".\n")
    cat("  * blacklisted nodes: '", blacklisted, "'\n")
    cat("  * whitelisted nodes: '", whitelisted, "'\n")
    cat("  * starting with neighbourhood: '", candidate.neighbours, "'\n")

  }#THEN

  # nothing much to do, just return.
  if (length(candidate.neighbours) <= 1)
    return(list(mb = mb[[target]], nbr = candidate.neighbours))

  # do not even try to remove whitelisted nodes; on the other hand.
  for (y in setdiff(candidate.neighbours, whitelisted)) {

    if (debug)
      cat("  * checking node", y, "for neighbourhood.\n")

    # choose the smaller set of possible d-separating nodes.
    if (markov)
      dsep.set = smaller(setdiff(mb[[target]], y), setdiff(mb[[y]], target))
    else
      dsep.set = setdiff(mb[[target]], y)

    if (debug)
      cat("    > dsep.set = '", dsep.set, "'\n")

    a = allsubs.test(x = target, y = y, sx = dsep.set,
          min = ifelse(empty.dsep, 0, 1), max = min(length(dsep.set), max.sx),
          data = data, test = test, alpha = alpha, extra.args = extra.args,
          debug = debug)

    # update the neighbourhood.
    if (a["p.value"] > alpha)
      candidate.neighbours = candidate.neighbours[candidate.neighbours != y]

    if (debug)
      cat("    > node", y, "is", ifelse(a["p.value"] > alpha, "not", "still"),
        "a neighbour of", target, ". ( p-value:", a["p.value"], ")\n")

  }#FOR

  return(list(mb = mb[[target]], nbr = candidate.neighbours))

}#NEIGHBOUR

# detect v-structures in the graph.
vstruct.detect = function(nodes, arcs, mb, data, alpha, extra.args = list(),
    test, blacklist, max.sx = ncol(data), debug = FALSE) {

  vstruct.centered.on = function(x, mb, data, dsep.set) {

    if (debug) {

      cat("----------------------------------------------------------------\n")
      cat("* v-structures centered on", x, ".\n")

    }#THEN

    tos = arcs[(arcs[, "to"] == x), "from"]

    if (length(tos) < 2)
      return(NULL)

    # build a list of possible parents for the node x, i.e. all the subsets
    # of size 2 of the nodes connected to x by incoming arcs.
    tos.combs = subsets(tos, 2)
    vs = NULL

    for (j in 1:nrow(tos.combs)) {

      y = tos.combs[j, 1]
      z = tos.combs[j, 2]

      if (debug)
        cat("  * checking", y, "->", x, "<-", z, "\n")

      # check there's no arc from y to z and vice versa.
      if (is.listed(arcs, c(y, z), either = TRUE))
        next
      # check whether this is because the arc is blacklisted; d-separating sets
      # are not learned at all in that case.
      if (is.null(blacklist))
        blacklisted.arc = FALSE
      else
        blacklisted.arc = is.listed(blacklist, c(y, z), both = TRUE)

      # if d-separating sets have been saved during the first part of structure
      # learning, it is possible to detect v-structures without additional tests;
      # if not (because of the algorithm or because of blacklists) we need to
      # test all possible subsets of the intersection of the markov blankets.
      if (!is.null(dsep.set) && !blacklisted.arc) {

        el = dsep.set[[which(sapply(dsep.set, function(x) setequal(x$arc, c(y, z))))]]

        # an unshielded triplet is a v-structure if and only if the central node
        # (x) is not part of the d-separating set of the other nodes (y and z).
        if (x %!in% el$dsep.set) {

          if (debug)
            cat("    @ detected v-structure", y, "->", x, "<-", z, "from d-separating set.\n")

          vs = rbind(vs, data.frame(max_a = el$p.value, y, x, z))

        }#THEN

      }#THEN
      else {

        # choose the smallest of mb(y) - {x,z} and mb(z) - {x,y} to cut down
        # the number of subsets to test.
        sx = smaller(setdiff(mb[[y]][['mb']], c(x, z)),
                     setdiff(mb[[z]][['mb']], c(x, y)))

        if (debug)
          cat("    > chosen d-separating set: '", sx, "'\n")

        a = allsubs.test(x = y, y = z, fixed = x, sx = sx, data = data,
              test = test, extra.args = extra.args, alpha = alpha,
              max = min(max.sx, length(sx)), debug = debug)

        if (a["p.value"] <= alpha) {

          if (debug)
            cat("    @ detected v-structure", y, "->", x, "<-", z, "\n")

          vs = rbind(vs, data.frame(max_a = a["max.p.value"], y, x, z))

        }#THEN

      }#ELSE

    }#FOR

    return(vs)

  }#VSTRUCT.CENTERED.ON

  sapply(nodes, vstruct.centered.on, mb = mb, data = data,
    dsep.set = attr(mb, "dsep.set"), simplify = FALSE)

}#VSTRUCT.DETECT

# include v-structures in the network, setting the corresponding arc directions.
vstruct.apply = function(arcs, vs, nodes, debug = FALSE) {

  if (debug)
    cat("----------------------------------------------------------------\n")

  for (i in seq(nrow(vs))) {

    x = vs[i, "x"]
    y = vs[i, "y"]
    z = vs[i, "z"]
    max_a = vs[i, "max_a"]

    # check whether the network already includes conflicting v-structures.
    if (!(is.listed(arcs, c(y, x)) && is.listed(arcs, c(z, x)))) {

      if (debug)
        cat("* not applying v-structure", y, "->", x, "<-", z, "(", max_a, ")\n")

      warning("vstructure ", y, " -> ", x, " <- ", z, " is not applicable, ",
        "because one or both arcs are oriented in the opposite direction.")

      next

    }#THEN

    # tentatively add the arcs that make up the v-structure.
    temp = set.arc.direction(y, x, arcs)
    temp = set.arc.direction(z, x, temp)

    # check whether the network is acyclic.
    if (!is.acyclic(temp, nodes, directed = TRUE)) {

      if (debug)
        cat("* not applying v-structure", y, "->", x, "<-", z, "(", max_a, ")\n")

      warning("vstructure ", y, " -> ", x, " <- ", z, " is not applicable, ",
        "because one or both arcs introduce cycles in the graph.")

      next

    }#THEN

    if (debug)
      cat("* applying v-structure", y, "->", x, "<-", z, "(", max_a, ")\n")

    # save the updated arc set.
    arcs = temp

  }#FOR

  return(arcs)

}#VSTRUCT.APPLY

# check markov blanket / neighbourhood sets for symmetry, and repair them.
enforce.symmetry = function(bn, mb = FALSE, debug = FALSE) {

  nodes = names(bn)
  n = length(nodes)

  if (mb)
    what = "markov blankets"
  else
    what = "neighbourhood sets"

  if (debug) {

    cat("----------------------------------------------------------------\n")
    cat(paste0("* checking consistency of ", what, ".\n"))
    for (node in nodes)
      cat("  > checking node", node, ".\n")

  }#THEN

  # the set to check is the markov blanket or the neighbours of each node.
  if (mb)
    sets = bn
  else
    sets = lapply(bn, `[[`, "nbr")

  # m[j, i] flags whether node j is in node i's set; folding both directions
  # gives 0 = unrelated, 2 = symmetric, 1 = asymmetric (a one-sided error).
  m = vapply(sets, function(s) as.integer(nodes %in% s), integer(n))
  pair = m + t(m)

  # nothing to fix if every pair is symmetric.
  if (!any(pair == 1L))
    return(bn)

  if (debug)
    for (i in seq_len(n))
      for (j in i:n)
        if (pair[i, j] == 1L)
          cat("@ asymmetry in the", what, "for", nodes[i], "and", nodes[j], ".\n")

  # rebuild keeping only the pairs both nodes agree on.
  fixed = lapply(seq_len(n), function(i) {

    keep = nodes[-i][pair[i, -i] == 2L]
    if (mb)
      keep
    else
      list(mb = bn[[i]][["mb"]], nbr = keep)

  })
  names(fixed) = nodes

  return(fixed)

}#ENFORCE.SYMMETRY

# explore the structure of the network using its arc set.
cache.structure = function(nodes, arcs, amat = NULL, debug = FALSE) {

  # rebuild the adjacency matrix only if it's not available.
  if (is.null(amat))
    amat = arcs2amat(arcs, nodes)

  .Call(call_cache_structure,
        nodes = nodes,
        amat = amat,
        debug = debug)

}#CACHE.STRUCTURE

# explore the structure of the neighbourhood of a target node.
cache.partial.structure = function(nodes, target, arcs, amat = NULL,
    debug = FALSE) {

  # rebuild the adjacency matrix only if it's not available.
  if (is.null(amat))
    amat = arcs2amat(arcs, nodes)

  .Call(call_cache_partial_structure,
        nodes = nodes,
        target = target,
        amat = amat,
        debug = debug)

}#CACHE.PARTIAL.STRUCTURE
