obikmer/src/obilayeredmap/src/layer.rs

use std::collections::HashMap;
use std::fs;
use std::path::Path;

use obicompactvec::{
    BinaryMatrix,
    PersistentBitMatrix, PersistentBitMatrixBuilder,
    PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
    PersistentSparseBitMatrix,
};
use obikseq::CanonicalKmer;
use obiskio::{UnitigFileReader, UnitigFileWriter};

use crate::error::{OLMError, OLMResult};
use crate::meta::IndexMode;
use crate::mphf_layer::MphfLayer;
pub(crate) use crate::mphf_layer::UNITIGS_FILE;

const COUNTS_DIR:   &str = "counts";
const PRESENCE_DIR: &str = "presence";

// ── Trait ─────────────────────────────────────────────────────────────────────

pub trait LayerData: Sized {
    type Item;
    fn open(layer_dir: &Path) -> OLMResult<Self>;
    fn read(&self, slot: usize) -> Self::Item;
}

impl LayerData for () {
    type Item = ();
    fn open(_layer_dir: &Path) -> OLMResult<Self> { Ok(()) }
    fn read(&self, _slot: usize) {}
}

impl LayerData for PersistentCompactIntMatrix {
    type Item = Box<[u32]>;
    fn open(layer_dir: &Path) -> OLMResult<Self> {
        PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io)
    }
    fn read(&self, slot: usize) -> Box<[u32]> { self.row(slot) }
}

impl LayerData for PersistentBitMatrix {
    type Item = Box<[bool]>;
    fn open(layer_dir: &Path) -> OLMResult<Self> {
        PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io)
    }
    fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) }
}

impl LayerData for PersistentSparseBitMatrix {
    type Item = Box<[bool]>;
    fn open(layer_dir: &Path) -> OLMResult<Self> {
        PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OLMError::Io)
    }
    fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) }
}

// ── Structures ────────────────────────────────────────────────────────────────

pub struct Layer<D: LayerData = ()> {
    mphf: MphfLayer,
    data: D,
}

pub struct Hit<T = ()> {
    pub slot: usize,
    pub data: T,
}

// ── Common read path ──────────────────────────────────────────────────────────

impl<D: LayerData> Layer<D> {
    pub fn open(path: &Path, mode: &IndexMode) -> OLMResult<Self> {
        let mphf = MphfLayer::open(path, mode)?;
        let data = D::open(path)?;
        Ok(Self { mphf, data })
    }

    pub fn query(&self, kmer: CanonicalKmer) -> Option<Hit<D::Item>> {
        self.mphf.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot) })
    }

    /// MPHF + evidence membership check only — no data read. For callers
    /// that batch many lookups before touching the matrix at all (e.g. to
    /// group hits by column for a later `sub_matrix`/`fill_sub_matrix`
    /// sweep), so a plain `query` reading — and discarding — a full row
    /// per lookup would be wasted work.
    pub fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
        self.mphf.find(kmer)
    }

    pub fn n(&self) -> usize { self.mphf.n() }

    /// Raw MPHF lookup: kmer → slot, no membership check.
    pub fn index(&self, kmer: CanonicalKmer) -> usize {
        self.mphf.index(kmer)
    }

    /// Batch raw MPHF lookup: kmers → slots, no membership check.
    pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
        self.mphf.index_batch(kmers)
    }

    /// Iterate over all canonical kmers in the layer, in deterministic order.
    pub fn iter_kmers(&self) -> crate::mphf_layer::KmerIter {
        self.mphf.iter_kmers()
    }

    /// Iterate over all canonical kmers, each paired with its zero-based
    /// sequence index in `unitigs.bin`.
    pub fn enumerate_kmers(&self) -> std::iter::Enumerate<crate::mphf_layer::KmerIter> {
        self.mphf.enumerate_kmers()
    }

    /// Iterate over the layer's canonical kmers in batches of `n`.
    pub fn iter_kmers_batch(&self, n: usize) -> crate::mphf_layer::KmerBatchIter {
        self.mphf.iter_kmers_batch(n)
    }

    /// Iterate over batches, each paired with the zero-based index of the
    /// first kmer in the batch.
    pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
        self.mphf.enumerate_kmers_batch(n)
    }

    pub fn unitig_writer(out_dir: &Path) -> OLMResult<UnitigFileWriter> {
        MphfLayer::unitig_writer(out_dir)
    }

    /// Build `unitigs.bin.idx` and `evidence.bin` from `unitigs.bin` and
    /// `mphf.bin` already present in `layer_dir`.
    /// `block_bits` controls the `.idx` block size (2^block_bits chunks/block).
    pub fn build_exact_evidence(layer_dir: &Path, block_bits: u8) -> OLMResult<usize> {
        MphfLayer::build_exact_evidence(layer_dir, block_bits)
    }

    /// Build `fingerprint.bin` from `unitigs.bin` and `mphf.bin` already
    /// present in `layer_dir`.  `b` — fingerprint bits (1..=64); `z` — Findere
    /// consecutive k-mer parameter (≥1).
    pub fn build_approx_evidence(layer_dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
        MphfLayer::build_approx_evidence(layer_dir, b, z)
    }

}

