//! Family presence-mask annex construction.
//!
//! See `docmd/theory/evolutionary_distances.md`, "Definitions: family, and
//! the canonical form of a family" and "Step 2b", for the full design
//! discussion this implements.
//!
//! For each distinct k-mer of each layer of the (already built/merged)
//! index, computes a 4-bit presence mask for its "family" (the up to 4
//! k-mers sharing its flanks, differing only at the central base —
//! well-defined for odd k): bit `b` set iff the family member whose own
//! canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed anywhere
//! in the current multi-genome index — a property of the whole index, not
//! of any one genome. Sibling count and minorant are *derived* from the
//! mask by callers, not stored (see `FamilyMask` and
//! [`sibling_annex_stats`](KmerIndex::sibling_annex_stats) below).
//!
//! Per layer, an `obipipeline` batch transform (throttled — see
//! `obipipeline::throttle`) generates a whole batch's central variants at
//! once (`BATCH_SIZE` source k-mers in, that batch's variants out as one
//! pipeline message), interleaved across many in-flight batches by the
//! scheduler's shared worker pool rather than processed on a single
//! thread. The actual cross-partition lookup reuses a `PartitionCache` of
//! every partition's already-open MPHF layers, built once for the whole
//! `build_sibling_annex` run, rather than reopening files per lookup or
//! per source layer. Two earlier, coarser-grained designs were tried and
//! measured (not guessed) to be worse, in order: (1) reopening/re-mmap'ing
//! every target partition's files on every single lookup — fine at toy
//! scale, ~90% system time against a real index; (2) a `Flat` pipeline
//! stage pushing one message per generated *variant* (up to 3 per source
//! k-mer) — cheaper than reopening files, but sampling a real run showed
//! most wall-clock time going into per-message channel send/notify
//! syscalls rather than the lookup itself, because a single k-mer's ≤3
//! variants is far too fine a granularity to amortise a pipeline's
//! synchronisation cost over. See `docmd/theory/evolutionary_distances.md`,
//! Step 2b, "Mechanism".

use std::path::Path;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::Arc;

use ndarray::Array2;
use rayon::prelude::*;

use obicompactvec::{
    FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder,
};
use obikpartitionner::KmerPartition;
use obipipeline::ThrottleGuard;
use obikseq::{CanonicalKmer, Minimizer};
use obilayeredmap::{MphfLayer, OLMError};
use obilayeredmap::meta::PartitionMeta;
use obiskbuilder::rolling_stat::RollingStat;
use obiskio::UnitigFileReader;
use obisys::progress_bar;

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

const INDEX_SUBDIR: &str = "index";
const ANNEX_FILE_NAME: &str = "siblings.psib";

fn olm_to_ok(e: OLMError) -> OKIError {
    match e {
        OLMError::Io(e) => OKIError::Io(e),
        other => OKIError::InvalidInput(format!("layered-map error: {other}")),
    }
}

/// Central-position base of a canonical k-mer, in the fixed 0=A/1=C/2=G/3=T
/// encoding — the mask's bit index. `k` must be odd (project invariant).
#[inline]
fn central_base(kmer: CanonicalKmer, k: usize) -> u8 {
    kmer.nucleotide((k - 1) / 2)
}

/// Is `kmer` the minorant of its family, given the family's presence mask?
/// Regenerates the family's 4 canonical forms from `kmer` itself (cheap, no
/// lookup — see the design doc's "Definitions" section for why this is
/// always safe: the set of 4 forms is invariant regardless of which member
/// you start from), and compares the raw encodings of whichever are marked
/// present in `mask`.
fn is_minorant(kmer: CanonicalKmer, mask: FamilyMask, k: usize) -> bool {
    kmer.central_canonical_neighbors().into_iter().all(|other| {
        other == kmer || !mask.has(central_base(other, k)) || kmer.raw() <= other.raw()
    })
}


/// Minimiser of a single, isolated canonical k-mer (not part of a streamed
/// sequence). `RollingStat` computes minimisers incrementally along a
/// sequence; this feeds one k-mer's bases through a fresh instance to get
/// the same selection for a single, disconnected k-mer. Not the leanest
/// possible primitive (an O(1)-amortised dedicated scan, as originally
/// sketched in the design doc's Step 0, would avoid the ASCII round-trip and
/// `RollingStat` allocation) but correct and reuses already-tested logic;
/// left as a follow-up optimisation.
fn lone_kmer_minimizer(kmer: CanonicalKmer) -> Minimizer {
    let ascii = kmer.to_ascii();
    let mut rs = RollingStat::new(0);
    for b in ascii {
        rs.push(b);
    }
    rs.canonical_minimizer()
        .expect("RollingStat must be ready after k bases of a valid k-mer")
}

/// Destination partition for a (possibly synthetic) canonical k-mer, using
/// the same routing rule as the rest of the index (`minimiser.seq_hash() &
/// mask`, `n_partitions` is a power of two).
fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize {
    let mask = (n_partitions as u64) - 1;
    (lone_kmer_minimizer(kmer).seq_hash() & mask) as usize
}

// ── obipipeline data types ─────────────────────────────────────────────────

/// A batch of this layer's distinct k-mers (local MPHF slot + k-mer), the
/// pipeline's source item — batched, not one k-mer per item, so that
/// pipeline messages and their synchronisation cost stay amortised over
/// thousands of lookups (see `build_layer_sibling_annex`'s comment on
/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved
/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform,
/// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
struct SourceBatch {
    items: Vec<(usize, CanonicalKmer)>,
    _permit: ThrottleGuard,
}

