phylo 5.1.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
    fn is_zeta_set(&self, node_id: TreeNodeID<Self>) -> bool {
        self.get_node(node_id).unwrap().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).unwrap().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.
    fn pairwise_distance(
        &self,
        oracle: &LcaOracle<'_, Self>,
        node_id_1: TreeNodeID<Self>,
        node_id_2: TreeNodeID<Self>,
    ) -> TreeNodeWeight<Self>;
}

/// 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()) {
            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).unwrap();
                    bp.flip_bit(*all_taxa_map.get(leaf_meta).unwrap());

                    let _ = bp_rev.apply_mask_xor(&bp);
                    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).unwrap())
                            .for_each(|x| {
                                let _ = bp.apply_mask_or(x);
                            });

                        let _ = bp_rev.apply_mask_xor(&bp);
                        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()) {
            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).unwrap();
                    bp.flip_bit(*all_taxa_map.get(leaf_meta).unwrap());

                    let _ = bp_rev.apply_mask_xor(&bp);

                    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).unwrap())
                            .for_each(|x| {
                                let _ = bp.apply_mask_or(x);
                            });

                        let _ = bp_rev.apply_mask_xor(&bp);

                        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)
                    .map(|id| self.get_node_taxa(id).unwrap())
                    .collect_vec()
            })
            .collect::<HashSet<_>>();
        let tree_clusters = tree
            .get_node_ids()
            .map(|node_id| {
                tree.get_cluster_ids(node_id)
                    .map(|id| tree.get_node_taxa(id).unwrap())
                    .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()) {
            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).unwrap())
                    .sum()
            };
            t1_size_map.insert(v, vsize);
            for c in tree.postord_ids(tree.get_root_id()) {
                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).unwrap() == tree.get_node_taxa(c).unwrap() {
                            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).unwrap();
                        intersection += intersection_map.get(&(v, cch)).unwrap();
                    }
                }
                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 zeta of leaf by taxa
    fn get_zeta_taxa(&self, taxa: &TreeNodeMeta<Self>) -> TreeNodeZeta<Self> {
        self.get_zeta(self.get_taxa_node_id(taxa).unwrap()).unwrap()
    }

    /// 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)
                    .unwrap()
                    .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)
                    .unwrap()
                    .powi(-1),
            )
    }

    /// Reurns the cophenetic distance between two trees using the naive algorithm (\Theta(n^2))
    fn cophen_dist<'a>(&'a self, tree: &'a Self, norm: u32) -> TreeNodeZeta<Self> {
        if !self.is_all_zeta_set() || !tree.is_all_zeta_set() {
            panic!("Zeta values not set");
        }
        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))
    fn cophen_dist_par<'a>(&'a self, tree: &'a Self, norm: u32) -> TreeNodeZeta<Self> {
        if !self.is_all_zeta_set() || !tree.is_all_zeta_set() {
            panic!("Zeta values not set");
        }
        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,
    ) -> TreeNodeZeta<Self> {
        // 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]);
                    (zeta_1 - zeta_2).abs()
                }
                false => {
                    let self_ids = x
                        .iter()
                        .map(|a| self.get_taxa_node_id(a).unwrap())
                        .collect_vec();
                    let tree_ids = x
                        .iter()
                        .map(|a| tree.get_taxa_node_id(a).unwrap())
                        .collect_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).unwrap();
                    let zeta_2 = tree.get_zeta(t_hat_lca_id).unwrap();
                    (zeta_1 - zeta_2).abs()
                }
            })
            .collect();

        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>>,
    ) -> TreeNodeZeta<Self> {
        // 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]);
                    (zeta_1 - zeta_2).abs()
                }
                false => {
                    let self_ids = x
                        .iter()
                        .map(|a| self.get_taxa_node_id(a).unwrap())
                        .collect_vec();
                    let tree_ids = x
                        .iter()
                        .map(|a| tree.get_taxa_node_id(a).unwrap())
                        .collect_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).unwrap();
                    let zeta_2 = tree.get_zeta(t_hat_lca_id).unwrap();
                    (zeta_1 - zeta_2).abs()
                }
            });

        Self::compute_norm(cophen_vec, norm)
    }
}