phylo 6.0.0

use std::{
    fmt::{Debug, Display},
    hash::Hash,
};

use crate::node::simple_rnode::{NodeTaxa, RootedMetaNode};
use crate::prelude::{Clusters, EulerWalk, PreOrder, RootedMetaTree, RootedTree, DFS};
use crate::{
    error::TreeError,
    iter::lca::LcaOracle,
    iter::node_iter::Ancestors,
    node::simple_rnode::RootedTreeNode,
    tree::simple_rtree::{TreeNodeID, TreeNodeMeta},
};

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

/// A trait describing subtree-prune-regraft operations
pub trait SPR: RootedTree + DFS + Sized {
    /// Attaches input tree to self by spliting an edge
    fn graft(
        &mut self,
        tree: Self,
        edge: (TreeNodeID<Self>, TreeNodeID<Self>),
    ) -> Result<(), TreeError>;

    /// Returns subtree starting at given node, while corresponding nodes from self.
    fn prune(&mut self, node_id: TreeNodeID<Self>) -> Result<Self, TreeError>;

    /// SPR function
    fn spr(
        &mut self,
        edge1: (TreeNodeID<Self>, TreeNodeID<Self>),
        edge2: (TreeNodeID<Self>, TreeNodeID<Self>),
    ) -> Result<(), TreeError> {
        let pruned_tree = SPR::prune(self, edge1.1)?;
        SPR::graft(self, pruned_tree, edge2)
    }
}

/// A trait describing Nearest Neighbour interchange operations
pub trait NNI
where
    Self: RootedTree + Sized,
{
    /// Performs an NNI operation
    fn nni(&mut self, node_id: TreeNodeID<Self>, left_ch: bool) -> Result<(), TreeError>;
}

/// A trait describing rerooting a tree
pub trait Reroot<'a>
where
    Self: RootedTree + Sized,
{
    /// Reroots tree at node. **Note: this changes the degree of a node**
    fn reroot_at_node(&mut self, node_id: TreeNodeID<Self>) -> Result<(), TreeError>;
    /// Reroots tree at a split node.
    fn reroot_at_edge(
        &mut self,
        edge: (TreeNodeID<Self>, TreeNodeID<Self>),
    ) -> Result<(), TreeError>;
}

/// A trait describing balancing a binary tree
pub trait Balance: Clusters + SPR + Sized
where
    TreeNodeID<Self>: Display + Debug + Hash + Clone + Ord,
{
    /// Balances a binary tree
    fn balance_subtree(&mut self) -> Result<(), TreeError>;
}

/// A trait describing subtree queries of a tree
pub trait Subtree: Ancestors + DFS + Sized
where
    TreeNodeID<Self>: Display + Debug + Hash + Clone + Ord,
{
    /// Returns a subtree consisting of only provided nodes
    fn induce_tree(
        &self,
        node_id_list: impl IntoIterator<
            Item = TreeNodeID<Self>,
            IntoIter = impl ExactSizeIterator<Item = TreeNodeID<Self>>,
        >,
    ) -> Result<Self, TreeError> {
        let mut subtree = Self::new();
        // As in `subtree`: drop the `Self::new()` placeholder if the induced
        // node set never overwrites it (happens when `self`'s root id is not
        // the placeholder's, e.g. inducing on an already-extracted subtree).
        let placeholder = subtree.get_root_id();
        let mut placeholder_kept = self.get_root_id() == placeholder;
        subtree.set_root(self.get_root_id());
        subtree.set_node(self.get_root().clone());
        for node_id in node_id_list.into_iter() {
            let ancestors: Vec<Self::Node> = self.root_to_node(node_id)?.cloned().collect();
            placeholder_kept |= ancestors.iter().any(|node| node.get_id() == placeholder);
            subtree.set_nodes(ancestors.into_iter());
        }
        subtree.clean();
        if !placeholder_kept {
            subtree.delete_node(placeholder);
        }
        Ok(subtree)
    }

    /// Returns subtree starting at provided node.
    fn subtree(&self, node_id: TreeNodeID<Self>) -> Result<Self, TreeError> {
        let mut subtree = Self::new();
        // `Self::new()` seeds a placeholder root node. Unless the extracted
        // subtree overwrites that slot, it would linger as an unreachable node
        // in the new arena, so track whether the real nodes cover it.
        let placeholder = subtree.get_root_id();
        subtree.set_root(node_id);
        let nodes: Vec<Self::Node> = self.dfs(node_id)?.cloned().collect();
        let placeholder_kept = nodes.iter().any(|node| node.get_id() == placeholder);
        subtree.set_nodes(nodes.into_iter());
        subtree
            .get_node_mut(node_id)
            .ok_or_else(|| TreeError::UnknownNode(node_id.into()))?
            .set_parent(None);
        if !placeholder_kept {
            subtree.delete_node(placeholder);
        }
        Ok(subtree)
    }
}