// ── Mode 1 — set membership ───────────────────────────────────────────────────

impl Layer<()> {
    pub fn build(out_dir: &Path, block_bits: u8, mode: &IndexMode) -> OLMResult<usize> {
        MphfLayer::build(out_dir, block_bits, mode, &mut |_, _| Ok(()))
    }

    /// Create a presence matrix for a set-membership layer (first merge).
    pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> {
        let presence_dir = layer_dir.join(PRESENCE_DIR);
        fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?;
        let mut mb = PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?;
        mb.add_col_ones().map_err(OLMError::Io)?.close().map_err(OLMError::Io)?;
        mb.close().map_err(OLMError::Io)
    }
}

// ── Mode 2 — count matrix ─────────────────────────────────────────────────────

impl Layer<PersistentCompactIntMatrix> {
    pub fn build(
        out_dir: &Path,
        block_bits: u8,
        mode: &IndexMode,
        count_of: impl Fn(CanonicalKmer) -> u32,
    ) -> OLMResult<usize> {
        let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
        let counts_dir = out_dir.join(COUNTS_DIR);
        let mut mb = PersistentCompactIntMatrixBuilder::new(n, &counts_dir)
            .map_err(OLMError::Io)?;
        let mut col = mb.add_col().map_err(OLMError::Io)?;
        let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
            col.set(slot, count_of(kmer));
            Ok(())
        })?;
        col.close().map_err(OLMError::Io)?;
        mb.close().map_err(OLMError::Io)?;
        Ok(n_built)
    }

    pub fn build_from_map(
        out_dir: &Path,
        block_bits: u8,
        mode: &IndexMode,
        counts: &HashMap<CanonicalKmer, u32>,
    ) -> OLMResult<usize> {
        Self::build(out_dir, block_bits, mode, |kmer| counts.get(&kmer).copied().unwrap_or(0))
    }
}

// ── Mode 2 — count matrix column append ──────────────────────────────────────

impl Layer<PersistentCompactIntMatrix> {
    pub fn append_genome_column(
        layer_dir: &Path,
        value_of: impl Fn(usize) -> u32,
    ) -> OLMResult<()> {
        PersistentCompactIntMatrix::append_column(&layer_dir.join(COUNTS_DIR), value_of)
            .map_err(OLMError::Io)
    }

    /// Number of genome columns in this layer's count matrix.
    pub fn n_cols(&self) -> usize {
        self.data.n_cols()
    }

    /// Extract a sub-matrix of counts for the rows at `slots`.
    ///
    /// Returns a column-first `Vec<Vec<u32>>`: one inner `Vec` per genome
    /// column, containing the counts for the requested slots in order.
    /// Column access is sequential for cache efficiency.
    pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<u32>> {
        self.data.sub_matrix(slots)
    }