/// One batch's worth of central-substitution variants (up to 3 per source
/// k-mer), each already routed to its destination partition and carrying
/// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a
/// hit. `(dest_partition, variant, source_slot, base)` per entry.
struct VariantBatch {
    items: Vec<(usize, CanonicalKmer, usize, u8)>,
    _permit: ThrottleGuard,
}

enum SibData {
    Batch(SourceBatch),
    Variants(VariantBatch),
}

/// Every partition's already-open MPHF layers, built **once** for the whole
/// `build_sibling_annex` run and shared (read-only) across every lookup, in
/// every source layer, for the rest of the run — not reopened/re-mmap'd per
/// query, nor per source layer.
///
/// Confirmed necessary by sampling a real run: routing lookups through
/// `KmerPartition::query_partition_with` (the same batching `obikmer query`
/// uses) still reopens+re-mmaps every target partition's files on every
/// call, and it is called once per destination partition **per source
/// layer** — for an index with many layers this repeats the same
/// `MphfLayer::open`/`Evidence::open`/`PersistentBitMatrix::open` work over
/// and over. Parallelising those calls (see the gather step below) spread
/// the redundant work across more cores but did not reduce it: sampling
/// showed Rayon workers spending their time inside repeated `open()`
/// syscalls, not computation. This cache amortises that cost to once per
/// partition for the entire run, regardless of how many source layers or
/// lookups follow.
/// A cached layer's opened presence/count matrix, alongside its `MphfLayer`.
enum Mat {
    Count(PersistentCompactIntMatrix),
    Presence(PersistentBitMatrix),
}

impl Mat {
    fn n_cols(&self) -> usize {
        match self {
            Mat::Count(m) => m.n_cols(),
            Mat::Presence(m) => m.n_cols(),
        }
    }
    fn carries(&self, g: usize, slot: usize) -> bool {
        match self {
            Mat::Count(m) => m.col_view(g).get(slot) != 0,
            Mat::Presence(m) => m.get(g, slot) != 0,
        }
    }
}

struct PartitionCache {
    /// `layers[partition][layer]` = that partition's opened MPHF layers,
    /// paired 1:1 with `mats[partition][layer]`; empty if the partition
    /// directory doesn't exist. Used by both [`KmerIndex::build_sibling_annex`]
    /// (`layers` only) and [`KmerIndex::sibling_annex_stats`] (both).
    layers: Vec<Vec<MphfLayer>>,
    mats: Vec<Vec<Mat>>,
}

impl PartitionCache {
    fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult<Self> {
        let pb = progress_bar("open_partitions", n_parts as u64, "partitions");
        let built: Vec<(Vec<MphfLayer>, Vec<Mat>)> = (0..n_parts)
            .into_par_iter()
            .map(|part| -> OKIResult<(Vec<MphfLayer>, Vec<Mat>)> {
                let index_dir = partition.part_dir(part).join(INDEX_SUBDIR);
                if !index_dir.exists() {
                    pb.inc(1);
                    return Ok((Vec::new(), Vec::new()));
                }
                let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
                let mut layers = Vec::with_capacity(meta.n_layers);
                let mut mats = Vec::with_capacity(meta.n_layers);
                for l in 0..meta.n_layers {
                    let layer_dir = index_dir.join(format!("layer_{l}"));
                    let Ok(mphf) = MphfLayer::open(&layer_dir, &meta.mode) else { continue };
                    let use_counts = with_counts && layer_dir.join("counts").exists();
                    let mat = if use_counts {
                        PersistentCompactIntMatrix::open(&layer_dir).ok().map(Mat::Count)
                    } else {
                        PersistentBitMatrix::open(&layer_dir).ok().map(Mat::Presence)
                    };
                    let Some(mat) = mat else { continue };
                    layers.push(mphf);
                    mats.push(mat);
                }
                pb.inc(1);
                Ok((layers, mats))
            })
            .collect::<OKIResult<Vec<_>>>()?;
        pb.finish_and_clear();
        let (layers, mats) = built.into_iter().unzip();
        Ok(Self { layers, mats })
    }

    /// Existence-only lookup of `variant` in partition `dest_partition`:
    /// tries each of the partition's already-open layers in turn, stopping
    /// at the first hit.
    fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> bool {
        self.layers
            .get(dest_partition)
            .is_some_and(|layers| layers.iter().any(|mphf| mphf.find(variant).is_some()))
    }

    /// Per-genome presence vector for `variant` in partition `dest_partition`
    /// (`true` iff that genome carries it), `None` on a miss. Same shape as
    /// `find`, but also reads the cached matrix instead of just the MPHF.
    fn find_presence(&self, dest_partition: usize, variant: CanonicalKmer, n_genomes: usize) -> Option<Vec<bool>> {
        let layers = self.layers.get(dest_partition)?;
        let mats = self.mats.get(dest_partition)?;
        for (mphf, mat) in layers.iter().zip(mats.iter()) {
            if let Some(slot) = mphf.find(variant) {
                let n_cols = mat.n_cols().min(n_genomes);
                return Some((0..n_cols).map(|g| mat.carries(g, slot)).collect());
            }
        }
        None
    }
}