/// A trait describing tree contraction operations
pub trait ContractTree: EulerWalk + DFS {
    /// Contracts tree that from post_ord node_id iterator.
    fn contracted_tree_nodes_from_iter(
        &self,
        new_tree_root_id: TreeNodeID<Self>,
        leaf_ids: &[TreeNodeID<Self>],
        node_iter: impl Iterator<Item = TreeNodeID<Self>>,
    ) -> impl Iterator<Item = Self::Node> {
        let mut node_map: HashMap<TreeNodeID<Self>, Self::Node> = HashMap::from_iter(vec![(
            new_tree_root_id,
            self.get_node(new_tree_root_id)
                .expect("new_tree_root_id is not a node of this tree")
                .clone(),
        )]);
        let mut remove_list: HashSet<TreeNodeID<Self>> = HashSet::from_iter(vec![]);

        let leaf_ids: HashSet<&TreeNodeID<Self>> = leaf_ids.iter().collect();
        node_iter
            .map(|x| {
                self.get_node(x)
                    .cloned()
                    .expect("node_iter yielded an id that is not in this tree")
            })
            .for_each(|mut node| {
                match node.is_leaf() {
                    true => {
                        if leaf_ids.contains(&node.get_id()) {
                            node_map.insert(node.get_id(), node);
                        }
                    }
                    false => {
                        let node_children_ids = node.get_children().to_vec();
                        for child_id in &node_children_ids {
                            match node_map.contains_key(child_id) {
                                true => {}
                                false => node.remove_child(child_id),
                            }
                        }
                        let node_children_ids = node.get_children().to_vec();
                        match node_children_ids.len() {
                            0 => {}
                            1 => {
                                // the node is a unifurcation
                                // node should be added to both node_map and remove_list
                                // if child of node is already in remove list, attach node children to node first
                                let child_node_id = node_children_ids[0];

                                if remove_list.contains(&child_node_id) {
                                    node.remove_child(&child_node_id);
                                    let grandchildren_ids = node_map
                                        .get(&child_node_id)
                                        .expect("invariant: node inserted during the post-order pass")
                                        .get_children()
                                        .to_vec();
                                    for grandchild_id in grandchildren_ids {
                                        node_map
                                            .get_mut(&grandchild_id)
                                            .expect("invariant: node inserted during the post-order pass")
                                            .set_parent(Some(node.get_id()));
                                        node.add_child(grandchild_id);
                                    }
                                }
                                remove_list.insert(node.get_id());
                                node_map.insert(node.get_id(), node);
                            }
                            _ => {
                                // node has multiple children
                                // for each child, suppress child if child is in remove list
                                node_children_ids.into_iter().for_each(|chid| {
                                    if remove_list.contains(&chid) {
                                        // suppress chid
                                        // remove chid from node children
                                        // children of chid are node grandchildren
                                        // add grandchildren to node children
                                        // set grandchildren parent to node
                                        node.remove_child(&chid);
                                        let node_grandchildren =
                                            node_map.get(&chid).expect("invariant: node inserted during the post-order pass").get_children().to_vec();
                                        for grandchild in node_grandchildren {
                                            node.add_child(grandchild);
                                            node_map
                                                .get_mut(&grandchild)
                                                .expect("invariant: node inserted during the post-order pass")
                                                .set_parent(Some(node.get_id()))
                                        }
                                    }
                                });
                                if node.get_id() == new_tree_root_id {
                                    node.set_parent(None);
                                }
                                node_map.insert(node.get_id(), node);
                            }
                        };
                    }
                }
            });
        remove_list.into_iter().for_each(|x| {
            node_map.remove(&x);
        });
        node_map.into_values()
    }