    /// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
    /// buffers to avoid allocating the outer `Vec`.
    ///
    /// `out` must have length equal to the number of genome columns.  Each
    /// `out[c]` is cleared, resized to `slots.len()`, and filled with the
    /// counts for column `c` in the same order as `slots`.
    pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<u32>]) {
        self.data.fill_sub_matrix(slots, out)
    }
}

// ── Mode 3 — presence/absence matrix ─────────────────────────────────────────

// ── Mode 3 — presence/absence matrix, generic over dense/sparse storage ──────
//
// `n_cols`/`sub_matrix`/`fill_sub_matrix` are identical in shape for any
// `D` satisfying `BinaryMatrix` — genuinely generic (not just two
// near-identical impl blocks) so `Layer<PersistentSparseBitMatrix>` gets
// them for free, matching `Layer<PersistentBitMatrix>` at the same call
// sites (see `obikphylo::siblings::cache::Mat`). Construction
// (`append_genome_column`/`build_presence`) stays `PersistentBitMatrix`-only
// below: sparse matrices aren't built column-by-column, they're built
// row-by-row from an already-built dense layer
// (`PersistentSparseBitMatrixBuilder::build_from_dense`).
impl<D: LayerData<Item = Box<[bool]>> + BinaryMatrix> Layer<D> {
    /// Number of genome columns in this layer's presence matrix — see
    /// `PersistentBitMatrix::n_cols`'s docs for the `Implicit` mono-genome
    /// special case (always reports `1`, regardless of the index's real
    /// genome count).
    pub fn n_cols(&self) -> usize {
        self.data.n_cols()
    }

    /// Extract a sub-matrix of presence/absence for the rows at `slots`.
    ///
    /// Returns a column-first `Vec<Vec<bool>>`: one inner `Vec` per genome
    /// column, containing the presence bits for the requested slots in order.
    /// Column access is sequential for cache efficiency on the dense
    /// storage — on sparse storage it's a naive row-by-row decode, see
    /// `docmd/architecture/siblings.md`'s sparse-matrix section.
    pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<bool>> {
        self.data.sub_matrix(slots)
    }

    /// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
    /// buffers to avoid allocating the outer `Vec`.
    ///
    /// `out` must have length equal to the number of genome columns.  Each
    /// `out[c]` is cleared, resized to `slots.len()`, and filled with the
    /// presence bits for column `c` in the same order as `slots`.
    pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<bool>]) {
        self.data.fill_sub_matrix(slots, out)
    }
}

impl Layer<PersistentBitMatrix> {
    pub fn append_genome_column(
        layer_dir: &Path,
        value_of: impl Fn(usize) -> bool,
    ) -> OLMResult<()> {
        PersistentBitMatrix::append_column(&layer_dir.join(PRESENCE_DIR), value_of)
            .map_err(OLMError::Io)
    }

    pub fn build_presence(
        out_dir: &Path,
        block_bits: u8,
        mode: &IndexMode,
        n_genomes: usize,
        present_in: impl Fn(CanonicalKmer, usize) -> bool,
    ) -> OLMResult<usize> {
        let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
        let presence_dir = out_dir.join(PRESENCE_DIR);
        let mut mb = PersistentBitMatrixBuilder::new(n, &presence_dir).map_err(OLMError::Io)?;
        let mut cols: Vec<_> = (0..n_genomes)
            .map(|_| mb.add_col().map_err(OLMError::Io))
            .collect::<OLMResult<_>>()?;
        let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
            for (g, col) in cols.iter_mut().enumerate() {
                col.set(slot, present_in(kmer, g));
            }
            Ok(())
        })?;
        for col in cols {
            col.close().map_err(OLMError::Io)?;
        }
        mb.close().map_err(OLMError::Io)?;
        Ok(n_built)
    }
}

#[cfg(test)]
#[path = "tests/layer.rs"]
mod tests;