phylo 6.0.0

use itertools::Itertools;
use num::{Float, NumCast, One, Zero};
use std::fmt::Debug;
use vers_vecs::BitVec;

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

#[cfg(feature = "parallel")]
use rayon::prelude::*;

use crate::prelude::*;

/// A trait describing the path functions in a tree.
pub trait PathFunction: RootedTree
where
    <Self as RootedTree>::Node: RootedZetaNode,
{
    /// Sets zeta value of all nodes in a tree by function
    fn set_zeta(
        &mut self,
        zeta_func: fn(&Self, TreeNodeID<Self>) -> TreeNodeZeta<Self>,
    ) -> Option<()> {
        let node_ids = self.get_node_ids().collect_vec();
        for node_id in node_ids {
            let zeta = zeta_func(self, node_id);
            self.set_node_zeta(node_id, Some(zeta))?;
        }
        Some(())
    }

    /// Returns zeta value of a node in a tree. None is no zeta value is set
    fn get_zeta(&self, node_id: TreeNodeID<Self>) -> Option<TreeNodeZeta<Self>> {
        self.get_node(node_id)?.get_zeta()
    }

    /// Returns true if node zeta value is not None
    ///
    /// # Panics
    ///
    /// Panics if `node_id` is not a node of this tree.
    fn is_zeta_set(&self, node_id: TreeNodeID<Self>) -> bool {
        self.get_node(node_id)
            .expect("node_id is not a node of this tree")
            .is_zeta_set()
    }

    /// Returns true if all node zeta value is not None    
    fn is_all_zeta_set(&self) -> bool {
        // Iterate ids so the allocating `get_nodes` default isn't materialised.
        self.get_node_ids().all(|id| {
            self.get_node(id)
                .expect("invariant: id came from get_node_ids")
                .is_zeta_set()
        })
    }

    /// Sets zeta value of a node in a tree by value.
    fn set_node_zeta(
        &mut self,
        node_id: TreeNodeID<Self>,
        zeta: Option<TreeNodeZeta<Self>>,
    ) -> Option<()> {
        self.get_node_mut(node_id)?.set_zeta(zeta);
        Some(())
    }
}

/// A trait describing efficient computation of vertex to vertex distances
pub trait DistanceMatrix: RootedWeightedTree + EulerWalk
where
    <Self as RootedTree>::Node: RootedWeightedNode,
{
    /// Return the symmetrical pairwise distance matrix.
    fn matrix(&self) -> Vec<Vec<TreeNodeWeight<Self>>>;

    /// Distance between two nodes, using a prebuilt LCA oracle.
    ///
    /// The oracle is passed in rather than rebuilt per call so a full matrix
    /// shares one euler-tour index across all `O(n^2)` pairs. Build it once
    /// with [`EulerWalk::lca`] and borrow it here.
    /// # Errors
    ///
    /// [`TreeError::UnknownNode`] if either id is not a node of this tree.
    fn pairwise_distance(
        &self,
        oracle: &LcaOracle<'_, Self>,
        node_id_1: TreeNodeID<Self>,
        node_id_2: TreeNodeID<Self>,
    ) -> Result<TreeNodeWeight<Self>, TreeError>;
}