    /// Returns a deep copy of the nodes in the contracted tree.
    ///
    /// Resolves the contracted tree's root by [`EulerWalk::get_lca_id`], which
    /// builds a throwaway [`LcaOracle`]. Callers
    /// that already know the root — or that contract repeatedly against one
    /// tree — should use [`Self::contracted_tree_nodes_from_root`] instead.
    fn contracted_tree_nodes(
        &self,
        leaf_ids: &[TreeNodeID<Self>],
    ) -> Result<impl Iterator<Item = Self::Node>, TreeError> {
        let root = self.get_lca_id(leaf_ids)?;
        Ok(self.contracted_tree_nodes_from_root(root, leaf_ids))
    }

    /// Returns a deep copy of the nodes in the contracted tree, given its root.
    ///
    /// `new_tree_root_id` must be the LCA of `leaf_ids`. Taking it as a
    /// parameter — as [`Self::contracted_tree_nodes_from_iter`] already does —
    /// is what lets a caller resolve it once and reuse it, rather than paying
    /// for an Euler tour and RMQ build per contraction.
    fn contracted_tree_nodes_from_root(
        &self,
        new_tree_root_id: TreeNodeID<Self>,
        leaf_ids: &[TreeNodeID<Self>],
    ) -> impl Iterator<Item = Self::Node> {
        let node_postord_iter = self
            .postord_nodes(new_tree_root_id)
            .expect("new_tree_root_id is not a node of this tree");
        let mut node_map: HashMap<TreeNodeID<Self>, Self::Node> = HashMap::from_iter(vec![(
            new_tree_root_id,
            self.get_node(new_tree_root_id)
                .expect("new_tree_root_id is not a node of this tree")
                .clone(),
        )]);

        let leaf_ids: HashSet<&TreeNodeID<Self>> = leaf_ids.iter().collect();
        let mut remove_list: HashSet<TreeNodeID<Self>> = HashSet::default();
        node_postord_iter.for_each(|orig_node| {
            let mut node = orig_node.clone();
            match node.is_leaf() {
                true => {
                    if leaf_ids.contains(&node.get_id()) {
                        node_map.insert(node.get_id(), node);
                    }
                }
                false => {
                    let node_children_ids = node.get_children().to_vec();
                    for child_id in &node_children_ids {
                        match node_map.contains_key(child_id) {
                            true => {}
                            false => node.remove_child(child_id),
                        }
                    }
                    let node_children_ids = node.get_children().to_vec();
                    match node_children_ids.len() {
                        0 => {}
                        1 => {
                            // the node is a unifurcation
                            // node should be added to both node_map and remove_list
                            // if child of node is already in remove list, attach node children to node first
                            let child_node_id = node_children_ids[0];

                            if remove_list.contains(&child_node_id) {
                                node.remove_child(&child_node_id);
                                let grandchildren_ids = node_map
                                    .get(&child_node_id)
                                    .expect("invariant: node inserted during the post-order pass")
                                    .get_children()
                                    .to_vec();
                                for grandchild_id in grandchildren_ids {
                                    node_map
                                        .get_mut(&grandchild_id)
                                        .expect("invariant: node inserted during the post-order pass")
                                        .set_parent(Some(node.get_id()));
                                    node.add_child(grandchild_id);
                                }
                            }
                            remove_list.insert(node.get_id());
                            node_map.insert(node.get_id(), node);
                        }
                        _ => {
                            // node has multiple children
                            // for each child, suppress child if child is in remove list
                            node_children_ids.into_iter().for_each(|chid| {
                                if remove_list.contains(&chid) {
                                    // suppress chid
                                    // remove chid from node children
                                    // children of chid are node grandchildren
                                    // add grandchildren to node children
                                    // set grandchildren parent to node
                                    node.remove_child(&chid);
                                    let node_grandchildren =
                                        node_map.get(&chid).expect("invariant: node inserted during the post-order pass").get_children().to_vec();
                                    for grandchild in node_grandchildren {
                                        node.add_child(grandchild);
                                        node_map
                                            .get_mut(&grandchild)
                                            .expect("invariant: node inserted during the post-order pass")
                                            .set_parent(Some(node.get_id()))
                                    }
                                }
                            });
                            if node.get_id() == new_tree_root_id {
                                node.set_parent(None);
                            }
                            node_map.insert(node.get_id(), node.clone());
                        }
                    };
                }
            }
        });
        remove_list.into_iter().for_each(|x| {
            node_map.remove(&x);
        });
        node_map.into_values()
    }

