obikmer/src/obikindex/src/stats.rs

use std::fs;
use std::path::Path;

use obicompactvec::{LayerMeta, PersistentBitMatrix, PersistentCompactIntMatrix};
use obicompactvec::traits::ColumnWeights;
use obilayeredmap::meta::PartitionMeta;
use rayon::prelude::*;

use crate::error::OKIResult;
use crate::index::KmerIndex;

/// Bits per kmer broken down by index component.
pub struct IndexBitsPerKmer {
    /// Total distinct k-mers across all partitions and layers.
    pub n_kmers:          usize,
    /// Number of genomes in the index.
    pub n_genomes:        usize,
    /// Bits used by the minimal perfect hash function (`mphf.bin`).
    pub mphf:             f64,
    /// Bits used by the evidence files (`evidence.bin`, `unitigs.bin*`,
    /// `fingerprint.bin`).
    pub evidence:         f64,
    /// Bits used by the count/presence matrices (`counts/` and `presence/`),
    /// normalised by k-mers only.
    pub matrix:           f64,
    /// `matrix` divided by the number of genomes — intrinsic encoding
    /// efficiency, independent of index size.
    pub matrix_per_genome: f64,
    /// Sum of mphf + evidence + matrix.
    pub total:            f64,
}

// ── File-size helpers ─────────────────────────────────────────────────────────

fn file_bytes(path: &Path) -> u64 {
    fs::metadata(path).map(|m| m.len()).unwrap_or(0)
}

fn dir_bytes(dir: &Path) -> u64 {
    if !dir.exists() { return 0; }
    fs::read_dir(dir)
        .map(|entries| {
            entries
                .filter_map(|e| e.ok())
                .filter_map(|e| e.metadata().ok())
                .filter(|m| m.is_file())
                .map(|m| m.len())
                .sum()
        })
        .unwrap_or(0)
}

// ── Per-layer accounting ──────────────────────────────────────────────────────

struct LayerBytes {
    n_kmers:  usize,
    mphf:     u64,
    evidence: u64,
    matrix:   u64,
}

fn layer_bytes(layer_dir: &Path) -> LayerBytes {
    let n_kmers = LayerMeta::load(layer_dir).map(|m| m.n).unwrap_or(0);

    let mphf = file_bytes(&layer_dir.join("mphf.bin"));

    let evidence = file_bytes(&layer_dir.join("unitigs.bin"))
        + file_bytes(&layer_dir.join("unitigs.bin.idx"))
        + file_bytes(&layer_dir.join("evidence.bin"))
        + file_bytes(&layer_dir.join("fingerprint.bin"));

    let matrix = dir_bytes(&layer_dir.join("counts"))
        + dir_bytes(&layer_dir.join("presence"));

    LayerBytes { n_kmers, mphf, evidence, matrix }
}

// ── KmerIndex::bits_per_kmer ──────────────────────────────────────────────────

impl KmerIndex {
    /// Compute bits-per-kmer statistics for the built index.
    ///
    /// File sizes are read directly from disk; kmer counts come from
    /// `layer_meta.json` (no need to scan the MPHF or unitig files).
    /// Computation is parallelised across partitions.
    pub fn bits_per_kmer(&self) -> OKIResult<IndexBitsPerKmer> {
        let n          = self.n_partitions();
        let n_genomes  = self.meta().genomes.len().max(1);

        let (n_kmers, mphf_b, evidence_b, matrix_b) = (0..n)
            .into_par_iter()
            .map(|i| {
                let index_dir = self.partition.part_dir(i).join("index");
                if !index_dir.exists() { return (0usize, 0u64, 0u64, 0u64); }

                let n_layers = PartitionMeta::load(&index_dir)
                    .map(|m| m.n_layers)
                    .unwrap_or(0);

                (0..n_layers).fold((0usize, 0u64, 0u64, 0u64), |acc, l| {
                    let lb = layer_bytes(&index_dir.join(format!("layer_{l}")));
                    (acc.0 + lb.n_kmers, acc.1 + lb.mphf, acc.2 + lb.evidence, acc.3 + lb.matrix)
                })
            })
            .reduce(|| (0, 0, 0, 0), |a, b| (a.0 + b.0, a.1 + b.1, a.2 + b.2, a.3 + b.3));

        if n_kmers == 0 {
            return Ok(IndexBitsPerKmer {
                n_kmers: 0, n_genomes,
                mphf: 0.0, evidence: 0.0,
                matrix: 0.0, matrix_per_genome: 0.0, total: 0.0,
            });
        }

        let bpk    = |bytes: u64| bytes as f64 * 8.0 / n_kmers as f64;
        let matrix = bpk(matrix_b);

        Ok(IndexBitsPerKmer {
            n_kmers,
            n_genomes,
            mphf:              bpk(mphf_b),
            evidence:          bpk(evidence_b),
            matrix,
            matrix_per_genome: matrix / n_genomes as f64,
            total:             bpk(mphf_b + evidence_b + matrix_b),
        })
    }

    /// Return `(total_distinct_kmers, per_genome_kmer_counts)`.
    ///
    /// For each genome, the count is the number of distinct k-mers for which
    /// that genome has a non-zero value (presence = 1, count > 0).
    /// Partitions are scanned in parallel; results are summed across partitions.
    pub fn genome_kmer_counts(&self) -> OKIResult<(usize, Vec<u64>)> {
        let n          = self.n_partitions();
        let n_genomes  = self.meta.genomes.len();

        let partials: Vec<(usize, Vec<u64>)> = (0..n)
            .into_par_iter()
            .map(|i| {
                let mut counts  = vec![0u64; n_genomes];
                let mut n_kmers = 0usize;

                let index_dir = self.partition.part_dir(i).join("index");
                if !index_dir.exists() { return (0, counts); }

                let n_layers = PartitionMeta::load(&index_dir)
                    .map(|m| m.n_layers)
                    .unwrap_or(0);

                for l in 0..n_layers {
                    let layer_dir = index_dir.join(format!("layer_{l}"));
                    if !layer_dir.exists() { continue; }

                    n_kmers += LayerMeta::load(&layer_dir).map(|m| m.n).unwrap_or(0);

                    let mat: Box<dyn ColumnWeights> =
                        if layer_dir.join("counts").exists()
                            && !layer_dir.join("presence").exists()
                        {
                            match PersistentCompactIntMatrix::open(&layer_dir) {
                                Ok(m)  => Box::new(m),
                                Err(_) => continue,
                            }
                        } else {
                            match PersistentBitMatrix::open(&layer_dir) {
                                Ok(m)  => Box::new(m),
                                Err(_) => continue,
                            }
                        };
                    let col_counts = mat.partial_kmer_counts();

                    for (c, &v) in col_counts.iter().enumerate() {
                        if c < n_genomes { counts[c] += v; }
                    }
                }

                (n_kmers, counts)
            })
            .collect();

        let total_kmers: usize = partials.iter().map(|(n, _)| n).sum();
        let mut total_counts   = vec![0u64; n_genomes];
        for (_, counts) in partials {
            for (i, v) in counts.into_iter().enumerate() {
                total_counts[i] += v;
            }
        }

        Ok((total_kmers, total_counts))
    }
}