impl KmerIndex {
    /// Build the sibling-count/minorant annex for every layer of every
    /// partition of this (already built) index, writing one annex file per
    /// layer alongside its existing index files. Safe to call again later
    /// (e.g. after a fresh `merge`) — each run simply overwrites the annex
    /// files of the index it is called on.
    ///
    /// Construction only — no statistics gathered here on purpose: this is
    /// meant to run routinely (it is the artefact the SNP-family distances
    /// will consume), while the sibling-count distribution
    /// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate,
    /// occasional diagnostic pass over the result, not run every time.
    ///
    /// Cross-partition/cross-layer lookups are required (a k-mer's siblings
    /// can live in any partition), but the layer loop itself — and thus the
    /// annex file this produces — stays local to one layer at a time.
    pub fn build_sibling_annex(&self) -> OKIResult<()> {
        let n_parts = self.n_partitions();
        let n_bits = n_parts.trailing_zeros() as usize;

        let partition = KmerPartition::open_with_config(
            &self.root_path,
            self.kmer_size(),
            self.minimizer_size(),
            n_bits,
        )
        .map_err(OKIError::Partition)?;

        tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
        let cache = Arc::new(PartitionCache::build(&partition, n_parts, self.meta.config.with_counts)?);

        let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
        let mut total_slots: u64 = 0;
        for part in 0..n_parts {
            let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
            if !index_dir.exists() {
                pb.inc(1);
                continue;
            }
            let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;

            let mut part_slots: u64 = 0;
            for l in 0..meta.n_layers {
                let layer_dir = index_dir.join(format!("layer_{l}"));
                part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?;
            }
            total_slots += part_slots;
            pb.inc(1);
            pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)"));
        }
        pb.finish_and_clear();
        tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions");

        Ok(())
    }

    /// Returns the number of distinct k-mers (annex slots) processed, for
    /// progress reporting.
    fn build_layer_sibling_annex(
        &self,
        layer_dir: &Path,
        n_parts: usize,
        cache: &Arc<PartitionCache>,
    ) -> OKIResult<u64> {
        let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
        let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
        let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
        let n_slots = mphf.n();

        // ── Enumerate this layer's distinct k-mers, one per slot ────────────
        let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; n_slots];
        let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
            .map_err(OKIError::Partition)?;
        for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
            if let Some(slot) = mphf.find(kmer) {
                slot_kmer[slot] = Some(kmer);
            }
        }

        let k = self.kmer_size();

        // ── Reconciliation state, initialised with each slot's own base —
        // that member is trivially present, no lookup needed. Built before
        // the pipeline runs, from the same enumeration, since `sources`
        // below is consumed as a throttled iterator, not collected.
        // `AtomicU8`, not `FamilyMask`, because the gather phase below
        // parallelises across destination partitions (independent
        // `query_partition_with` calls, safe to run concurrently) and their
        // `Found` hits can land on arbitrary, possibly-shared slots — a
        // lock-free `fetch_or` avoids needing any synchronisation beyond
        // that. ─────────────────────────────────────────────────────────
        let mask: Vec<AtomicU8> = (0..n_slots).map(|_| AtomicU8::new(0)).collect();
        for (slot, kmer) in slot_kmer.iter().enumerate().filter_map(|(s, k)| k.map(|k| (s, k))) {
            mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
        }

        // ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one —
        // the actual cross-partition lookup reuses
        // `KmerPartition::query_partition_with` (the same batching mechanism
        // `obikmer query` already uses: open a partition's files once,
        // answer a whole batch of queries against it) instead of one lookup
        // per pipeline item. A per-item lookup (tried first) reopened/
        // re-mmap'd every target partition's files on every single variant
        // — fine at toy scale, but ~90% system time against a real index,
        // observed in practice. A *later* attempt still pushed one pipeline
        // message per generated variant (a `Flat` stage, `SourceItem` =>
        // `VariantQuery`, one k-mer in => up to 3 variants out as separate
        // messages) — cheaper than reopening files, but sampling a real run
        // showed most wall-clock time going into per-message channel
        // send/notify syscalls instead of the lookup itself: the pipeline's
        // whole point is amortising synchronisation over a batch, and a
        // single k-mer's ≤3 variants is far too fine a granularity for
        // that. Batching `BATCH_SIZE` source k-mers into one pipeline item
        // — a plain 1-to-1 (`|`, not `||`) transform, batch in, batch of
        // variants out, one message either way — keeps the per-message
        // synchronisation cost amortised over thousands of lookups instead
        // of one to three. ──────────────────────────────────────────────
        const BATCH_SIZE: usize = 4096;
        let n_workers = obisys::effective_parallelism();
        let capacity = 256;

        // Throttling limits how many *batches* are in flight at once — the
        // permit is acquired per batch (not per k-mer) in the source
        // thread, and released once its `VariantBatch` has been read out of
        // the pipeline by the accumulation loop below. See
        // `obipipeline::throttle`'s docs for why this is required, not
        // optional, once a `Flat`-style stage sits in the pipeline.
        let sources: Vec<(usize, CanonicalKmer)> = slot_kmer
            .iter()
            .enumerate()
            .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer)))
            .collect();
        let batches: Vec<Vec<(usize, CanonicalKmer)>> = sources
            .chunks(BATCH_SIZE)
            .map(|chunk| chunk.to_vec())
            .collect();
        let throttled = obipipeline::throttle(batches.into_iter(), n_workers).map(|t| SourceBatch {
            items: t.item,
            _permit: t.guard,
        });

        let pipe = obipipeline::make_pipe! {
            SibData : SourceBatch => VariantBatch,
            | {
                move |batch: SourceBatch| -> VariantBatch {
                    let mut items = Vec::with_capacity(batch.items.len() * 3);
                    for (slot, kmer) in batch.items {
                        for variant in kmer.central_canonical_neighbors() {
                            if variant == kmer {
                                continue;
                            }
                            items.push((
                                partition_of(variant, n_parts),
                                variant,
                                slot,
                                central_base(variant, k),
                            ));
                        }
                    }
                    VariantBatch { items, _permit: batch._permit }
                }
            } : Batch => Variants,
        };

        // ── Group generated variants by destination partition. `cache`
        // holds every partition already mmap'd (no more `open()` cost), but
        // `mmap` pages are still loaded on demand and can be evicted — a
        // lookup is not free just because the file isn't reopened. Grouping
        // keeps one partition's pages hot while its whole batch is resolved,
        // instead of faulting pages in and out as lookups jump between
        // partitions in whatever order the pipeline happens to produce
        // them. Each batch's throttle permit drops here, once accumulated.
        let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
        for vb in pipe.apply(throttled, n_workers, capacity) {
            for (dest_partition, variant, source_slot, base) in vb.items {
                outgoing[dest_partition].push((variant, source_slot, base));
            }
        }

        // ── Resolve each partition's batch against the cache in one
        // contiguous pass; parallelised across partitions (independent,
        // read-only) so this keeps using multiple cores without giving up
        // the per-partition locality above. ─────────────────────────────
        outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
            for &(variant, source_slot, base) in queries {
                if cache.find(dest, variant) {
                    mask[source_slot].fetch_or(1 << base, Ordering::Relaxed);
                }
            }
        });

        // ── Write the layer's annex file ─────────────────────────────────────
        let annex_path = layer_dir.join(ANNEX_FILE_NAME);
        let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?;
        for (slot, m) in mask.iter().enumerate() {
            if slot_kmer[slot].is_none() {
                continue; // unused MPHF slot, if any — leave at the sentinel
            }
            builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed)));
        }
        builder.close()?;

        Ok(n_slots as u64)
    }
}

