obikmer/src/obikphylo/src/siblings/tests.rs

use std::io::Write;
use std::path::Path;

use obikseq::{CanonicalKmer, Kmer, Sequence};
use obikindex::layer::MphfLayer;
use obisys::Reporter;
use tempfile::tempdir;

use obikindexer::algorithms::counter::Counter;
use obikindexer::algorithms::dereplicator::Dereplicator;
use obikindexer::algorithms::layer_builder::LayerBuilder;
use obikindexer::algorithms::partitionner::PartitionRouter;
use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode};

use super::alignment::SnpAlignmentExt;
use super::build::SiblingAnnexBuildExt;
use super::cardinality::CardinalityExt;
use super::distance::DistanceExt;
use super::entropy::ShannonEntropyExt;
use super::entropy_annex::{EntropyAnnex, ENTROPY_ANNEX_FILE_NAME};
use super::helpers::is_minorant;
use super::sankoff_bundle::SankoffBundleExt;
use super::stats::SiblingStatsExt;
use super::subsample::EntropyBias;
use super::{FamilyMask, SiblingAnnex, ANNEX_FILE_NAME};

// 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: obikindex::layer::IndexMode::Exact,
        block_bits: 0,
    };
    let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)))
        .expect("create");

    let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
    let mut router = PartitionRouter::new(&mut idx);
    for page in stream {
        let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
        router.write_batch(batch).expect("write_batch");
    }
    router.close().expect("close partition writers"); // also marks scatter done
    drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope

    Dereplicator::new(&idx).run(None::<fn(obisys::Progress)>).expect("dereplicate");
    Counter::new(&idx).run(None::<fn(obisys::Progress)>).expect("count_kmer"); // also writes the spectrum + marks counted
    LayerBuilder::new(&idx)
        .run(None::<fn(obisys::Progress)>)
        .expect("build_layers"); // also marks indexed
    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 meta = idx.partition_meta(0).unwrap();
    for l in 0..meta.n_layers {
        let layer_dir = idx.layer_dir(0, 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* presence bits (1 and 2 set) — but
    // only g1's slot should carry the stored minorant flag.
    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.bits(), expected_mask.bits(), "AACCGCTTAAG");
    assert_eq!(a.siblings(), 1);
    assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant");
    assert!(a.is_minorant(), "g1's stored minorant flag should be set at build time");

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

#[test]
fn sibling_annex_works_after_pack_sparse() {
    // Same fixture as `sibling_annex_one_sibling_each`, but with
    // `pack --sparse` (`KmerIndex::pack_matrices(true)`) run on the merged
    // index before building the sibling annex — proves `PartitionCache`'s
    // sparse-detection (`Mat::SparsePresence`, gated on `presence/
    // is_multi.prsb`) and the generic `TypedLayer<D>` methods it relies on
    // actually round-trip through the real build pipeline, not just the
    // unit-level `TypedLayer<PersistentSparseBitMatrix>` tests in
    // `obikindex::layer`.
    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.pack_matrices(true).expect("pack_matrices(sparse)");

    let index_dir = merged.index_dir(0);
    assert!(
        index_dir.join("layer_0").join("presence").join("is_multi.prsb").exists(),
        "pack_matrices(true) must leave the sparse marker file behind"
    );

    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.bits(), expected_mask.bits(), "AACCGCTTAAG");
    assert_eq!(a.siblings(), 1);
    assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant");
    assert!(a.is_minorant(), "g1's stored minorant flag should be set at build time");

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

#[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"
        );
    }
}

#[test]
fn sibling_family_size_histogram_matches_full_stats_global_counts() {
    // Same fixture as `sibling_annex_stats_counts_each_family_once_and_per_genome`
    // — the cheap, annex-only histogram must agree with the `counts` half of
    // the full (cross-partition) stats pass, without needing a `PartitionCache`
    // at all.
    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 histogram = merged.sibling_family_size_histogram().expect("sibling_family_size_histogram");
    assert_eq!(histogram, [0, 1, 0, 0], "one family of size 2, counted once");

    let stats = merged.sibling_annex_stats().expect("sibling_annex_stats");
    assert_eq!(histogram, stats.counts, "must agree with the full stats pass's global counts");
}

