open-hypergraphs 0.3.1

use super::relation::converse;
use crate::array::{Array, ArrayKind, NaturalArray};
use crate::category::Arrow;
use crate::finite_function::FiniteFunction;
use crate::indexed_coproduct::IndexedCoproduct;
use crate::strict::hypergraph::Hypergraph;
use num_traits::{One, Zero};

/// Return the adjacency map for a [`Hypergraph`] `h`.
///
/// If `X` is the finite set of operations in `h`, then `operation_adjacency(h)` computes the
/// indexed coproduct `adjacency : X → X*`, where the list `adjacency(x)` is all operations reachable in
/// a single step from operation `x`.
pub(crate) fn operation_adjacency<K: ArrayKind, O, A>(
    h: &Hypergraph<K, O, A>,
) -> IndexedCoproduct<K, FiniteFunction<K>>
where
    K::Type<K::I>: NaturalArray<K>,
{
    h.t.flatmap(&converse(&h.s))
}

/// Return the node-level adjacency map for a [`Hypergraph`].
///
/// If `W` is the finite set of nodes in `h`, then `node_adjacency(h)` computes the
/// indexed coproduct `adjacency : W → W*`, where the list `adjacency(w)` is all nodes reachable in
/// a single step from node `w`.
pub(crate) fn node_adjacency<K: ArrayKind, O, A>(
    h: &Hypergraph<K, O, A>,
) -> IndexedCoproduct<K, FiniteFunction<K>>
where
    K::Type<K::I>: NaturalArray<K>,
{
    node_adjacency_from_incidence(&h.s, &h.t)
}

/// Return the node-level adjacency map from incidence data.
///
/// If `W` is the finite set of nodes, then the result is an indexed coproduct
/// `adjacency : W → W*`, where `adjacency(w)` is all nodes reachable in a single step from `w`.
pub(crate) fn node_adjacency_from_incidence<K: ArrayKind>(
    s: &IndexedCoproduct<K, FiniteFunction<K>>,
    t: &IndexedCoproduct<K, FiniteFunction<K>>,
) -> IndexedCoproduct<K, FiniteFunction<K>>
where
    K::Type<K::I>: NaturalArray<K>,
{
    converse(s).flatmap(t)
}

/// A kahn-ish algorithm for topological sorting of an adjacency relation, encoded as an
/// [`IndexedCoproduct`] (see [`converse`] for details)
///
pub(crate) fn kahn<K: ArrayKind>(
    adjacency: &IndexedCoproduct<K, FiniteFunction<K>>,
) -> (K::Index, K::Type<K::I>)
where
    K::Type<K::I>: NaturalArray<K>,
{
    // The layering assignment to each node.
    // A mutable array of length n with values in {0..n}
    let mut order: K::Type<K::I> = K::Type::<K::I>::fill(K::I::zero(), adjacency.len());
    // Predicate determining if a node has been visited.
    // 1 = unvisited
    // 0 = visited
    // NOTE: we store this as "NOT visited" so we can efficiently filter using "repeat".
    let mut unvisited: K::Type<K::I> = K::Type::<K::I>::fill(K::I::one(), adjacency.len());
    // Indegree of each node.
    let mut indegree = indegree(adjacency);
    // the set of nodes on the frontier, initialized to those with zero indegree.
    let mut frontier: K::Index = zero(&indegree);

    // Loop until frontier is empty, or at max possible layering depth.
    let mut depth = K::I::zero();

    // Implementation outline:
    // 1. Compute *sparse* relative indegree, which is:
    //      - idxs of reachable nodes
    //      - counts of reachability from a given set
    // 2. Subtract from global indegree array using scatter_sub_assign
    //      - scatter_sub_assign::<K>(&mut indegree.table, &reachable_ix, &reachable_count.table);
    // 3. Compute new frontier:
    //      - Numpy-esque: `reachable_ix[indegree[reachable_ix] == 0 && unvisited[reachable_ix]]`
    while !frontier.is_empty() && depth <= adjacency.len() {
        // Mark nodes in the current frontier as visited
        // unvisited[frontier] = 0;
        unvisited.scatter_assign_constant(&frontier, K::I::zero());
        // Set the order of nodes in the frontier to the current depth.
        // order[frontier] = depth;
        order.scatter_assign_constant(&frontier, depth.clone());

        // For each node, compute the number of incoming edges from nodes in the frontier,
        // and count paths to each.
        let (reachable_ix, reachable_count) = sparse_relative_indegree(
            adjacency,
            &FiniteFunction::new(frontier, adjacency.len()).unwrap(),
        );

        // indegree = indegree - dense_relative_indegree(a, f)
        indegree
            .table
            .as_mut()
            .scatter_sub_assign(&reachable_ix.table, &reachable_count.table);

        // Reachable nodes with zero indegree...
        // frontier = reachable_ix[indegree[reachable_ix] == 0]
        frontier = {
            // *indices* i of reachable_ix such that indegree[reachable_ix[i]] == 0
            let reachable_ix_indegree_zero_ix = indegree
                .table
                .gather(reachable_ix.table.get_range(..))
                .zero();

            // only nodes in reachable_ix with indegree 0
            reachable_ix
                .table
                .gather(reachable_ix_indegree_zero_ix.get_range(..))
        };

        // .. and filter out those which have been visited.
        // frontier = frontier[unvisited[frontier]]
        frontier = filter::<K>(
            &frontier,
            &unvisited.as_ref().gather(frontier.get_range(..)),
        );

        // Increment depth
        depth = depth + K::I::one();
    }

    (order.into(), unvisited)
}