/// A trait describing naive computation of Robinson Foulds distance
pub trait RobinsonFoulds
where
    Self: RootedTree + RootedMetaTree + Clusters,
    <Self as RootedTree>::Node: RootedMetaNode,
{
    /// Returns Robinson Foulds distance between tree and self.
    fn rf(&self, tree: &Self) -> usize {
        let mut dist = 0;
        let mut all_taxa: HashSet<&TreeNodeMeta<Self>> = self.get_taxa_space().collect();
        all_taxa.extend(tree.get_taxa_space());
        let num_taxa = all_taxa.len();
        let all_taxa_map: HashMap<&TreeNodeMeta<Self>, usize> = all_taxa
            .into_iter()
            .enumerate()
            .map(|x| (x.1, x.0))
            .collect();

        let mut self_bps: HashMap<TreeNodeID<Self>, BitVec> = vec![].into_iter().collect();
        let mut self_out_bps: HashSet<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_taxa);
            let mut bp_rev = BitVec::from_ones(num_taxa);
            match self.is_leaf(n_id) {
                true => {
                    let leaf_meta = self
                        .get_node_taxa(n_id)
                        .expect("invariant: leaves carry a taxon");
                    bp.flip_bit(
                        *all_taxa_map
                            .get(leaf_meta)
                            .expect("invariant: all_taxa_map spans both trees' taxa"),
                    );

                    bp_rev
                        .apply_mask_xor(&bp)
                        .expect("invariant: every bitvector is sized to num_taxa");
                    self_bps.insert(n_id, bp.clone());
                    self_out_bps.insert(bp);
                    self_out_bps.insert(bp_rev);
                }
                false => {
                    if n_id == self.get_root_id() {
                        continue;
                    } else {
                        self.get_node_children_ids(n_id)
                            .map(|x| {
                                self_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_taxa");
                            });

                        bp_rev
                            .apply_mask_xor(&bp)
                            .expect("invariant: every bitvector is sized to num_taxa");
                        self_bps.insert(n_id, bp.clone());
                        self_out_bps.insert(bp);
                        self_out_bps.insert(bp_rev);
                    }
                }
            };
        }

        let mut tree_bps: HashMap<TreeNodeID<Self>, BitVec> = vec![].into_iter().collect();
        let mut tree_out_bps: HashSet<BitVec> = vec![].into_iter().collect();
        for n_id in tree
            .postord_ids(tree.get_root_id())
            .expect("invariant: the root id always names a node")
        {
            let mut bp = BitVec::from_zeros(num_taxa);
            let mut bp_rev = BitVec::from_ones(num_taxa);

            match tree.is_leaf(n_id) {
                true => {
                    let leaf_meta = tree
                        .get_node_taxa(n_id)
                        .expect("invariant: leaves carry a taxon");
                    bp.flip_bit(
                        *all_taxa_map
                            .get(leaf_meta)
                            .expect("invariant: all_taxa_map spans both trees' taxa"),
                    );

                    bp_rev
                        .apply_mask_xor(&bp)
                        .expect("invariant: every bitvector is sized to num_taxa");

                    tree_bps.insert(n_id, bp.clone());
                    tree_out_bps.insert(bp);
                    tree_out_bps.insert(bp_rev);
                }
                false => {
                    if n_id == tree.get_root_id() {
                        continue;
                    } else {
                        tree.get_node_children_ids(n_id)
                            .map(|x| {
                                tree_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_taxa");
                            });

                        bp_rev
                            .apply_mask_xor(&bp)
                            .expect("invariant: every bitvector is sized to num_taxa");

                        tree_bps.insert(n_id, bp.clone());
                        tree_out_bps.insert(bp);
                        tree_out_bps.insert(bp_rev);
                    }
                }
            };
        }
        for i in self_out_bps.iter() {
            if tree_out_bps.contains(i) {
                continue;
            } else {
                dist += 1;
            }
        }
        for i in tree_out_bps.iter() {
            if self_out_bps.contains(i) {
                continue;
            } else {
                dist += 1;
            }
        }

        dist / 2
    }
}

/// A trait describing naive computation of Cluster Matching distance
pub trait ClusterMatching
where
    Self: RootedTree + RootedMetaTree + Clusters,
    <Self as RootedTree>::Node: RootedMetaNode,
{
    /// Returns Cluster Matching distance between tree and self.
    fn cm(&self, tree: &Self) -> usize {
        let self_clusters = self
            .get_node_ids()
            .map(|node_id| {
                self.get_cluster_ids(node_id)
                    .expect("invariant: node_id came from get_node_ids")
                    .map(|id| {
                        self.get_node_taxa(id)
                            .expect("invariant: cluster members are labelled leaves")
                    })
                    .collect_vec()
            })
            .collect::<HashSet<_>>();
        let tree_clusters = tree
            .get_node_ids()
            .map(|node_id| {
                tree.get_cluster_ids(node_id)
                    .expect("invariant: node_id came from get_node_ids")
                    .map(|id| {
                        tree.get_node_taxa(id)
                            .expect("invariant: cluster members are labelled leaves")
                    })
                    .collect_vec()
            })
            .collect::<HashSet<_>>();

        self_clusters.difference(&tree_clusters).collect_vec().len()
            + tree_clusters.difference(&self_clusters).collect_vec().len()
    }
}