    /// Returns a contracted tree from slice containing NodeID's
    ///
    /// Builds a throwaway [`LcaOracle`] to find the
    /// contracted tree's root — an Euler tour plus RMQ, linear in the size of
    /// the original tree. Callers contracting the same tree more than once
    /// should build one oracle with [`EulerWalk::lca`] and hand it to
    /// [`Self::contract_tree_with_oracle`], which is the same work amortised.
    fn contract_tree(&self, leaf_ids: &[TreeNodeID<Self>]) -> Result<Self, TreeError> {
        self.contract_tree_with_oracle(leaf_ids, &self.lca())
    }

    /// Returns a contracted tree, reusing a prebuilt LCA oracle.
    ///
    /// `oracle` must have been built from this tree.
    fn contract_tree_with_oracle(
        &self,
        leaf_ids: &[TreeNodeID<Self>],
        oracle: &LcaOracle<'_, Self>,
    ) -> Result<Self, TreeError>;

    /// Returns a contracted tree from an iterator containing NodeID's
    fn contract_tree_from_iter(
        &self,
        leaf_ids: &[TreeNodeID<Self>],
        node_iter: impl Iterator<Item = TreeNodeID<Self>>,
    ) -> Result<Self, TreeError>;
}

/// A struct representing an Ordered Leaf Array tree
#[derive(Clone)]
pub struct OLATree<T: NodeTaxa> {
    /// Taxa labels in leaf ordering σ
    pub taxa: Vec<T>,
    /// OLA indices: non-negative values are leaf indices, negative values are internal node indices
    pub indices: Vec<i64>,
}

impl<T: NodeTaxa> Default for OLATree<T> {
    fn default() -> Self {
        OLATree {
            taxa: Vec::new(),
            indices: Vec::new(),
        }
    }
}

/// Returns the child of `ancestor` on the path toward `descendant`.
fn child_of<Tr>(tree: &Tr, ancestor: TreeNodeID<Tr>, descendant: TreeNodeID<Tr>) -> TreeNodeID<Tr>
where
    Tr: RootedTree,
    Tr::Node: RootedTreeNode,
{
    let mut current = descendant;
    loop {
        let parent = tree
            .get_node_parent_id(current)
            .expect("invariant: walk stops at `ancestor`, which is above the root path");
        if parent == ancestor {
            return current;
        }
        current = parent;
    }
}

