phylo 6.0.0

#[cfg(feature = "non_crypto_hash")]
use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
#[cfg(not(feature = "non_crypto_hash"))]
use std::collections::{HashMap, HashSet};

use itertools::Itertools;
use std::collections::VecDeque;
use vers_vecs::BitVec;

use crate::{
    error::TreeError,
    iter::lca::LcaOracle,
    node::simple_rnode::RootedTreeNode,
    tree::simple_rtree::{RootedTree, TreeNodeID},
};

/// Trait describing depth-first iteration of nodes in a tree
pub trait DFS
where
    Self: RootedTree + Sized,
{
    /// Returns an iterator of immutable reference of nodes in a tree in postfix order
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node` is not a node of this tree.
    fn postord_nodes(
        &self,
        start_node: TreeNodeID<Self>,
    ) -> Result<impl Iterator<Item = &Self::Node>, TreeError>;

    /// Returns an iterator of NodeID's in a tree in postfix order
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node` is not a node of this tree.
    fn postord_ids(
        &self,
        start_node: TreeNodeID<Self>,
    ) -> Result<impl Iterator<Item = TreeNodeID<Self>>, TreeError>;

    /// Returns a DFS iterator of immutable node references a tree
    ///
    /// The start node is checked once, here; the returned iterator cannot fail.
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn dfs(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = &Self::Node>, TreeError> {
        let start = self
            .get_node(start_node_id)
            .ok_or_else(|| TreeError::UnknownNode(start_node_id.into()))?;
        // A tree reaches every node exactly once through the parent/child
        // structure, so no visited-set is needed to guard against revisits.
        let mut stack = VecDeque::from([start]);
        let mut out_vec = vec![];
        while let Some(x) = stack.pop_front() {
            out_vec.push(x);
            for &child_id in x.get_children().iter().rev() {
                stack.push_front(
                    self.get_node(child_id)
                        .expect("invariant: child id came from the arena"),
                );
            }
        }
        Ok(out_vec.into_iter())
    }
}

/// Trait describing breadth-first iteration of nodes in a tree
pub trait BFS
where
    Self: RootedTree + Sized,
{
    /// Returns an iterator of immutable reference of nodes in a tree in level order
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn bfs_nodes(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl Iterator<Item = &Self::Node>, TreeError>;

    /// Returns an iterator of NodeID's in a tree in level order
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn bfs_ids(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl Iterator<Item = TreeNodeID<Self>>, TreeError>;
}

/// Trait describing breadth-first iteration of nodes in a tree
pub trait PreOrder
where
    Self: RootedTree + Sized,
{
    /// Returns an iterator of immutable reference of nodes in a tree in prefix order
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn preord_nodes(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = &Self::Node>, TreeError> {
        let start = self
            .get_node(start_node_id)
            .ok_or_else(|| TreeError::UnknownNode(start_node_id.into()))?;
        // A tree reaches every node exactly once, so the visited-set the
        // previous body carried was pure overhead.
        let mut stack = VecDeque::from([start]);
        let mut out_vec = vec![];
        while let Some(x) = stack.pop_front() {
            out_vec.push(x);
            for &child_id in x.get_children().iter().rev() {
                stack.push_front(
                    self.get_node(child_id)
                        .expect("invariant: child id came from the arena"),
                );
            }
        }
        Ok(out_vec.into_iter())
    }

    /// Returns an iterator of NodeID's in a tree in prefix order
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn preord_ids(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = TreeNodeID<Self>>, TreeError> {
        if self.get_node(start_node_id).is_none() {
            return Err(TreeError::UnknownNode(start_node_id.into()));
        }
        // Reversing the children slice in place avoids both the visited-set and
        // the per-node `collect_vec` the previous body allocated.
        let mut stack = VecDeque::from([start_node_id]);
        let mut out_vec = vec![];
        while let Some(x) = stack.pop_front() {
            out_vec.push(x);
            for &child_id in self
                .get_node(x)
                .expect("invariant: id came from the traversal stack")
                .get_children()
                .iter()
                .rev()
            {
                stack.push_front(child_id);
            }
        }
        Ok(out_vec.into_iter())
    }
}

/// Trait describing iteration of nodes along a path
pub trait Ancestors
where
    Self: RootedTree + Sized,
{
    /// Returns an iterator of immutable references to nodes in a tree from root to node
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn root_to_node(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = &Self::Node>, TreeError> {
        let start = self
            .get_node(start_node_id)
            .ok_or_else(|| TreeError::UnknownNode(start_node_id.into()))?;
        let mut stack = VecDeque::from([start]);
        while let Some(x) = stack.pop_front() {
            stack.push_front(x);
            match x.get_parent() {
                Some(pid) => {
                    stack.push_front(
                        self.get_node(pid)
                            .expect("invariant: parent id came from the arena"),
                    );
                }
                None => {
                    break;
                }
            }
        }
        Ok(stack.into_iter())
    }

    /// Returns an iterator of NodeID's in a tree from root to node
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn root_to_node_ids(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = TreeNodeID<Self>>, TreeError> {
        if self.get_node(start_node_id).is_none() {
            return Err(TreeError::UnknownNode(start_node_id.into()));
        }
        let mut stack = VecDeque::from([start_node_id]);
        while let Some(x) = stack.pop_front() {
            stack.push_front(x);
            match self.get_node_parent_id(x) {
                Some(pid) => {
                    stack.push_front(pid);
                }
                None => {
                    break;
                }
            }
        }
        Ok(stack.into_iter())
    }

    /// Returns an iterator of immutable references to nodes in a tree from node to root
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn node_to_root(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = &Self::Node>, TreeError> {
        let start = self
            .get_node(start_node_id)
            .ok_or_else(|| TreeError::UnknownNode(start_node_id.into()))?;
        let mut stack = VecDeque::from([start]);
        while let Some(x) = stack.pop_front() {
            stack.push_back(x);
            match x.get_parent() {
                Some(pid) => {
                    stack.push_front(
                        self.get_node(pid)
                            .expect("invariant: parent id came from the arena"),
                    );
                }
                None => {
                    break;
                }
            }
        }
        Ok(stack.into_iter())
    }

    /// Returns an iterator of NodeID's in a tree from node to root
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn node_to_root_ids(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = TreeNodeID<Self>>, TreeError> {
        if self.get_node(start_node_id).is_none() {
            return Err(TreeError::UnknownNode(start_node_id.into()));
        }
        let mut stack = VecDeque::from([start_node_id]);
        while let Some(x) = stack.pop_front() {
            stack.push_back(x);
            match self.get_node_parent_id(x) {
                Some(pid) => {
                    stack.push_front(pid);
                }
                None => {
                    break;
                }
            }
        }
        Ok(stack.into_iter())
    }

    /// Returns depth of a node as number of edges in the path from node to root
    fn depth(&self, node_id: TreeNodeID<Self>) -> usize {
        // Count edges by walking to the root; the previous body materialised
        // the whole path in a `VecDeque` just to read its length.
        RootedTree::get_node_depth(self, node_id)
    }
}

/// Trait describing an Euler Tour of a tree
pub trait EulerWalk
where
    Self: RootedTree + Sized,
{
    /// Returns euler tour of tree as iterator of immutable references to nodes
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn euler_walk_nodes(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = &Self::Node>, TreeError> {
        let start = self
            .get_node(start_node_id)
            .ok_or_else(|| TreeError::UnknownNode(start_node_id.into()))?;
        let mut stack = VecDeque::from([start]);
        let mut visited: HashSet<TreeNodeID<Self>> = HashSet::default();
        let mut out_vec = vec![];
        while let Some(node) = stack.pop_front() {
            let id = node.get_id();
            if !visited.insert(id) {
                if let Some(parent_id) = node.get_parent() {
                    out_vec.push(
                        self.get_node(parent_id)
                            .expect("invariant: parent id came from the arena"),
                    )
                }
            } else {
                out_vec.push(node);
                stack.push_front(node);
                for &child_id in node.get_children().iter().rev() {
                    stack.push_front(
                        self.get_node(child_id)
                            .expect("invariant: child id came from the arena"),
                    )
                }
            }
        }
        Ok(out_vec.into_iter())
    }

    /// Returns euler tour of tree as iterator of NodeID's
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `start_node_id` is not a node of this tree.
    fn euler_walk_ids(
        &self,
        start_node_id: TreeNodeID<Self>,
    ) -> Result<impl ExactSizeIterator<Item = TreeNodeID<Self>>, TreeError> {
        if self.get_node(start_node_id).is_none() {
            return Err(TreeError::UnknownNode(start_node_id.into()));
        }
        let mut stack = VecDeque::from([start_node_id]);
        let mut visited: HashSet<TreeNodeID<Self>> = HashSet::default();
        let mut out_vec = vec![];
        while let Some(node_id) = stack.pop_front() {
            if !visited.insert(node_id) {
                if let Some(parent_id) = self.get_node_parent_id(node_id) {
                    out_vec.push(parent_id)
                }
            } else {
                out_vec.push(node_id);
                stack.push_front(node_id);
                for child_id in self
                    .get_node_children_ids(node_id)
                    .collect_vec()
                    .iter()
                    .rev()
                {
                    stack.push_front(*child_id)
                }
            }
        }
        Ok(out_vec.into_iter())
    }

    /// Builds a constant-time LCA oracle borrowing this tree immutably.
    ///
    /// The returned [`LcaOracle`] holds a shared borrow of `self`, so the tree
    /// cannot be mutated while it is alive. Build one, run every query against
    /// it, then drop it before mutating the tree again.
    fn lca(&self) -> LcaOracle<'_, Self> {
        LcaOracle::build(self)
    }

    /// Lowest common ancestor of a slice of nodes, by NodeID.
    ///
    /// Fallback for one-off queries: it builds a throwaway [`LcaOracle`] for a
    /// single lookup. Callers that query repeatedly should build one oracle
    /// with [`Self::lca`] and reuse it.
    ///
    /// # Errors
    ///
    /// [`TreeError::EmptyNodeSet`] if `node_id_vec` is empty;
    /// [`TreeError::UnknownNode`] if any id is not a node of this tree.
    fn get_lca_id(&self, node_id_vec: &[TreeNodeID<Self>]) -> Result<TreeNodeID<Self>, TreeError> {
        if node_id_vec.is_empty() {
            return Err(TreeError::EmptyNodeSet);
        }
        for id in node_id_vec {
            if self.get_node(*id).is_none() {
                return Err(TreeError::UnknownNode((*id).into()));
            }
        }
        Ok(self.lca().get_lca_id(node_id_vec))
    }

    /// Lowest common ancestor of a slice of nodes, by immutable reference.
    ///
    /// # Errors
    ///
    /// As [`Self::get_lca_id`].
    fn get_lca<'a>(
        &'a self,
        node_id_vec: &[TreeNodeID<Self>],
    ) -> Result<&'a Self::Node, TreeError> {
        let lca_id = self.get_lca_id(node_id_vec)?;
        Ok(self
            .get_node(lca_id)
            .expect("invariant: the LCA is a node of this tree"))
    }
}

/// Trait describing iteration of clusters and bipartitions in a tree.
pub trait Clusters: DFS + BFS + Sized {
    /// Returns cluster of a node in a rooted tree (smallest cluster in an unrooted tree) as iterator of immutable reference to a node
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `node_id` is not a node of this tree.
    fn get_cluster(
        &self,
        node_id: TreeNodeID<Self>,
    ) -> Result<impl Iterator<Item = &Self::Node>, TreeError> {
        // `dfs` already materialises its walk, so filtering it lazily avoids a
        // second `Vec` that only existed to promise `ExactSizeIterator`.
        Ok(self.dfs(node_id)?.filter(|x| x.is_leaf()))
    }

    /// Returns cluster of a node in a rooted tree (smallest cluster in an unrooted tree) as iterator of NodeID's
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `node_id` is not a node of this tree.
    fn get_cluster_ids(
        &self,
        node_id: TreeNodeID<Self>,
    ) -> Result<impl Iterator<Item = TreeNodeID<Self>>, TreeError> {
        Ok(self.get_cluster(node_id)?.map(move |x| x.get_id()))
    }

    /// Returns all clusters of a tree as iterator of NodeID's
    fn get_clusters_ids(
        &self,
    ) -> impl ExactSizeIterator<
        Item = (
            TreeNodeID<Self>,
            impl ExactSizeIterator<Item = TreeNodeID<Self>>,
        ),
    > {
        let mut clusters: HashMap<TreeNodeID<Self>, Vec<TreeNodeID<Self>>> =
            vec![].into_iter().collect();
        for n_id in self
            .postord_ids(self.get_root_id())
            .expect("invariant: the root id always names a node")
        {
            match self.is_leaf(n_id) {
                true => {
                    clusters.insert(n_id, vec![n_id]);
                }
                false => {
                    let node_cluster = self
                        .get_node_children_ids(n_id)
                        .flat_map(|x| {
                            clusters
                                .get(&x)
                                .cloned()
                                .expect("invariant: children precede parents in post-order")
                        })
                        .collect_vec();
                    clusters.insert(n_id, node_cluster);
                }
            };
        }

        clusters
            .into_iter()
            .map(|(n_id, cluster)| (n_id, cluster.into_iter()))
    }

    /// Returns size of a cluster of nodes
    ///
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if `node_id` is not a node of this tree.
    fn get_cluster_size(&self, node_id: TreeNodeID<Self>) -> Result<usize, TreeError> {
        Ok(self.get_cluster_ids(node_id)?.count())
    }

    /// Returns bipartition of an edge in a tree as iterator of immutable reference to a node
    fn get_bipartition(
        &self,
        edge: (TreeNodeID<Self>, TreeNodeID<Self>),
    ) -> Result<
        (
            impl Iterator<Item = &Self::Node>,
            impl Iterator<Item = &Self::Node>,
        ),
        TreeError,
    > {
        let c2 = self.get_cluster(edge.1)?;
        // Hash the opposite cluster so membership is O(1); the previous body
        // scanned `c2_ids` linearly for every element of `c1`. Both sides are
        // now lazy -- no `Vec` just to hand back an `ExactSizeIterator`.
        let c2_ids: HashSet<TreeNodeID<Self>> = self.get_cluster_ids(edge.1)?.collect();
        let c1 = self
            .get_cluster(edge.0)?
            .filter(move |x| !c2_ids.contains(&x.get_id()));
        Ok((c1, c2))
    }

    /// Returns bipartition of an edge in a tree as iterator of NodeID's
    fn get_bipartition_ids(
        &self,
        edge: (TreeNodeID<Self>, TreeNodeID<Self>),
    ) -> Result<
        (
            impl Iterator<Item = TreeNodeID<Self>>,
            impl Iterator<Item = TreeNodeID<Self>>,
        ),
        TreeError,
    > {
        let c2 = self.get_cluster_ids(edge.1)?;
        // O(1) membership instead of a linear scan per element of `c1`.
        let c2_ids: HashSet<TreeNodeID<Self>> = self.get_cluster_ids(edge.1)?.collect();
        let c1 = self
            .get_cluster_ids(edge.0)?
            .filter(move |x| !c2_ids.contains(x));
        Ok((c1, c2))
    }

    /// Returns all bipartitions of a tree as iterator of NodeID's
    fn get_bipartitions_ids(
        &self,
    ) -> impl ExactSizeIterator<
        Item = (
            impl ExactSizeIterator<Item = TreeNodeID<Self>>,
            impl ExactSizeIterator<Item = TreeNodeID<Self>>,
        ),
    > {
        let leaf_ids: HashMap<TreeNodeID<Self>, usize> = self
            .get_leaf_ids()
            .enumerate()
            .map(|(idx, id)| (id, idx))
            .collect();
        let leaf_ids_rev: Vec<TreeNodeID<Self>> = leaf_ids.keys().copied().collect();
        let num_leaves = leaf_ids.len();
        let mut bps: HashMap<TreeNodeID<Self>, BitVec> = vec![].into_iter().collect();
        for n_id in self
            .postord_ids(self.get_root_id())
            .expect("invariant: the root id always names a node")
        {
            let mut bp = BitVec::from_zeros(num_leaves);
            match self.is_leaf(n_id) {
                true => {
                    bp.flip_bit(
                        *leaf_ids
                            .get(&n_id)
                            .expect("invariant: n_id is a leaf, so it is in leaf_ids"),
                    );
                    bps.insert(n_id, bp.clone());
                }
                false => {
                    if n_id == self.get_root_id() {
                        continue;
                    }
                    self.get_node_children_ids(n_id)
                        .map(|x| {
                            bps.get(&x)
                                .expect("invariant: children precede parents in post-order")
                        })
                        .for_each(|x| {
                            bp.apply_mask_or(x)
                                .expect("invariant: every bitvector is sized to num_leaves");
                        });
                    if self.get_node_parent_id(n_id) != Some(self.get_root_id()) {
                        bps.insert(n_id, bp);
                    }
                }
            };
        }

        bps.into_values().map(move |bit_bp| {
            let mut bp1 = Vec::with_capacity(leaf_ids.len());
            let mut bp2 = Vec::with_capacity(leaf_ids.len());
            for (idx, bit) in leaf_ids_rev.iter().enumerate().take(bit_bp.len()) {
                match bit_bp
                    .is_bit_set(idx)
                    .expect("invariant: idx is bounded by bit_bp.len()")
                {
                    true => {
                        bp1.push(bit.to_owned());
                    }
                    false => {
                        bp2.push(bit.to_owned());
                    }
                }
            }
            (bp1.into_iter(), bp2.into_iter())
        })
    }

    /// Returns median NodeID of a set of leaves in a tree.
    ///
    /// # Errors
    ///
    /// [`TreeError::EmptyNodeSet`] if `taxa_set` yields nothing.
    fn get_median_node_id_for_leaves(
        &self,
        taxa_set: impl Iterator<Item = TreeNodeID<Self>>,
    ) -> Result<TreeNodeID<Self>, TreeError> {
        let mut cluster_sizes: HashMap<TreeNodeID<Self>, usize> = vec![].into_iter().collect();
        let mut median_node_id: TreeNodeID<Self> = self.get_root_id();
        let leaf_ids: HashSet<TreeNodeID<Self>> = taxa_set.collect();
        if leaf_ids.is_empty() {
            return Err(TreeError::EmptyNodeSet);
        }
        for n_id in self
            .postord_ids(self.get_root_id())
            .expect("invariant: the root id always names a node")
        {
            if self.is_leaf(n_id) && leaf_ids.contains(&n_id) {
                cluster_sizes.insert(n_id, 1);
            } else {
                let mut cluster_size = 0;
                for c_id in self.get_node_children_ids(n_id) {
                    cluster_size += cluster_sizes
                        .get(&c_id)
                        .expect("invariant: children precede parents in post-order");
                }
                cluster_sizes.insert(n_id, cluster_size);
            }
        }
        loop {
            median_node_id = self
                .get_node_children_ids(median_node_id)
                .max_by(|x, y| {
                    let x_cluster_size = cluster_sizes
                        .get(x)
                        .expect("invariant: every node was sized in the post-order pass");
                    let y_cluster_size = cluster_sizes
                        .get(y)
                        .expect("invariant: every node was sized in the post-order pass");
                    x_cluster_size.cmp(y_cluster_size)
                })
                .expect("invariant: the loop only descends into internal nodes");
            if cluster_sizes
                .get(&median_node_id)
                .expect("invariant: every node was sized in the post-order pass")
                <= &(leaf_ids.len() / 2)
            {
                break;
            }
        }
        Ok(median_node_id)
    }

    /// Returns immutable reference to median node of a set of leaves in a tree.
    ///
    /// # Errors
    ///
    /// As [`Self::get_median_node_id_for_leaves`].
    fn get_median_node_for_leaves(
        &self,
        taxa_set: impl Iterator<Item = TreeNodeID<Self>>,
    ) -> Result<&Self::Node, TreeError> {
        let id = self.get_median_node_id_for_leaves(taxa_set)?;
        Ok(self
            .get_node(id)
            .expect("invariant: the median node is a node of this tree"))
    }

    /// Returns an immutable reference to median node of all leaves in a tree.
    ///
    /// # Errors
    ///
    /// [`TreeError::EmptyNodeSet`] if the tree has no leaves.
    fn get_median_node(&self) -> Result<&Self::Node, TreeError> {
        let leaves = self.get_leaves().map(|x| x.get_id()).collect_vec();
        self.get_median_node_for_leaves(leaves.into_iter())
    }

    /// Returns median NodeID of all leaves in a tree.
    ///
    /// # Errors
    ///
    /// [`TreeError::EmptyNodeSet`] if the tree has no leaves.
    fn get_median_node_id(&self) -> Result<TreeNodeID<Self>, TreeError> {
        let leaves = self.get_leaf_ids().collect_vec();
        self.get_median_node_id_for_leaves(leaves.into_iter())
    }
}