/// A trait describing computation of Cluster Affinity
pub trait ClusterAffinity
where
    Self: RootedTree + RootedMetaTree + Clusters,
    <Self as RootedTree>::Node: RootedMetaNode,
{
    /// Returns Cluster Affinity cost from self to tree..
    fn ca(&self, tree: &Self) -> usize {
        let mut dist = 0;
        let mut t1_size_map: HashMap<TreeNodeID<Self>, usize> =
            [].into_iter().collect::<HashMap<_, _>>();
        let mut t2_size_map: HashMap<TreeNodeID<Self>, usize> =
            [].into_iter().collect::<HashMap<_, _>>();
        let mut intersection_map: HashMap<(TreeNodeID<Self>, TreeNodeID<Self>), usize> =
            [].into_iter().collect::<HashMap<_, _>>();
        for v in self
            .postord_ids(self.get_root_id())
            .expect("invariant: the root id always names a node")
        {
            let mut mindist = usize::MAX;
            let vsize = if self.is_leaf(v) {
                1
            } else {
                self.get_node_children_ids(v)
                    .map(|x| {
                        t1_size_map
                            .get(&x)
                            .expect("invariant: children precede parents in post-order")
                    })
                    .sum()
            };
            t1_size_map.insert(v, vsize);
            for c in tree
                .postord_ids(tree.get_root_id())
                .expect("invariant: the root id always names a node")
            {
                let mut size = 0;
                let mut intersection = 0;
                if tree.is_leaf(c) {
                    size = 1;
                    if self.is_leaf(v) {
                        if self.get_node_taxa(v).expect(
                            "invariant: branch is guarded by is_leaf, and leaves carry a taxon",
                        ) == tree.get_node_taxa(c).expect(
                            "invariant: branch is guarded by is_leaf, and leaves carry a taxon",
                        ) {
                            intersection = 1
                        } else {
                            intersection = 0
                        }
                    } else {
                        for ch in self.get_node_children_ids(v) {
                            if *intersection_map.get(&(ch, c)).unwrap_or(&0) > 0 {
                                intersection = 1;
                                break;
                            }
                        }
                    }
                } else {
                    for cch in tree.get_node_children_ids(c) {
                        size += t2_size_map
                            .get(&cch)
                            .expect("invariant: children precede parents in post-order");
                        intersection += intersection_map
                            .get(&(v, cch))
                            .expect("invariant: (v, child) was inserted on an earlier pass");
                    }
                }
                t2_size_map.insert(c, size);
                intersection_map.insert((v, c), intersection);
                let cdist = size + vsize - (2 * intersection);
                if mindist > cdist {
                    mindist = cdist;
                }
            }
            dist += mindist
        }
        dist
    }
}

/// A trait describing naive computation of Weighted Robinson Foulds distance
pub trait WeightedRobinsonFoulds
where
    Self: RootedWeightedTree + RootedMetaTree + Clusters,
    <Self as RootedTree>::Node: RootedWeightedNode + RootedMetaNode,
{
    /// Returns weighted Robinson Foulds distance between tree and self.
    fn wrfs(&self, tree: &Self) -> TreeNodeWeight<Self>;
}

/// A trait describing naive computation of cophenetic distance
pub trait CopheneticDistance:
    PathFunction + RootedMetaTree + Clusters + Ancestors + ContractTree + Debug