/// Distribution of family sizes (1-4), read back from an already-built
/// annex (see [`KmerIndex::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]>,
}

impl KmerIndex {
    /// 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`] has not been run on
    /// this index first.
    ///
    /// [`build_sibling_annex`]: Self::build_sibling_annex
    pub fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats> {
        let n_parts = self.n_partitions();
        let n_genomes = self.meta.genomes.len();
        let with_counts = self.meta.config.with_counts;
        let k = self.kmer_size();
        let n_bits = n_parts.trailing_zeros() as usize;

        // 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 partition = KmerPartition::open_with_config(
            &self.root_path,
            self.kmer_size(),
            self.minimizer_size(),
            n_bits,
        )
        .map_err(OKIError::Partition)?;
        let cache = PartitionCache::build(&partition, n_parts, with_counts)?;

        // Gather the (partition, layer) pairs to process — cheap metadata
        // reads only, checking every annex file exists up front so a
        // missing one is reported before any real work starts.
        let mut layer_dirs = Vec::new();
        for part in 0..n_parts {
            let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
            if !index_dir.exists() {
                continue;
            }
            let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
            for l in 0..meta.n_layers {
                let layer_dir = index_dir.join(format!("layer_{l}"));
                let annex_path = layer_dir.join(ANNEX_FILE_NAME);
                if !annex_path.exists() {
                    return Err(OKIError::InvalidInput(format!(
                        "no sibling annex at {} — run build_sibling_annex first",
                        annex_path.display()
                    )));
                }
                layer_dirs.push(layer_dir);
            }
        }

        // One layer's worth of work, parallelised across layers with Rayon
        // — independent, read-only, each producing its own partial tally
        // merged at the end.
        let pb = progress_bar("sibling_annex_stats", layer_dirs.len() as u64, "layers");
        let partials: Vec<SiblingAnnexStats> = layer_dirs
            .par_iter()
            .map(|layer_dir| -> OKIResult<SiblingAnnexStats> {
                let mut stats = SiblingAnnexStats {
                    per_genome: vec![[0u64; 4]; n_genomes],
                    ..Default::default()
                };

                let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
                let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
                let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;

                // Need each slot's own k-mer to derive minorant — same
                // enumeration as construction.
                let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
                let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
                let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
                    .map_err(OKIError::Partition)?;
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    if let Some(slot) = mphf.find(kmer) {
                        slot_kmer[slot] = Some(kmer);
                    }
                }

                let use_counts = with_counts && layer_dir.join("counts").exists();
                let mat = if use_counts {
                    Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
                } else {
                    Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
                };
                let n_cols = mat.n_cols().min(n_genomes);

                for slot in 0..annex.len() {
                    let Some(mask) = annex.get(slot) else { continue };
                    let Some(kmer) = slot_kmer[slot] else { continue };
                    if !is_minorant(kmer, mask, k) {
                        continue; // this family is tallied at its minorant's slot only
                    }
                    let s = mask.siblings() as usize;
                    stats.counts[s] += 1;

                    // "Genome g represents this family" means g carries
                    // *any* of its members, not just the minorant's own —
                    // start from the minorant's own presence (already
                    // open, no lookup) and OR in every other present
                    // member's presence vector, resolved against the
                    // whole-run cache (no I/O) — exactly `mask.siblings()`
                    // of them, the mask tells us precisely which to fetch.
                    let mut carries = vec![false; n_cols];
                    for g in 0..n_cols {
                        carries[g] = mat.carries(g, slot);
                    }
                    for other in kmer.central_canonical_neighbors() {
                        if other == kmer {
                            continue;
                        }
                        let base = central_base(other, k);
                        if !mask.has(base) {
                            continue;
                        }
                        let dest = partition_of(other, n_parts);
                        if let Some(other_presence) = cache.find_presence(dest, other, n_genomes) {
                            for (g, &present) in other_presence.iter().enumerate() {
                                if present {
                                    carries[g] = true;
                                }
                            }
                        }
                    }

                    for (g, &carried) in carries.iter().enumerate() {
                        if carried {
                            stats.per_genome[g][s] += 1;
                        }
                    }
                }

                pb.inc(1);
                Ok(stats)
            })
            .collect::<OKIResult<Vec<_>>>()?;
        pb.finish_and_clear();

        let mut stats = SiblingAnnexStats {
            per_genome: vec![[0u64; 4]; n_genomes],
            ..Default::default()
        };
        for part in partials {
            for s in 0..4 {
                stats.counts[s] += part.counts[s];
            }
            for g in 0..n_genomes {
                for s in 0..4 {
                    stats.per_genome[g][s] += part.per_genome[g][s];
                }
            }
        }
        Ok(stats)
    }
}

