obikmer/src/obikphylo/src/siblings_old/stats.rs

use std::sync::Arc;

use rayon::prelude::*;

use obisys::progress_bar;

use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};

use super::ANNEX_FILE_NAME;
use super::SiblingAnnex;
use super::cache::PartitionCache;
use super::family_scan::{Selection, scan_layer_families};

/// Distribution of family sizes (1-4), read back from an already-built
/// annex (see [`super::build::SiblingAnnexBuildExt::build_sibling_annex`])
/// plus the index's presence/count data — a separate, occasional diagnostic
/// pass, not fused into construction.
///
/// Every count here is **per family, not per slot**: a family with `F`
/// members occupies `F` annex slots (one per observed member), all sharing
/// the same mask. Counting every slot would count each family up to 4
/// times over; only the minorant's slot is tallied (minorant is derived on
/// the fly — see `is_minorant` — not stored, but cheap: no lookup, pure
/// bit arithmetic on already-in-hand data).
#[derive(Debug, Clone, Default)]
pub struct SiblingAnnexStats {
    /// `counts[s]` = number of *families* of size `s + 1` (index 0 = size 1,
    /// i.e. 0 siblings, ... index 3 = size 4, i.e. 3 siblings).
    pub counts: [u64; 4],
    /// `per_genome[g][s]` = number of families of size `s + 1` for which
    /// genome `g` (index into `KmerIndex::meta().genomes`) carries at least
    /// one member.
    pub per_genome: Vec<[u64; 4]>,
}

/// Adds [`sibling_family_size_histogram`](Self::sibling_family_size_histogram)
/// and [`sibling_annex_stats`](Self::sibling_annex_stats) to `KmerIndex`.
pub trait SiblingStatsExt {
    /// The global family-size histogram alone (`SiblingAnnexStats::counts`,
    /// no `per_genome`) — reads only the already-built annex (`mask.siblings()`
    /// + `mask.is_minorant()`, 1 byte/slot, mmap'd), nothing else: no
    /// `unitigs.bin` scan, no MPHF lookup, no `PartitionCache`. Cost is a
    /// linear scan of one small file per layer, independent of the rest of
    /// the index's size or how much of it is paged in. An earlier version
    /// re-derived minorant-ness per slot (re-reading `unitigs.bin`, hashing
    /// every k-mer through the MPHF again to reconstruct what
    /// `build_sibling_annex` already knew) — sampling a real run showed
    /// `MphfLayer::find` alone at 71% of wall-clock time for what should be
    /// a near-instant four-bucket count. `mask.is_minorant()` is that same
    /// fact, computed once at construction (see `build_layer_sibling_annex`)
    /// and stored in the annex's own spare bits — free to read back.
    /// [`sibling_annex_stats`](Self::sibling_annex_stats) computes the same
    /// `counts` but pays the full per-genome cross-partition resolution
    /// cost to get there too; use this when only the global histogram is
    /// needed. Requires an annex built after the minorant flag was added —
    /// re-run `build_sibling_annex` if this reads as all-zero on an older one.
    fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]>;

    /// Tally the family-size distribution of an already-built annex
    /// (globally, and per genome), counting each family once (at its
    /// minorant slot). Errors if
    /// [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex)
    /// has not been run on this index first.
    fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats>;
}

impl SiblingStatsExt for KmerIndex {
    fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]> {
        let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;

        // Each layer's annex file is independent — no shared `PartitionCache`
        // and no `scan_layer_families` partition-grouped locality to protect
        // (unlike `sibling_annex_stats`/`distance`/`cardinality_tally`), so
        // layers can be scanned concurrently.
        layer_dirs
            .par_iter()
            .try_fold(
                || [0u64; 4],
                |mut counts, layer_dir| -> OKIResult<[u64; 4]> {
                    let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
                    for slot in 0..annex.len() {
                        let Some(mask) = annex.get(slot) else {
                            continue;
                        };
                        if !mask.is_minorant() {
                            continue; // family tallied once, at its minorant
                        }
                        counts[mask.siblings() as usize] += 1;
                    }
                    Ok(counts)
                },
            )
            .try_reduce(
                || [0u64; 4],
                |a, b| {
                    let mut sum = a;
                    for i in 0..4 {
                        sum[i] += b[i];
                    }
                    Ok(sum)
                },
            )
    }

    fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats> {
        let n_parts = self.n_partitions();
        let n_genomes = self.meta().genomes().map_err(OKIError::Io)?.len();
        let with_counts = self.meta().config.with_counts;
        let k = self.kmer_size();
        // Same whole-run cache as `build_sibling_annex` — see its docs for
        // why re-opening per lookup (or per call to a batching helper) is
        // not good enough on a real index.
        let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
        let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;

        // One layer at a time, not parallelised across layers — see
        // `snp_pseudo_alignment`'s comment for why `par_iter()` over layers
        // would defeat `scan_layer_families`'s partition-grouped locality.
        let pb = progress_bar("sibling_annex_stats", layer_dirs.len() as u64, "layers");
        let mut stats = SiblingAnnexStats {
            per_genome: vec![[0u64; 4]; n_genomes],
            ..Default::default()
        };
        for layer_dir in &layer_dirs {
            scan_layer_families(
                layer_dir,
                n_parts,
                n_genomes,
                with_counts,
                k,
                &cache,
                &Selection::All,
                |_family_idx, mask, genome_mask| {
                    // "Genome g represents this family" means g carries
                    // *any* of its members, not just the minorant's own —
                    // `genome_mask[g] != 0` is exactly that.
                    let s = mask.siblings() as usize;
                    stats.counts[s] += 1;
                    for (g, &m) in genome_mask.iter().enumerate() {
                        if m != 0 {
                            stats.per_genome[g][s] += 1;
                        }
                    }
                },
            )?;
            pb.inc(1);
        }
        pb.finish_and_clear();

        Ok(stats)
    }
}