/// Compute indegree of all nodes in a multigraph.
pub(crate) fn indegree<K: ArrayKind>(
    adjacency: &IndexedCoproduct<K, FiniteFunction<K>>,
) -> FiniteFunction<K>
where
    K::Type<K::I>: NaturalArray<K>,
{
    // Indegree is *relative* indegree with respect to all nodes.
    // PERFORMANCE: can compute this more efficiently by just bincounting adjacency directly.
    dense_relative_indegree(adjacency, &FiniteFunction::<K>::identity(adjacency.len()))
}

/// Using the adjacency information in `adjacency`, compute the indegree of all nodes reachable from `f`.
///
/// More formally, define:
///
/// ```text
/// a : Σ_{n ∈ A} s(n) → N  // the adjacency information of each
/// f : K → N               // a subset of `K` nodes
/// ```
///
/// Then `dense_relative_indegree(a, f)` computes the indegree from `f` of all `N` nodes.
///
/// # Returns
///
/// A finite function `N → E+1` denoting indegree of each node in `N` relative to `f`.
pub(crate) fn dense_relative_indegree<K: ArrayKind>(
    adjacency: &IndexedCoproduct<K, FiniteFunction<K>>,
    f: &FiniteFunction<K>,
) -> FiniteFunction<K>
where
    K::Type<K::I>: NaturalArray<K>,
{
    // Must have that the number of nodes `adjacency.len()`
    assert_eq!(adjacency.len(), f.target());

    // Operations reachable from those in the set f.
    let reached = adjacency.indexed_values(f).unwrap();
    // Counts in `table` can be as large as the number of reached incidences, which may exceed
    // `adjacency.len()` when multiplicity is present.
    let target = reached.table.len() + K::I::one();
    let table = (reached.table.as_ref() as &K::Type<K::I>).bincount(adjacency.len());
    FiniteFunction::new(table, target).unwrap()
}

/// Using the adjacency information in `adjacency`, compute the indegree of all nodes reachable from `f`.
///
/// More formally, let:
///
/// - `a : Σ_{n ∈ A} s(n) → N` denote the adjacency information of each
/// - `f : K → N` be a subset of `K` nodes
///
/// Then `sparse_relative_indegree(a, f)` computes:
///
/// - `g : R → N`, the subset of (R)eachable nodes reachable from `f`
/// - `i : R → E+1`, the *indegree* of nodes in `R`.
///
pub(crate) fn sparse_relative_indegree<K: ArrayKind>(
    a: &IndexedCoproduct<K, FiniteFunction<K>>,
    f: &FiniteFunction<K>,
) -> (FiniteFunction<K>, FiniteFunction<K>)
where
    K::Type<K::I>: NaturalArray<K>,
{
    // Must have that the number of nodes `adjacency.len()`
    assert_eq!(a.len(), f.target());

    // Indices of operations reachable from those in the set f.
    // Indices may appear more than once.
    let g = a.indexed_values(f).unwrap();
    let (i, c) = g.table.sparse_bincount();
    // Counts in `c` can be as large as the number of reached incidences, which may exceed `a.len()`
    // when multiplicity is present.
    let target = g.table.len() + K::I::one();

    (
        FiniteFunction::new(i, a.len()).unwrap(),
        FiniteFunction::new(c, target).unwrap(),
    )
}

/// Given:
///
/// - `values : K → N`
/// - `predicate : K → 2`
///
/// Return the subset of `values` for which `predicate(i) = 1`
pub(crate) fn filter<K: ArrayKind>(values: &K::Index, predicate: &K::Index) -> K::Index {
    predicate.repeat(values.get_range(..))
}

/// Given an array of indices `values` in `{0..N}` and a predicate `N → 2`, select select values `i` for
/// which `predicate(i) = 1`.
#[allow(dead_code)]
fn filter_by_dense<K: ArrayKind>(values: &K::Index, predicate: &K::Index) -> K::Index
where
    K::Type<K::I>: NaturalArray<K>,
{
    predicate
        .gather(values.get_range(..))
        .repeat(values.get_range(..))
}

/// Given a FiniteFunction `X → L`, compute its converse,
/// a relation `r : L → X*`, and return the result as an array of arrays,
/// where `r_i` is the list of elements in `X` mapping to i.
pub(crate) fn converse_iter<K: ArrayKind>(
    order: FiniteFunction<K>,
) -> impl Iterator<Item = K::Index>
where
    K::Type<K::I>: NaturalArray<K>,
    K::I: Into<usize>,
{
    let c = converse(&IndexedCoproduct::elements(order));
    c.into_iter().map(|x| x.table)
}

////////////////////////////////////////////////////////////////////////////////
// Array trait helpers

// FiniteFunction helpers
pub(crate) fn zero<K: ArrayKind>(f: &FiniteFunction<K>) -> K::Index
where
    K::Type<K::I>: NaturalArray<K>,
{
    (f.table.as_ref() as &K::Type<K::I>).zero()
}