/// Raw p-distance restricted to loci that are single-copy in **both**
/// genomes of a pair — the "stringent / paralogy-aware" locus eligibility
/// rule (`docmd/theory/evolutionary_distances.md`, "Locus eligibility"),
/// without the JC/K2P/LogDet finalisation math: `snp[i,j]` + `shared[i,j]`
/// is the number of eligible loci, `snp[i,j] / (snp[i,j] + shared[i,j])` is
/// `p_hat`. A quick, self-contained way to sanity-check the estimator
/// against a real index before the full `SnpTally` design is built.
///
/// A locus (family, tallied once at its minorant) is eligible for pair
/// `(i, j)` iff genome `i` carries exactly one of the family's observed
/// forms **and** genome `j` carries exactly one (possibly a different one)
/// — presence-only: a genome carrying the same form twice (a same-allele
/// duplicate) is indistinguishable from carrying it once when only a
/// presence matrix is available, so such cases are not excluded here even
/// when a count index exists. See "Locus eligibility", stringent rule, for
/// why this matters and how a count index would close the gap — left as a
/// follow-up, not applied here.
pub struct RawSnpDistanceOutput {
    /// n×n count of eligible loci where the two genomes' single forms differ.
    pub snp: Array2<u64>,
    /// n×n count of eligible loci where the two genomes' single forms agree.
    pub shared: Array2<u64>,
}

impl KmerIndex {
    /// Shared traversal behind [`raw_snp_distance`](Self::raw_snp_distance)
    /// and [`base_pair_tally`](Self::base_pair_tally): for every family
    /// (tallied once, at its minorant) of every layer of the already-built
    /// sibling annex, resolves each genome's single observed form (`None`
    /// if absent or ambiguous/multi-copy), then calls `on_pair(acc, i, j,
    /// bi, bj)` for every genome pair `(i, j)` where both are unambiguous
    /// and single-copy (`bi == bj` means shared at that locus, `bi != bj`
    /// means a SNP). Layers are processed in parallel (rayon); each gets
    /// its own accumulator from `zero()`, combined pairwise via `combine`.
    fn scan_family_pairs<Acc, F, C>(
        &self,
        label: &str,
        zero: impl Fn() -> Acc + Sync,
        on_pair: F,
        combine: C,
    ) -> OKIResult<Acc>
    where
        Acc: Send,
        F: Fn(&mut Acc, usize, usize, u8, u8) + Sync,
        C: Fn(Acc, Acc) -> Acc,
    {
        let n_parts = self.n_partitions();
        let n_genomes = self.meta.genomes.len();
        let with_counts = self.meta.config.with_counts;
        let k = self.kmer_size();
        let n_bits = n_parts.trailing_zeros() as usize;

        let partition = KmerPartition::open_with_config(
            &self.root_path,
            self.kmer_size(),
            self.minimizer_size(),
            n_bits,
        )
        .map_err(OKIError::Partition)?;
        let cache = PartitionCache::build(&partition, n_parts, with_counts)?;

        let mut layer_dirs = Vec::new();
        for part in 0..n_parts {
            let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
            if !index_dir.exists() {
                continue;
            }
            let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
            for l in 0..meta.n_layers {
                let layer_dir = index_dir.join(format!("layer_{l}"));
                let annex_path = layer_dir.join(ANNEX_FILE_NAME);
                if !annex_path.exists() {
                    return Err(OKIError::InvalidInput(format!(
                        "no sibling annex at {} — run build_sibling_annex first",
                        annex_path.display()
                    )));
                }
                layer_dirs.push(layer_dir);
            }
        }

        let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
        let partials: Vec<Acc> = layer_dirs
            .par_iter()
            .map(|layer_dir| -> OKIResult<Acc> {
                let mut acc = zero();

                let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
                let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
                let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
                let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;

                let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
                let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
                    .map_err(OKIError::Partition)?;
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    if let Some(slot) = mphf.find(kmer) {
                        slot_kmer[slot] = Some(kmer);
                    }
                }

                let use_counts = with_counts && layer_dir.join("counts").exists();
                let mat = if use_counts {
                    Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
                } else {
                    Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
                };
                let n_cols = mat.n_cols().min(n_genomes);

                // Per family: which single form (if exactly one) each genome
                // carries — `None` once a second form is seen (ambiguous,
                // not single-copy, ineligible for either side of a pair).
                let mut single_form: Vec<Option<u8>> = Vec::with_capacity(n_cols);
                let mut ambiguous: Vec<bool> = Vec::with_capacity(n_cols);

                for slot in 0..annex.len() {
                    let Some(mask) = annex.get(slot) else { continue };
                    let Some(kmer) = slot_kmer[slot] else { continue };
                    if !is_minorant(kmer, mask, k) {
                        continue; // family tallied once, at its minorant
                    }

                    single_form.clear();
                    single_form.resize(n_cols, None);
                    ambiguous.clear();
                    ambiguous.resize(n_cols, false);

                    for other in kmer.central_canonical_neighbors() {
                        let base = central_base(other, k);
                        if !mask.has(base) {
                            continue;
                        }
                        let presence: Option<Vec<bool>> = if other == kmer {
                            Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
                        } else {
                            let dest = partition_of(other, n_parts);
                            cache.find_presence(dest, other, n_genomes)
                        };
                        let Some(presence) = presence else { continue };
                        for (g, &present) in presence.iter().enumerate() {
                            if !present {
                                continue;
                            }
                            if single_form[g].is_some() {
                                ambiguous[g] = true;
                            } else {
                                single_form[g] = Some(base);
                            }
                        }
                    }

                    for i in 0..n_cols {
                        if ambiguous[i] {
                            continue;
                        }
                        let Some(bi) = single_form[i] else { continue };
                        for j in (i + 1)..n_cols {
                            if ambiguous[j] {
                                continue;
                            }
                            let Some(bj) = single_form[j] else { continue };
                            on_pair(&mut acc, i, j, bi, bj);
                        }
                    }
                }

                pb.inc(1);
                Ok(acc)
            })
            .collect::<OKIResult<Vec<_>>>()?;
        pb.finish_and_clear();