where
    <Self as RootedTree>::Node: RootedMetaNode + RootedZetaNode,
    TreeNodeZeta<Self>: NodeWeight,
{
    /// Returns the first node whose zeta annotation is unset, if any.
    fn first_unset_zeta(&self) -> Option<TreeNodeID<Self>> {
        self.get_node_ids().find(|id| !self.is_zeta_set(*id))
    }

    /// Returns zeta of leaf by taxa
    ///
    /// # Errors
    ///
    /// [`TreeError::TaxaSetMismatch`] if `taxa` does not label a node of this
    /// tree; [`TreeError::MissingZeta`] if that node has no zeta set.
    fn get_zeta_taxa(&self, taxa: &TreeNodeMeta<Self>) -> Result<TreeNodeZeta<Self>, TreeError> {
        let id = self
            .get_taxa_node_id(taxa)
            .ok_or(TreeError::TaxaSetMismatch)?;
        self.get_zeta(id).ok_or(TreeError::MissingZeta(id.into()))
    }

    /// Reurns the nth norm of an iterator composed of floating point values
    fn compute_norm(
        vector: impl Iterator<Item = TreeNodeZeta<Self>>,
        norm: u32,
    ) -> TreeNodeZeta<Self> {
        if norm == 0 {
            return vector.fold(<TreeNodeZeta<Self>>::zero(), |acc, x| acc.max(x));
        }
        if norm == 1 {
            return vector.sum();
        }
        vector
            .map(|x| {
                let mut out = <TreeNodeZeta<Self>>::one();
                for _ in 0..norm {
                    out = out * x;
                }
                out
            })
            .sum::<TreeNodeZeta<Self>>()
            .powf(
                <TreeNodeZeta<Self> as NumCast>::from(norm)
                    .expect("invariant: a u32 norm is representable in the zeta type")
                    .powi(-1),
            )
    }

    #[cfg(feature = "parallel")]
    /// Returns the vector norm for an iterator
    fn compute_norm_par(vector: Vec<TreeNodeZeta<Self>>, norm: u32) -> TreeNodeZeta<Self> {
        if norm == 1 {
            return vector.into_iter().sum();
        }
        vector
            .par_iter()
            .map(|x| (*x).powi(norm as i32))
            .sum::<TreeNodeZeta<Self>>()
            .powf(
                <TreeNodeZeta<Self> as NumCast>::from(norm)
                    .expect("invariant: a u32 norm is representable in the zeta type")
                    .powi(-1),
            )
    }

    /// Reurns the cophenetic distance between two trees using the naive algorithm (\Theta(n^2))
    ///
    /// # Errors
    ///
    /// [`TreeError::MissingZeta`] if either tree has a node with no zeta set;
    /// call [`PathFunction::set_zeta`] first.
    fn cophen_dist<'a>(
        &'a self,
        tree: &'a Self,
        norm: u32,
    ) -> Result<TreeNodeZeta<Self>, TreeError> {
        if let Some(id) = self.first_unset_zeta().or_else(|| tree.first_unset_zeta()) {
            return Err(TreeError::MissingZeta(id.into()));
        }
        let binding1 = self
            .get_taxa_space()
            .collect::<HashSet<&TreeNodeMeta<Self>>>();
        let binding2 = tree
            .get_taxa_space()
            .collect::<HashSet<&TreeNodeMeta<Self>>>();
        let taxa_set = binding1.intersection(&binding2).cloned();

        self.cophen_dist_by_taxa(tree, norm, taxa_set)
    }

    #[cfg(feature = "parallel")]
    /// Returns the cophenetic distance between two trees using the naive algorithm (\Theta(n^2))
    ///
    /// # Errors
    ///
    /// [`TreeError::MissingZeta`] if either tree has a node with no zeta set;
    /// call [`PathFunction::set_zeta`] first.
    fn cophen_dist_par<'a>(
        &'a self,
        tree: &'a Self,
        norm: u32,
    ) -> Result<TreeNodeZeta<Self>, TreeError> {
        if let Some(id) = self.first_unset_zeta().or_else(|| tree.first_unset_zeta()) {
            return Err(TreeError::MissingZeta(id.into()));
        }
        let binding1 = self
            .get_taxa_space()
            .collect::<HashSet<&TreeNodeMeta<Self>>>();
        let binding2 = tree
            .get_taxa_space()
            .collect::<HashSet<&TreeNodeMeta<Self>>>();
        let taxa_set = binding1.intersection(&binding2).cloned().collect_vec();

        self.cophen_dist_by_taxa_par(tree, norm, taxa_set.into_iter())
    }

    #[cfg(feature = "parallel")]
    /// Returns the Cophenetic distance between two trees restricted to a taxa set using the \theta(n^2) naive algorithm.
    fn cophen_dist_by_taxa_par<'a>(
        &'a self,
        tree: &'a Self,
        norm: u32,
        taxa_set: impl Iterator<Item = &'a TreeNodeMeta<Self>> + Send,
    ) -> Result<TreeNodeZeta<Self>, TreeError> {
        // One euler-tour index per tree, shared across every pair, rather than
        // the naive per-query rebuild the trait fallback would do.
        let self_oracle = self.lca();
        let tree_oracle = tree.lca();
        let cophen_vec: Vec<TreeNodeZeta<Self>> = taxa_set
            .combinations_with_replacement(2)
            .par_bridge()
            .map(|x| match x[0] == x[1] {
                true => {
                    let zeta_1 = self.get_zeta_taxa(x[0])?;
                    let zeta_2 = tree.get_zeta_taxa(x[0])?;
                    Ok((zeta_1 - zeta_2).abs())
                }
                false => {
                    let self_ids = x
                        .iter()
                        .map(|a| self.get_taxa_node_id(a).ok_or(TreeError::TaxaSetMismatch))
                        .collect::<Result<Vec<_>, _>>()?;
                    let tree_ids = x
                        .iter()
                        .map(|a| tree.get_taxa_node_id(a).ok_or(TreeError::TaxaSetMismatch))
                        .collect::<Result<Vec<_>, _>>()?;
                    let t_lca_id = self_oracle.get_lca_id(self_ids.as_slice());
                    let t_hat_lca_id = tree_oracle.get_lca_id(tree_ids.as_slice());
                    let zeta_1 = self
                        .get_zeta(t_lca_id)
                        .ok_or(TreeError::MissingZeta(t_lca_id.into()))?;
                    let zeta_2 = tree
                        .get_zeta(t_hat_lca_id)
                        .ok_or(TreeError::MissingZeta(t_hat_lca_id.into()))?;
                    Ok((zeta_1 - zeta_2).abs())
                }
            })
            .collect::<Result<Vec<_>, TreeError>>()?;

        Ok(Self::compute_norm_par(cophen_vec, norm))
    }

    /// Returns the Cophenetic distance between two trees restricted to a taxa set using the \theta(n^2) naive algorithm.
    fn cophen_dist_by_taxa<'a>(
        &'a self,
        tree: &'a Self,
        norm: u32,
        taxa_set: impl Iterator<Item = &'a TreeNodeMeta<Self>>,
    ) -> Result<TreeNodeZeta<Self>, TreeError> {
        // One euler-tour index per tree, shared across every pair, rather than
        // the naive per-query rebuild the trait fallback would do.
        let self_oracle = self.lca();
        let tree_oracle = tree.lca();
        let cophen_vec = taxa_set
            .combinations_with_replacement(2)
            .map(|x| match x[0] == x[1] {
                true => {
                    let zeta_1 = self.get_zeta_taxa(x[0])?;
                    let zeta_2 = tree.get_zeta_taxa(x[0])?;
                    Ok((zeta_1 - zeta_2).abs())
                }
                false => {
                    let self_ids = x
                        .iter()
                        .map(|a| self.get_taxa_node_id(a).ok_or(TreeError::TaxaSetMismatch))
                        .collect::<Result<Vec<_>, _>>()?;
                    let tree_ids = x
                        .iter()
                        .map(|a| tree.get_taxa_node_id(a).ok_or(TreeError::TaxaSetMismatch))
                        .collect::<Result<Vec<_>, _>>()?;
                    let t_lca_id = self_oracle.get_lca_id(self_ids.as_slice());
                    let t_hat_lca_id = tree_oracle.get_lca_id(tree_ids.as_slice());
                    let zeta_1 = self
                        .get_zeta(t_lca_id)
                        .ok_or(TreeError::MissingZeta(t_lca_id.into()))?;
                    let zeta_2 = tree
                        .get_zeta(t_hat_lca_id)
                        .ok_or(TreeError::MissingZeta(t_hat_lca_id.into()))?;
                    Ok((zeta_1 - zeta_2).abs())
                }
            })
            .collect::<Result<Vec<_>, TreeError>>()?;

        Ok(Self::compute_norm(cophen_vec.into_iter(), norm))
    }
}