//! Directed-acyclic graph implementation.
#![warn(missing_docs)]
use std::{
    borrow::Borrow,
    cmp::Ordering,
    collections::{BTreeMap, BTreeSet, VecDeque},
    fmt,
    ops::{ControlFlow, Deref, Index},
};

/// A node in the graph.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Node<K, V> {
    /// The node value, stored by the user.
    pub value: V,
    /// Nodes depended on.
    pub dependencies: BTreeSet<K>,
    /// Nodes depending on this node.
    pub dependents: BTreeSet<K>,
}

impl<K, V> Node<K, V> {
    fn new(value: V) -> Self {
        Self {
            value,
            dependencies: BTreeSet::new(),
            dependents: BTreeSet::new(),
        }
    }
}

impl<K, V> Borrow<V> for &Node<K, V> {
    fn borrow(&self) -> &V {
        &self.value
    }
}

impl<K, V> Deref for Node<K, V> {
    type Target = V;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

/// A directed acyclic graph.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Dag<K, V> {
    graph: BTreeMap<K, Node<K, V>>,
    tips: BTreeSet<K>,
    roots: BTreeSet<K>,
}

impl<K: Ord + Copy, V> Dag<K, V> {
    /// Create a new empty DAG.
    pub fn new() -> Self {
        Self {
            graph: BTreeMap::new(),
            tips: BTreeSet::new(),
            roots: BTreeSet::new(),
        }
    }

    /// Create a DAG with a root node.
    pub fn root(key: K, value: V) -> Self {
        Self {
            graph: BTreeMap::from_iter([(key, Node::new(value))]),
            tips: BTreeSet::from_iter([key]),
            roots: BTreeSet::from_iter([key]),
        }
    }

    /// Check whether there are any nodes in the graph.
    pub fn is_empty(&self) -> bool {
        self.graph.is_empty()
    }

    /// Return the number of nodes in the graph.
    pub fn len(&self) -> usize {
        self.graph.len()
    }

    /// Add a node to the graph.
    pub fn node(&mut self, key: K, value: V) -> Option<Node<K, V>> {
        self.tips.insert(key);
        self.roots.insert(key);
        self.graph.insert(
            key,
            Node {
                value,
                dependencies: BTreeSet::new(),
                dependents: BTreeSet::new(),
            },
        )
    }

    /// Add a dependency from one node to the other.
    pub fn dependency(&mut self, from: K, to: K) {
        if let Some(node) = self.graph.get_mut(&from) {
            node.dependencies.insert(to);
            self.roots.remove(&from);
        }
        if let Some(node) = self.graph.get_mut(&to) {
            node.dependents.insert(from);
            self.tips.remove(&to);
        }
    }

    /// Check if the graph contains a node.
    pub fn contains(&self, key: &K) -> bool {
        self.graph.contains_key(key)
    }

    /// Get a node.
    pub fn get(&self, key: &K) -> Option<&Node<K, V>> {
        self.graph.get(key)
    }

    /// Check whether there is a dependency between two nodes.
    pub fn has_dependency(&self, from: &K, to: &K) -> bool {
        self.graph
            .get(from)
            .map(|n| n.dependencies.contains(to))
            .unwrap_or_default()
    }