        let mut total = zero();
        for partial in partials {
            total = combine(total, partial);
        }
        Ok(total)
    }

    /// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex
    /// (run [`build_sibling_annex`](Self::build_sibling_annex) first).
    pub fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> {
        let n_genomes = self.meta.genomes.len();
        let (snp, shared) = self.scan_family_pairs(
            "raw_snp_distance",
            || (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::zeros((n_genomes, n_genomes))),
            |(snp, shared), i, j, bi, bj| {
                if bi == bj {
                    shared[[i, j]] += 1;
                    shared[[j, i]] += 1;
                } else {
                    snp[[i, j]] += 1;
                    snp[[j, i]] += 1;
                }
            },
            |(mut snp, mut shared), (s, sh)| {
                snp += &s;
                shared += &sh;
                (snp, shared)
            },
        )?;
        Ok(RawSnpDistanceOutput { snp, shared })
    }

    /// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG,
    /// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs
    /// whose overall SNP ratio in `raw` is at or below `ratio_ceiling` —
    /// same saturation-exclusion discipline as `calibrate_p_hat`, for the
    /// same reason: a saturated pair's observed base-pair mix trends toward
    /// neutral base composition, not the true point-mutation spectrum.
    ///
    /// A second full pass over the annex, sharing
    /// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by
    /// it, not a blind re-scan) — needed because `raw_snp_distance` only
    /// keeps aggregate SNP/shared counts per genome pair, not which bases
    /// were actually involved at each locus, and the ratio-ceiling filter
    /// can only be evaluated once the aggregate counts are known.
    pub fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally> {
        let n_genomes = self.meta.genomes.len();
        let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
            if i == j {
                return false;
            }
            let snp = raw.snp[[i, j]];
            let total = snp + raw.shared[[i, j]];
            total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
        });

        let counts = self.scan_family_pairs(
            "base_pair_tally",
            || [[0u64; 4]; 4],
            |counts, i, j, bi, bj| {
                if bi != bj && included[[i, j]] {
                    counts[bi as usize][bj as usize] += 1;
                    counts[bj as usize][bi as usize] += 1;
                }
            },
            |mut total, partial| {
                for a in 0..4 {
                    for b in 0..4 {
                        total[a][b] += partial[a][b];
                    }
                }
                total
            },
        )?;
        Ok(BasePairTally { counts })
    }
}

/// See [`KmerIndex::base_pair_tally`].
pub struct BasePairTally {
    /// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled
    /// over included genome pairs, where the two genomes' single forms are
    /// `a` and `b` (0=A, 1=C, 2=G, 3=T). Diagonal always `0` (an `a == b`
    /// locus is `shared`, not tallied here).
    pub counts: [[u64; 4]; 4],
}

/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set
/// iff the genome carries the member whose own central base is `b`):
/// single bit -> the plain base; 2 or 3 bits -> the matching IUPAC
/// ambiguity code (preserves partial information instead of collapsing to
/// `N`, the same convention used for diploid heterozygous VCF/FASTA sites);
/// all 4 bits -> `N`; no bits (genome carries none of the family's observed
/// members) -> `-` (no data at this locus for this genome).
fn iupac_code(mask: u8) -> u8 {
    match mask & 0b1111 {
        0b0000 => b'-',
        0b0001 => b'A',
        0b0010 => b'C',
        0b0100 => b'G',
        0b1000 => b'T',
        0b0101 => b'R', // A/G
        0b1010 => b'Y', // C/T
        0b0110 => b'S', // C/G
        0b1001 => b'W', // A/T
        0b1100 => b'K', // G/T
        0b0011 => b'M', // A/C
        0b1110 => b'B', // C/G/T
        0b1101 => b'D', // A/G/T
        0b1011 => b'H', // A/C/T
        0b0111 => b'V', // A/C/G
        0b1111 => b'N',
        _ => unreachable!("masked to 4 bits"),
    }
}

/// A SNP-only pseudo-alignment: one row (byte sequence, IUPAC-coded) per
/// genome, one column per variable family (`family_size() >= 2` — monomorphic
/// families carry no signal and are skipped, unlike `raw_snp_distance`'s
/// tally which does count them as `shared`). Column order is the same,
/// deterministic sweep order as the annex build (partition, then layer, then
/// slot) — arbitrary but stable and identical across genomes, which is all a
/// pseudo-alignment needs (there is no natural genomic coordinate to sort by
/// once flanks are dropped). See `docmd/theory/evolutionary_distances.md`,
/// "Multi-genome framing: family as pseudo-alignment column".
pub struct SnpAlignment {
    /// `sequences[g]` = genome `g`'s IUPAC-coded row, same length for every
    /// genome (`sequences.len()` columns).
    pub sequences: Vec<Vec<u8>>,
}