/// Exercises the four sibling-annex consumers that were rewritten to share
/// `family_scan::scan_layer_families` (partition-grouped lookups instead of
/// one lookup per family) — same one-sibling-each fixture as the tests
/// above (g1 carries the family's centre-C member, g2 the centre-G one),
/// hand-verified expected output for each.
#[test]
fn family_scan_consumers_agree_on_one_sibling_each() {
    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");

    // Merge doesn't promise to preserve source order, so resolve each
    // genome's index by label rather than assuming g1 -> 0, g2 -> 1.
    let idx_of = |label: &str| merged.meta().genomes().unwrap().iter().position(|g| g.label == label).unwrap();
    let (i1, i2) = (idx_of("g1"), idx_of("g2"));

    // snp_pseudo_alignment: one variable family, one column — g1's row
    // reads 'C' (its own member), g2's reads 'G'.
    let alignment = merged.snp_pseudo_alignment(None, None).expect("snp_pseudo_alignment");
    assert_eq!(alignment.sequences[i1], vec![b'C']);
    assert_eq!(alignment.sequences[i2], vec![b'G']);

    // raw_snp_distance: g1's single form (C) != g2's (G) at the family's
    // one eligible locus -> a SNP, not a shared site.
    let raw = merged.raw_snp_distance().expect("raw_snp_distance");
    assert_eq!(raw.snp[[i1, i2]], 1);
    assert_eq!(raw.snp[[i2, i1]], 1);
    assert_eq!(raw.shared[[i1, i2]], 0);
    assert_eq!(raw.shared[[i2, i1]], 0);

    // cardinality_tally: both genomes carry exactly one member of the
    // family (cardinality 1 each) -> one co-occurrence at [1][1].
    let cardinality = merged.cardinality_tally(&raw, 1.0).expect("cardinality_tally");
    assert_eq!(cardinality.counts[1][1], 1);
    let total: u64 = cardinality.counts.iter().flatten().sum();
    assert_eq!(total, 1, "no other cardinality pair should be tallied");

    // base_pair_tally: the one eligible, differing locus is C (base 1) vs
    // G (base 2).
    let base_pairs = merged.base_pair_tally(&raw, 1.0).expect("base_pair_tally");
    assert_eq!(base_pairs.counts[1][2], 1);
    assert_eq!(base_pairs.counts[2][1], 1);
    let total: u64 = base_pairs.counts.iter().flatten().sum();
    assert_eq!(total, 2, "no other base pair should be tallied");
    assert_eq!(base_pairs.same, [0, 0, 0, 0], "the two genomes never agree at this locus");
}

#[test]
fn sibling_annex_no_empty_masks_after_build() {
    let dir = tempdir().unwrap();
    let seq = b"ACGTACGTACGT".repeat(500);
    let g1 = build_single_genome_index(dir.path(), "g1", &seq);
    g1.build_sibling_annex().expect("build_sibling_annex");

    let n_layers = g1.n_layers(0).expect("partition meta");
    for l in 0..n_layers {
        let layer_dir = g1.layer_dir(0, l);
        let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).expect("annex open");
        for slot in 0..annex.len() {
            let mask = annex.get(slot).expect("slot must have an entry");
            assert!(
                mask.family_size() >= 1,
                "layer {l} slot {slot}: empty family mask (bits={:04b}) after build — \
                 indicates batch-offset bug or unseeded mask slot",
                mask.bits()
            );
        }
    }
}

#[test]
fn sibling_histogram_does_not_panic_on_partial_last_batch() {
    // Regression test for the bug where enumerate_kmers_batch().enumerate()
    // used the batch index instead of the sequence offset, leaving the
    // trailing slots of a non-multiple-of-BATCH_SIZE layer unseeded and
    // producing the minorant-only sentinel (0x10) that caused
    // sibling_family_size_histogram to panic with index u32::MAX.
    let dir = tempdir().unwrap();
    let seq = b"ACGTACGTACGT".repeat(500);
    let g1 = build_single_genome_index(dir.path(), "g1", &seq);
    g1.build_sibling_annex().expect("build_sibling_annex");
    let hist = g1.sibling_family_size_histogram().expect("sibling_family_size_histogram");
    let total: u64 = hist.iter().sum();
    assert!(total > 0, "histogram should contain at least one family");
}

/// `sibling_annex_one_sibling_each`'s fixture happens to put g2's k-mer in
/// layer 0 and g1's in layer 1 (a fresh single-genome index, then one
/// merge that adds exactly one new layer for content absent from the
/// first — confirmed empirically, not assumed) — a real two-layer index,
/// not a contrived one, so it exercises both write paths
/// `build_layer_sibling_annex` now has to get right: the seeding pass (a
/// slot's own base gets *this* layer's index directly) and the
/// cross-partition/cross-layer resolution pass (`PartitionCache::find`
/// reporting which layer the other family member actually lives in, not
/// just that it exists).
#[test]
fn sibling_annex_records_the_real_layer_of_each_family_member() {
    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 n_layers = merged.n_layers(0).unwrap();
    assert_eq!(n_layers, 2, "fixture assumption: one merge, one new layer");

    let g1_kmer = canonical(b"AACCGCTTAAG"); // own base C (1), sibling base G (2) — lives in layer 1
    let g2_kmer = canonical(b"AACCGGTTAAG"); // own base G (2), sibling base C (1) — lives in layer 0

    let a = annex_info_for(&merged, g1_kmer);
    assert_eq!(a.layer_value(1), Some(1), "g1's own base: seeded with its own layer (1)");
    assert_eq!(a.layer_value(2), Some(0), "g1's sibling (g2's form): resolved to its real layer (0)");

    let b = annex_info_for(&merged, g2_kmer);
    assert_eq!(b.layer_value(2), Some(0), "g2's own base: seeded with its own layer (0)");
    assert_eq!(b.layer_value(1), Some(1), "g2's sibling (g1's form): resolved to its real layer (1)");
}