    /// Get the graph's root nodes, ie. nodes which don't depend on other nodes.
    pub fn roots(&self) -> impl Iterator<Item = (&K, &Node<K, V>)> + '_ {
        self.roots
            .iter()
            .filter_map(|k| self.graph.get(k).map(|n| (k, n)))
    }

    /// Get the graph's tip nodes, ie. nodes which aren't depended on by other nodes.
    pub fn tips(&self) -> impl Iterator<Item = (&K, &Node<K, V>)> + '_ {
        self.tips
            .iter()
            .filter_map(|k| self.graph.get(k).map(|n| (k, n)))
    }

    /// Merge a DAG into this one.
    pub fn merge(&mut self, mut other: Self) {
        let Some((root, _)) = other.roots().next() else {
            return;
        };
        let mut visited = BTreeSet::new();
        let mut queue = VecDeque::<K>::from([*root]);

        while let Some(next) = queue.pop_front() {
            if !visited.insert(next) {
                continue;
            }
            if let Some(node) = other.graph.remove(&next) {
                if !self.contains(&next) {
                    self.node(next, node.value);
                }
                for k in &node.dependents {
                    self.dependency(*k, next);
                }
                for k in &node.dependencies {
                    self.dependency(next, *k);
                }
                queue.extend(node.dependents.iter());
            }
        }
    }

    /// Return a topological ordering of the graph's nodes.
    /// Uses a comparison function to sort partially ordered nodes.
    pub fn sorted<F>(&self, mut compare: F) -> Vec<K>
    where
        F: FnMut(&K, &K) -> Ordering,
    {
        let mut order = Vec::new(); // Stores the topological order.
        let mut visited = BTreeSet::new(); // Nodes that have been visited.
        let mut keys = self.graph.keys().collect::<Vec<_>>();

        keys.sort_by(|a, b| compare(a, b));

        for node in keys {
            self.visit(node, &mut visited, &mut order);
        }
        order
    }

    /// Fold over the graph in topological order, pruning branches along the way.
    ///
    /// To continue traversing a branch, return [`ControlFlow::Continue`] from the
    /// filter function. To stop traversal of a branch, return [`ControlFlow::Break`].
    pub fn fold<A, F>(&self, roots: &[K], mut acc: A, mut filter: F) -> A
    where
        F: for<'r> FnMut(A, &'r K, &'r Node<K, V>) -> ControlFlow<A, A>,
    {
        let mut visited = BTreeSet::new();
        let mut queue = VecDeque::<K>::from_iter(roots.iter().cloned());

        while let Some(next) = queue.pop_front() {
            if !visited.insert(next) {
                continue;
            }
            if let Some(node) = self.graph.get(&next) {
                match filter(acc, &next, node) {
                    ControlFlow::Continue(a) => {
                        queue.extend(node.dependents.iter().cloned());
                        acc = a;
                    }
                    ControlFlow::Break(a) => {
                        // When filtering out a node, we filter out all transitive dependents on
                        // that node by adding them to the already visited list.
                        visited.extend(self.descendants_of(node));
                        acc = a;
                    }
                }
            }
        }
        acc
    }

    fn descendants_of(&self, from: &Node<K, V>) -> Vec<K> {
        let mut visited = BTreeSet::new();
        let mut stack = VecDeque::new();
        let mut nodes = Vec::new();

        stack.extend(from.dependents.iter());

        while let Some(key) = stack.pop_front() {
            if let Some(node) = self.graph.get(&key) {
                if visited.insert(key) {
                    nodes.push(key);

                    for &neighbour in &node.dependents {
                        stack.push_back(neighbour);
                    }
                }
            }
        }
        nodes
    }

    /// Add nodes recursively to the topological order, starting from the given node.
    fn visit(&self, key: &K, visited: &mut BTreeSet<K>, order: &mut Vec<K>) {
        if visited.contains(key) {
            return;
        }
        visited.insert(*key);

        // Recursively visit all of the node's dependencies.
        if let Some(node) = self.graph.get(key) {
            for dependency in &node.dependencies {
                self.visit(dependency, visited, order);
            }
        }
        // Add the node to the topological order.
        order.push(*key);
    }
}

impl<K: Ord + Copy + fmt::Debug, V> Index<&K> for Dag<K, V> {
    type Output = Node<K, V>;