/// A trait for converting trees to and from an Ordered Leaf Array representation
pub trait OLA: RootedMetaTree + EulerWalk + PreOrder
where
    Self::Node: RootedMetaNode,
{
    /// Constructs a tree from an OLATree representation
    fn from_vec(ola: OLATree<TreeNodeMeta<Self>>) -> Self;

    /// Converts the tree into an OLATree representation.
    ///
    /// Leaves are ordered by pre-order DFS traversal to establish the leaf ordering σ.
    /// Each index entry is either a leaf index (≥ 0) or an internal node index (< 0).
    fn to_vec(&self) -> OLATree<TreeNodeMeta<Self>> {
        // Step 1: collect leaves in pre-order to fix leaf ordering σ
        let leaf_ids: Vec<TreeNodeID<Self>> = self
            .preord_ids(self.get_root_id())
            .expect("invariant: the root id always names a node")
            .filter(|id| self.is_leaf(*id))
            .collect();

        let n = leaf_ids.len();
        if n <= 1 {
            return OLATree::default();
        }

        let mut ola_indices: Vec<i64> = Vec::with_capacity(n - 1);

        // One euler-tour index shared across the O(n^2) LCA queries below.
        // Without it every `get_lca_id` call would rebuild the whole index.
        let oracle = self.lca();

        // LCA(l_i, l_j) for every j < i, held across both passes below. Each is
        // an O(1) oracle query, but there are O(n^2) of them, and the deepest-LCA
        // pass and the sibling filter want the same set -- recomputing it for the
        // second pass doubled the query count for no gain. Reused across
        // iterations of `i` so the growing buffer is allocated once.
        let mut lcas: Vec<TreeNodeID<Self>> = Vec::with_capacity(n);

        // Step 2: for each leaf l_i (i >= 1), find its sibling in the restricted tree T^i
        for i in 1..n {
            let li = leaf_ids[i];

            lcas.clear();
            lcas.extend((0..i).map(|j| oracle.get_lca_id(&[li, leaf_ids[j]])));

            // The parent of l_i in T^i is the LCA(l_i, l_j) with the greatest depth over all j < i
            let p_id = *lcas
                .iter()
                .max_by_key(|&&lca| oracle.get_node_depth(lca))
                .expect("invariant: i >= 1, so lcas holds at least one entry");

            // Sibling's leaves in T^{i-1}: those l_j (j < i) whose LCA with l_i is exactly p_id,
            // meaning they live on the opposite side of p_id from l_i
            let sibling_indices: Vec<usize> = (0..i).filter(|&j| lcas[j] == p_id).collect();

            let entry = if sibling_indices.len() == 1 {
                // Sibling is a single leaf: entry = its index in σ (non-negative)
                sibling_indices[0] as i64
            } else {
                // Sibling is an internal node.
                // index(v) = -max(μ(c1), μ(c2)), where μ(c) = min leaf index in child c's subtree.
                // The sibling node in T^i is the LCA of all sibling leaves in the original tree.
                // Split sibling_indices by which child of sib_id each leaf descends through.
                let sib_leaf_ids: Vec<TreeNodeID<Self>> =
                    sibling_indices.iter().map(|&j| leaf_ids[j]).collect();
                let sib_id = oracle.get_lca_id(&sib_leaf_ids);
                let mut child_min: HashMap<TreeNodeID<Self>, usize> = HashMap::default();
                for &j in &sibling_indices {
                    let child = child_of(self, sib_id, leaf_ids[j]);
                    let e = child_min.entry(child).or_insert(j);
                    if j < *e {
                        *e = j;
                    }
                }
                let mu_max = *child_min
                    .values()
                    .max()
                    .expect("invariant: sibling_indices is non-empty in this branch");
                -(mu_max as i64)
            };

            ola_indices.push(entry);
        }

        // Step 3: collect taxa labels in leaf ordering σ
        let taxa: Vec<TreeNodeMeta<Self>> = leaf_ids
            .iter()
            .map(|&id| {
                self.get_node_taxa_cloned(id)
                    .expect("invariant: id came from the leaf set, which is labelled")
            })
            .collect();

        OLATree {
            taxa,
            indices: ola_indices,
        }
    }
}