#[test]
fn subsample_and_shannon_on_one_variable_family() {
    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 idx_of = |label: &str| merged.meta().genomes().unwrap().iter().position(|g| g.label == label).unwrap();
    let (i1, i2) = (idx_of("g1"), idx_of("g2"));

    // This fixture has exactly one non-monomorphic family (see
    // `family_scan_consumers_agree_on_one_sibling_each`) — total_count = 1.
    // Asking for far more than that must fall back to "keep everything",
    // identical to the unsampled (`None`) result.
    let sampled = merged.snp_pseudo_alignment(Some(1000), None).expect("snp_pseudo_alignment (subsample >> total)");
    assert_eq!(sampled.sequences[i1], vec![b'C']);
    assert_eq!(sampled.sequences[i2], vec![b'G']);

    // Asking for 0 must keep nothing — an empty (but still one-row-per-genome)
    // alignment, not an error.
    let empty = merged.snp_pseudo_alignment(Some(0), None).expect("snp_pseudo_alignment (subsample 0)");
    assert!(empty.sequences[i1].is_empty());
    assert!(empty.sequences[i2].is_empty());

    // shannon_entropy_csv: one row, both entropy15 and entropy4 = 1 bit
    // exactly (2 genomes present, each carrying exactly one distinct base,
    // uniform split -> -2*(0.5*log2(0.5)) = 1 — the two definitions
    // coincide here since no genome carries more than one base).
    let csv_path = dir.path().join("shannon.csv");
    merged.shannon_entropy_csv(&csv_path, Some(1000), None).expect("shannon_entropy_csv");
    let csv = std::fs::read_to_string(&csv_path).unwrap();
    let mut lines = csv.lines();
    assert_eq!(lines.next(), Some("layer,family_idx,entropy15,entropy4,family_size,n_genomes_present"));
    let row = lines.next().expect("exactly one data row");
    assert!(lines.next().is_none(), "exactly one non-monomorphic family in this fixture");
    let fields: Vec<&str> = row.split(',').collect();
    let entropy15: f64 = fields[2].parse().unwrap();
    let entropy4: f64 = fields[3].parse().unwrap();
    assert!((entropy15 - 1.0).abs() < 1e-6, "entropy15 should be exactly 1 bit, got {entropy15}");
    assert!((entropy4 - 1.0).abs() < 1e-6, "entropy4 should be exactly 1 bit, got {entropy4}");
    assert_eq!(fields[4], "2", "family_size");
    assert_eq!(fields[5], "2", "n_genomes_present");
}

#[test]
fn sankoff_bundle_matches_old_separate_calls() {
    // Same fixture as `sibling_annex_one_sibling_each`/
    // `subsample_and_shannon_on_one_variable_family`: one non-monomorphic
    // family, g1='C' (minorant), g2='G'. Regression check that fusing
    // raw_snp_distance/base_pair_tally/cardinality_tally/snp_pseudo_alignment
    // into `sankoff_bundle`'s two passes produces identical results to the
    // old four independent calls on `Selection::All` (no `--subsample`,
    // no `--entropy`) — see `DevDocMD/architecture/siblings.md`, "Wired into
    // `pack`...".
    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 n_genomes = merged.meta().genomes().unwrap().len();
    let exclude_mask = vec![false; n_genomes];
    let ratio_ceiling = 0.5;

    let expected_raw = merged.raw_snp_distance().expect("raw_snp_distance");
    let expected_base = merged.base_pair_tally(&expected_raw, ratio_ceiling).expect("base_pair_tally");
    let expected_card = merged.cardinality_tally(&expected_raw, ratio_ceiling).expect("cardinality_tally");
    let expected_alignment = merged.snp_pseudo_alignment(None, None).expect("snp_pseudo_alignment");

    let bundle = merged.sankoff_bundle(None, None, ratio_ceiling, &exclude_mask).expect("sankoff_bundle");

    assert_eq!(bundle.raw.snp, expected_raw.snp);
    assert_eq!(bundle.raw.shared, expected_raw.shared);
    assert_eq!(bundle.base_pair_tally.counts, expected_base.counts);
    assert_eq!(bundle.base_pair_tally.same, expected_base.same);
    assert_eq!(bundle.cardinality_tally.counts, expected_card.counts);
    assert_eq!(bundle.alignment.sequences, expected_alignment.sequences);
}