    fn index(&self, key: &K) -> &Self::Output {
        self.get(key)
            .unwrap_or_else(|| panic!("Dag::index: node {key:?} not found in graph"))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_len() {
        let mut dag = Dag::new();

        dag.node(0, ());
        dag.node(1, ());
        dag.node(2, ());

        assert_eq!(dag.len(), 3);
    }

    #[test]
    fn test_is_empty() {
        let mut dag = Dag::new();
        assert!(dag.is_empty());

        dag.node(0, ());
        assert!(!dag.is_empty());
    }

    #[test]
    fn test_dependencies() {
        let mut dag = Dag::new();

        dag.node(0, ());
        dag.node(1, ());
        dag.dependency(0, 1);

        assert!(dag.has_dependency(&0, &1));
        assert!(!dag.has_dependency(&1, &0));
    }

    #[test]
    fn test_get() {
        let mut dag = Dag::new();

        dag.node(0, "rad");
        dag.node(1, "dar");

        assert_eq!(dag[&0].value, "rad");
        assert_eq!(dag[&1].value, "dar");
        assert!(dag.get(&2).is_none());
    }

    #[test]
    fn test_cycle() {
        let mut dag = Dag::new();

        dag.node(0, ());
        dag.node(1, ());

        dag.dependency(0, 1);
        dag.dependency(1, 0);

        let sorted = dag.sorted(|a, b| a.cmp(b));
        let expected: &[&[i32]] = &[&[0, 1], &[1, 0]];

        assert!(expected.contains(&sorted.as_slice()));
    }

    #[test]
    fn test_merge_1() {
        let mut a = Dag::new();
        let mut b = Dag::new();
        let mut c = Dag::new();

        a.node(0, ());
        a.node(1, ());
        a.dependency(1, 0);

        b.node(0, ());
        b.node(2, ());
        b.dependency(2, 0);

        c.merge(a);
        c.merge(b);

        assert!(c.get(&0).is_some());
        assert!(c.get(&1).is_some());
        assert!(c.get(&2).is_some());
        assert!(c.has_dependency(&1, &0));
        assert!(c.has_dependency(&2, &0));
    }

    #[test]
    fn test_merge_2() {
        let mut a = Dag::new();
        let mut b = Dag::new();

        a.node(0, ());
        a.node(1, ());
        a.node(2, ());
        a.dependency(1, 0);
        a.dependency(2, 0);

        b.node(0, ());
        b.node(1, ());
        b.node(2, ());
        b.node(3, ());
        b.node(4, ());
        b.dependency(1, 0);
        b.dependency(2, 0);
        b.dependency(3, 0);
        b.dependency(4, 2);

        assert!(a.tips.contains(&2));

        a.merge(b);

        assert!(a.get(&0).is_some());
        assert!(a.get(&1).is_some());
        assert!(a.get(&2).is_some());
        assert!(a.get(&3).is_some());
        assert!(a.get(&4).is_some());
        assert!(a.has_dependency(&4, &2));
        assert!(a.get(&2).unwrap().dependents.contains(&4));
        assert!(a.get(&0).unwrap().dependents.contains(&3));
        assert!(a.tips.contains(&1));
        assert!(!a.tips.contains(&2));
        assert!(a.tips.contains(&3));
        assert!(a.tips.contains(&4));
        assert!(a.roots.contains(&0));
    }

    #[test]
    fn test_diamond() {
        let mut dag = Dag::new();

        dag.node(0, ());
        dag.node(1, ());
        dag.node(2, ());
        dag.node(3, ());

        dag.dependency(1, 0);
        dag.dependency(2, 0);
        dag.dependency(3, 1);
        dag.dependency(3, 2);

        assert_eq!(dag.tips().map(|(k, _)| *k).collect::<Vec<_>>(), vec![3]);
        assert_eq!(dag.roots().map(|(k, _)| *k).collect::<Vec<_>>(), vec![0]);

        // All of the possible sort orders for the above graph.
        let expected: &[&[i32]] = &[&[0, 1, 2, 3], &[0, 2, 1, 3]];
        let actual = dag.sorted(|a, b| a.cmp(b));

        assert!(expected.contains(&actual.as_slice()), "{actual:?}");
    }

    #[test]
    fn test_complex() {
        let mut dag = Dag::new();

        dag.node(0, ());
        dag.node(1, ());
        dag.node(2, ());
        dag.node(3, ());
        dag.node(4, ());
        dag.node(5, ());

        dag.dependency(3, 2);
        dag.dependency(1, 3);
        dag.dependency(2, 5);
        dag.dependency(0, 5);
        dag.dependency(0, 4);
        dag.dependency(1, 4);

        assert_eq!(
            dag.tips().map(|(k, _)| *k).collect::<BTreeSet<_>>(),
            BTreeSet::from_iter([1, 0])
        );
        assert_eq!(
            dag.roots().map(|(k, _)| *k).collect::<BTreeSet<_>>(),
            BTreeSet::from_iter([4, 5])
        );

        // All of the possible sort orders for the above graph.
        let expected = &[
            [4, 5, 0, 2, 3, 1],
            [4, 5, 2, 0, 3, 1],
            [4, 5, 2, 3, 0, 1],
            [4, 5, 2, 3, 1, 0],
            [5, 2, 3, 4, 0, 1],
            [5, 2, 3, 4, 1, 0],
            [5, 2, 4, 0, 3, 1],
            [5, 2, 4, 3, 0, 1],
            [5, 2, 4, 3, 1, 0],
            [5, 4, 0, 2, 3, 1],
            [5, 4, 2, 0, 3, 1],
            [5, 4, 2, 3, 0, 1],
            [5, 4, 2, 3, 1, 0],
        ];
        let mut sorts = BTreeSet::new();
        let mut rng = fastrand::Rng::new();

        while sorts.len() < expected.len() {
            sorts.insert(dag.sorted(|a, b| if rng.bool() { a.cmp(b) } else { b.cmp(a) }));
        }
        for e in expected {
            assert!(sorts.remove(e.to_vec().as_slice()));
        }
        assert!(sorts.is_empty());
    }

    #[test]
    fn test_fold_sorting() {
        let mut dag = Dag::new();

        dag.node("R", ());
        dag.node("A1", ());
        dag.node("A2", ());
        dag.node("A3", ());
        dag.node("B1", ());
        dag.node("B2", ());
        dag.node("B3", ());
        dag.node("C1", ());

        dag.dependency("A1", "R");
        dag.dependency("A2", "R");
        dag.dependency("A3", "R");

        dag.dependency("B1", "A1");
        dag.dependency("B2", "A1");
        dag.dependency("B3", "A2");
        dag.dependency("B3", "A3");

        dag.dependency("C1", "B1");
        dag.dependency("C1", "B2");
        dag.dependency("C1", "B3");

        let acc = dag.fold(&["R"], Vec::new(), |mut acc, key, _| {
            acc.push(*key);
            ControlFlow::Continue(acc)
        });
        assert_eq!(acc, vec!["R", "A1", "A2", "A3", "B1", "B2", "B3", "C1"]);
    }

    #[test]
    fn test_fold_reject() {
        let mut dag = Dag::new();

        dag.node("R", true);
        dag.node("A1", false); // Reject.
        dag.node("A2", true);
        dag.node("B1", true);
        dag.node("C1", true);
        dag.node("D1", true);

        dag.dependency("A1", "R");
        dag.dependency("A2", "R");
        dag.dependency("B1", "A1");
        dag.dependency("C1", "B1");
        dag.dependency("D1", "C1");
        dag.dependency("D1", "A2");

        let a1 = dag.get(&"A1").unwrap();
        assert_eq!(dag.descendants_of(a1), vec!["B1", "C1", "D1"]);

        let acc = dag.fold(&["R"], Vec::new(), |mut acc, key, accept| {
            if !accept.value {
                ControlFlow::Break(acc)
            } else {
                acc.push(*key);
                ControlFlow::Continue(acc)
            }
        });
        assert_eq!(acc, vec!["R", "A2"]);
    }
}