impl KmerIndex {
    /// Build the SNP-only pseudo-alignment from an already-built sibling
    /// annex (run [`build_sibling_annex`](Self::build_sibling_annex) first).
    pub fn snp_pseudo_alignment(&self) -> OKIResult<SnpAlignment> {
        let n_parts = self.n_partitions();
        let n_genomes = self.meta.genomes.len();
        let with_counts = self.meta.config.with_counts;
        let k = self.kmer_size();
        let n_bits = n_parts.trailing_zeros() as usize;

        let partition = KmerPartition::open_with_config(
            &self.root_path,
            self.kmer_size(),
            self.minimizer_size(),
            n_bits,
        )
        .map_err(OKIError::Partition)?;
        let cache = PartitionCache::build(&partition, n_parts, with_counts)?;

        let mut layer_dirs = Vec::new();
        for part in 0..n_parts {
            let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
            if !index_dir.exists() {
                continue;
            }
            let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
            for l in 0..meta.n_layers {
                let layer_dir = index_dir.join(format!("layer_{l}"));
                let annex_path = layer_dir.join(ANNEX_FILE_NAME);
                if !annex_path.exists() {
                    return Err(OKIError::InvalidInput(format!(
                        "no sibling annex at {} — run build_sibling_annex first",
                        annex_path.display()
                    )));
                }
                layer_dirs.push(layer_dir);
            }
        }

        let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers");
        // `Vec<Vec<u8>>` per layer, one entry (column) per variable family;
        // `par_iter().map(...).collect()` on this indexed source preserves
        // input order, so concatenating the results below in order gives a
        // single deterministic column order across the whole index.
        let partials: Vec<Vec<Vec<u8>>> = layer_dirs
            .par_iter()
            .map(|layer_dir| -> OKIResult<Vec<Vec<u8>>> {
                let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
                let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
                let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
                let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;

                let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
                let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
                    .map_err(OKIError::Partition)?;
                for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
                    if let Some(slot) = mphf.find(kmer) {
                        slot_kmer[slot] = Some(kmer);
                    }
                }

                let use_counts = with_counts && layer_dir.join("counts").exists();
                let mat = if use_counts {
                    Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
                } else {
                    Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
                };
                let n_cols = mat.n_cols().min(n_genomes);

                let mut columns: Vec<Vec<u8>> = Vec::new();
                let mut genome_mask: Vec<u8> = Vec::with_capacity(n_genomes);

                for slot in 0..annex.len() {
                    let Some(mask) = annex.get(slot) else { continue };
                    let Some(kmer) = slot_kmer[slot] else { continue };
                    if !is_minorant(kmer, mask, k) {
                        continue; // family tallied once, at its minorant
                    }
                    if mask.family_size() < 2 {
                        continue; // monomorphic family — no signal, skip
                    }

                    genome_mask.clear();
                    genome_mask.resize(n_genomes, 0);

                    for other in kmer.central_canonical_neighbors() {
                        let base = central_base(other, k);
                        if !mask.has(base) {
                            continue;
                        }
                        let presence: Option<Vec<bool>> = if other == kmer {
                            Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
                        } else {
                            let dest = partition_of(other, n_parts);
                            cache.find_presence(dest, other, n_genomes)
                        };
                        let Some(presence) = presence else { continue };
                        for (g, &present) in presence.iter().enumerate() {
                            if present {
                                genome_mask[g] |= 1 << base;
                            }
                        }
                    }

                    columns.push(genome_mask.iter().map(|&m| iupac_code(m)).collect());
                }

                pb.inc(1);
                Ok(columns)
            })
            .collect::<OKIResult<Vec<_>>>()?;
        pb.finish_and_clear();

