obikmer/src/obikindex/src/distance.rs

use ndarray::Array2;
use obicompactvec::traits::{BitPartials, CountPartials};
use obilayeredmap::LayeredStore;
use rayon::prelude::*;

use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;

// ── Public API ────────────────────────────────────────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DistanceMetric {
    /// Jaccard distance on presence/absence data.
    Jaccard,
    /// Hamming distance (number of differing kmer positions) on presence/absence data.
    Hamming,
    /// Bray-Curtis dissimilarity on raw counts.
    BrayCurtis,
    /// Bray-Curtis dissimilarity normalised by per-genome total counts.
    RelfreqBrayCurtis,
    /// Euclidean distance on raw counts.
    Euclidean,
    /// Euclidean distance on relative frequencies.
    RelfreqEuclidean,
    /// Hellinger distance on counts.
    Hellinger,
    /// Euclidean distance in the Hellinger (√relative-frequency) space (unnormalised variant).
    HellingerEuclidean,
}

pub struct DistanceOutput {
    /// n×n pairwise distance matrix (genomes in index order).
    pub matrix: Array2<f64>,
    /// n×n shared-kmer count matrix (intersection), if requested.
    pub shared_kmers: Option<Array2<u64>>,
}

impl DistanceMetric {
    pub fn requires_counts(self) -> bool {
        matches!(
            self,
            DistanceMetric::BrayCurtis
            | DistanceMetric::RelfreqBrayCurtis
            | DistanceMetric::Euclidean
            | DistanceMetric::RelfreqEuclidean
            | DistanceMetric::Hellinger
            | DistanceMetric::HellingerEuclidean
        )
    }
}

// ── KmerIndex::distance ───────────────────────────────────────────────────────

impl KmerIndex {
    pub fn distance(&self, metric: DistanceMetric, shared_kmers: bool, presence_threshold: u32) -> OKIResult<DistanceOutput> {
        let n_genomes = self.meta.genomes.len();
        if n_genomes < 2 {
            return Err(OKIError::InvalidInput(
                "distance requires at least 2 genomes in the index".into(),
            ));
        }

        let use_counts = self.meta.config.with_counts;
        if metric.requires_counts() && !use_counts {
            return Err(OKIError::InvalidInput(format!(
                "{metric:?} requires a count index (with_counts = true)"
            )));
        }

        let n_parts = self.n_partitions();

        if use_counts {
            let stores: Vec<_> = (0..n_parts)
                .into_par_iter()
                .map(|i| self.partition.count_store(i).map_err(OKIError::Partition))
                .collect::<OKIResult<_>>()?;
            let global = LayeredStore::new(stores);

            let matrix = match metric {
                DistanceMetric::BrayCurtis        => CountPartials::bray_dist_matrix(&global),
                DistanceMetric::RelfreqBrayCurtis => CountPartials::relfreq_bray_dist_matrix(&global),
                DistanceMetric::Euclidean         => CountPartials::euclidean_dist_matrix(&global),
                DistanceMetric::RelfreqEuclidean  => CountPartials::relfreq_euclidean_dist_matrix(&global),
                DistanceMetric::Hellinger         => CountPartials::hellinger_dist_matrix(&global),
                DistanceMetric::HellingerEuclidean => CountPartials::hellinger_euclidean_dist_matrix(&global),
                DistanceMetric::Jaccard           => CountPartials::threshold_jaccard_dist_matrix(&global, presence_threshold),
                DistanceMetric::Hamming           => {
                    return Err(OKIError::InvalidInput(
                        "Hamming is only available for presence/absence indexes".into(),
                    ));
                }
            };

            let shared = if shared_kmers {
                let (inter, _) = CountPartials::partial_threshold_jaccard(&global, presence_threshold);
                Some(inter)
            } else {
                None
            };

            Ok(DistanceOutput { matrix, shared_kmers: shared })
        } else {
            let stores: Vec<_> = (0..n_parts)
                .into_par_iter()
                .map(|i| self.partition.presence_store(i).map_err(OKIError::Partition))
                .collect::<OKIResult<_>>()?;
            let global = LayeredStore::new(stores);

            let matrix = match metric {
                DistanceMetric::Jaccard => BitPartials::jaccard_dist_matrix(&global),
                DistanceMetric::Hamming => {
                    BitPartials::hamming_dist_matrix(&global).mapv(|v| v as f64)
                }
                other => {
                    return Err(OKIError::InvalidInput(format!(
                        "{other:?} requires a count index; use --metric jaccard or --metric hamming"
                    )));
                }
            };

            let shared = if shared_kmers {
                let (inter, _) = BitPartials::partial_jaccard(&global);
                Some(inter)
            } else {
                None
            };

            Ok(DistanceOutput { matrix, shared_kmers: shared })
        }
    }
}