#[test]
fn base_pair_tally_accumulates_base_a_diagnostic() {
    // Diagnostic for a user-reported observation on real data:
    // `composition_transitions`'s row/column for base 'A' (index 0) was
    // *entirely* zero, including the diagonal (`same[0]`, "A stays A"),
    // despite 'A' appearing at ~6% frequency in the pseudo-alignment
    // itself. g1/g2 are identical (both central base 'A'), g3 differs
    // (central base 'C') at the same family — makes the family variable
    // (family_size=2) while g1/g2 form an eligible (ratio 0), ordinary
    // pair both resolving `single_form` to `Some(0)` = 'A'.
    let dir = tempdir().unwrap();
    let g1 = build_single_genome_index(dir.path(), "g1", b"AAAAAAAAAAA");
    let g2 = build_single_genome_index(dir.path(), "g2", b"AAAAAAAAAAA");
    let g3 = build_single_genome_index(dir.path(), "g3", b"AAAAACAAAAA");
    let mut rep = Reporter::new();
    let merged = KmerIndex::merge(
        &dir.path().join("merged.idx"),
        &[&g1, &g2, &g3],
        MergeMode::Presence,
        false,
        false,
        1.0,
        &mut rep,
    ).expect("merge");
    merged.build_sibling_annex().expect("build_sibling_annex");

    let n_genomes = merged.meta().genomes().unwrap().len();
    let exclude_mask = vec![false; n_genomes];
    let bundle = merged.sankoff_bundle(None, None, 0.5, &exclude_mask).expect("sankoff_bundle");

    assert!(
        bundle.base_pair_tally.same[0] >= 1,
        "base 'A' (index 0) must accumulate in BasePairTally.same for an identical, included g1/g2 \
         pair both resolving to 'A' — got same={:?}, counts={:?}",
        bundle.base_pair_tally.same, bundle.base_pair_tally.counts,
    );
}

#[test]
fn entropy_annex_builds_on_demand_and_biases_selection() {
    // Same fixture, same single non-monomorphic family, known entropy15 =
    // 1.0 bit exactly (see `subsample_and_shannon_on_one_variable_family`).
    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 layer_dirs = super::family_scan::sibling_layer_dirs(&merged).expect("layer dirs");
    for layer_dir in &layer_dirs {
        assert!(
            !layer_dir.join(ENTROPY_ANNEX_FILE_NAME).exists(),
            "entropy annex must not exist before any entropy-biased selection is requested"
        );
    }

    // This fixture spans 2 layers (a single-genome bootstrap layer plus the
    // merged one, see `merge_two`) — the one non-monomorphic family lives
    // in whichever layer actually got a non-empty selection, not
    // necessarily `layer_dirs[0]`.
    let find_layer = |selections: &[super::subsample::LayerSelection]| {
        selections.iter().position(|s| matches!(s, Some(set) if !set.is_empty()))
    };

    // mu == the family's real entropy, tiny sigma -> weight == 1 exactly,
    // p0 == 1.0 (no `--subsample`) -> the family is selected with
    // probability 1 (up to a measure-zero event on the uniform draw).
    let bias_center = EntropyBias { mu: 1.0, sigma: 0.5 };
    let selections = super::subsample::compute_selections(&merged, &layer_dirs, None, Some(bias_center))
        .expect("compute_selections (entropy-biased)");
    for layer_dir in &layer_dirs {
        assert!(
            layer_dir.join(ENTROPY_ANNEX_FILE_NAME).exists(),
            "compute_selections must build the entropy annex for every layer on first use"
        );
    }
    let li = find_layer(&selections).expect("mu at the true entropy must select the one non-monomorphic family, in some layer");
    assert_eq!(selections[li], Some(std::collections::HashSet::from([0])));

    // The annex now exists — read it back directly and check the stored
    // value matches `--shannon`'s own entropy15 computation exactly.
    let entropy_annex = EntropyAnnex::open(&layer_dirs[li].join(ENTROPY_ANNEX_FILE_NAME)).expect("open entropy annex");
    let stored = entropy_annex.get(0).expect("family 0 has a real (non-sentinel) entropy value");
    assert!((stored as f64 - 1.0).abs() < 1e-6, "stored entropy should be 1 bit, got {stored}");

    // mu far from any achievable entropy, tiny sigma -> weight underflows
    // to exactly 0.0 -> never selected, deterministically, in any layer.
    let bias_far = EntropyBias { mu: 100.0, sigma: 0.001 };
    let selections_far = super::subsample::compute_selections(&merged, &layer_dirs, None, Some(bias_far))
        .expect("compute_selections (entropy-biased, far)");
    assert!(find_layer(&selections_far).is_none(), "mu far from the true entropy must never select it, in any layer");
}