        let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); n_genomes];
        for layer_columns in partials {
            for column in layer_columns {
                for (g, &code) in column.iter().enumerate() {
                    sequences[g].push(code);
                }
            }
        }
        Ok(SnpAlignment { sequences })
    }
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use std::path::Path;

    use obicompactvec::SiblingAnnex;
    use obikseq::{Kmer, Sequence};
    use obisys::Reporter;
    use tempfile::tempdir;

    use crate::meta::{GenomeInfo, IndexConfig};
    use crate::merge::MergeMode;

    use super::*;

    // k must be >= 11 (project constraint, "k ∈ [11,31]"); k=11, level_max=1,
    // theta=0.0 mirror `obiskbuilder`'s own tests (smaller k/level_max
    // combinations trip an unrelated pre-existing bug in `obikentropy`'s
    // sliding-window ring buffer — not this feature's concern).
    const K: usize = 11;
    const M: usize = 5;

    /// Build a single-genome index from one in-memory FASTA sequence, driving
    /// the same primitives `obikmer`'s `scatter` step uses (minus the
    /// multi-file `obipipeline` wrapper — a single sequence needs none of
    /// that): normalise -> build superkmers -> route -> write.
    /// `cargo test` doesn't install a `tracing` subscriber the way `obikmer`'s
    /// CLI does, so `debug!`/etc. are silent no-ops by default — including the
    /// `PartitionRunner` instrumentation that would matter most for
    /// re-diagnosing a hang here. `try_init` is idempotent across concurrently
    /// running tests (later calls just find a subscriber already installed).
    fn init_tracing() {
        let _ = tracing_subscriber::fmt()
            .with_env_filter(
                tracing_subscriber::EnvFilter::try_from_default_env()
                    .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
            )
            .with_writer(std::io::stderr)
            .try_init();
    }

    fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
        init_tracing();
        let fasta_path = dir.join(format!("{label}.fasta"));
        let mut f = std::fs::File::create(&fasta_path).unwrap();
        writeln!(f, ">{label}").unwrap();
        f.write_all(seq).unwrap();
        writeln!(f).unwrap();
        drop(f);

        let index_path = dir.join(format!("{label}.idx"));
        let config = IndexConfig {
            kmer_size: K,
            minimizer_size: M,
            n_bits: 0, // 1 partition — keeps the test deterministic and simple
            with_counts: false,
            evidence: obilayeredmap::IndexMode::Exact,
            block_bits: 0,
        };
        let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)), false)
            .expect("create");

        let mut rep = Reporter::new();
        let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
        for page in stream {
            let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
            idx.partition_mut().write_batch(batch).expect("write_batch");
        }
        idx.partition_mut().close().expect("close partition writers");
        idx.mark_scattered().expect("mark_scattered");
        idx.dereplicate_and_count(false, &mut rep).expect("dereplicate_and_count");
        idx.build_layers(1, None, false, &mut rep).expect("build_layers");
        idx
    }

    fn canonical(ascii: &[u8]) -> CanonicalKmer {
        Kmer::from_ascii(ascii).unwrap().canonical()
    }

    /// Read back the annex entry for a given canonical k-mer from the merged
    /// index's (single) partition/layer, asserting it was found at all.
    fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> FamilyMask {
        let index_dir = idx.partition().part_dir(0).join(INDEX_SUBDIR);
        let meta = PartitionMeta::load(&index_dir).unwrap();
        for l in 0..meta.n_layers {
            let layer_dir = index_dir.join(format!("layer_{l}"));
            let mphf = MphfLayer::open(&layer_dir, &meta.mode).unwrap();
            if let Some(slot) = mphf.find(kmer) {
                let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).unwrap();
                return annex.get(slot).expect("slot must have a computed annex entry");
            }
        }
        panic!("kmer not found in any layer of partition 0");
    }

    fn merge_two(dir: &Path, g1: &KmerIndex, g2: &KmerIndex) -> KmerIndex {
        let mut rep = Reporter::new();
        KmerIndex::merge(
            &dir.join("merged.idx"),
            &[g1, g2],
            MergeMode::Presence,
            false,
            false,
            1.0,
            &mut rep,
        )
        .expect("merge")
    }

    #[test]
    fn sibling_annex_one_sibling_each() {
        // k=11, centre = index 5 (0-based). Two genomes, each exactly one
        // k-mer, sharing every base except the centre:
        //   g1 = "AACCGCTTAAG" (centre 'C', base index 1)
        //   g2 = "AACCGGTTAAG" (centre 'G', base index 2)
        // Hand-verified: both stay forward-oriented under canonicalisation
        // (each is lexicographically smaller than its own reverse
        // complement, since both start with "AA"), and raw(g1) < raw(g2)
        // (only differing base: C=0b01 < G=0b10 at the centre) — so g1 is
        // the minorant, g2 is not. The mask is a family-wide value: both
        // slots must read back the *same* mask (bits 1 and 2 set).
        let dir = tempdir().unwrap();
        let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
        let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
        let merged = merge_two(dir.path(), &g1, &g2);
        merged.build_sibling_annex().expect("build_sibling_annex");

        let g1_kmer = canonical(b"AACCGCTTAAG");
        let g2_kmer = canonical(b"AACCGGTTAAG");
        let expected_mask = FamilyMask::EMPTY.with(1).with(2);

        let a = annex_info_for(&merged, g1_kmer);
        assert_eq!(a, expected_mask, "AACCGCTTAAG");
        assert_eq!(a.siblings(), 1);
        assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant");

        let b = annex_info_for(&merged, g2_kmer);
        assert_eq!(b, expected_mask, "AACCGGTTAAG");
        assert_eq!(b.siblings(), 1);
        assert!(!is_minorant(g2_kmer, b, K), "g2 should not be the minorant");
    }

    #[test]
    fn sibling_annex_zero_siblings_when_identical_and_no_variant_exists() {
        // Same k-mer in both genomes, no other genome around to carry a
        // variant -> 0 siblings, trivially its own minorant.
        let dir = tempdir().unwrap();
        let g1 = build_single_genome_index(dir.path(), "g1", b"GATTACAGATC");
        let g2 = build_single_genome_index(dir.path(), "g2", b"GATTACAGATC");
        let merged = merge_two(dir.path(), &g1, &g2);
        merged.build_sibling_annex().expect("build_sibling_annex");

        let kmer = canonical(b"GATTACAGATC");
        let mask = annex_info_for(&merged, kmer);
        assert_eq!(mask.siblings(), 0, "GATTACAGATC");
        assert_eq!(mask.family_size(), 1);
        assert!(is_minorant(kmer, mask, K));
    }

    #[test]
    fn sibling_annex_stats_counts_each_family_once_and_per_genome() {
        // Reuses the one-sibling-each fixture: a single family of size 2
        // (g1's centre-C form + g2's centre-G form), each genome carrying
        // exactly one of the two members. Stats must report exactly one
        // family of size 2 (`counts[1] == 1`, since index 1 = size 2), not
        // two (which naively summing both slots would give), and both
        // genomes represented at size 2, neither at any other size.
        let dir = tempdir().unwrap();
        let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
        let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
        let merged = merge_two(dir.path(), &g1, &g2);
        merged.build_sibling_annex().expect("build_sibling_annex");

        let stats = merged.sibling_annex_stats().expect("sibling_annex_stats");

        assert_eq!(stats.counts, [0, 1, 0, 0], "one family of size 2, counted once");
        assert_eq!(stats.per_genome.len(), 2);
        for g in 0..2 {
            assert_eq!(
                stats.per_genome[g], [0, 1, 0, 0],
                "genome {g} should represent exactly one size-2 family"
            );
        }
    }
}
