data-beans 0.6.12

#![allow(dead_code)]

use crate::sparse_data_visitors::*;
use crate::sparse_io_stack::SparseIoStack;
use crate::sparse_io_vector::SparseIoVec;
use legume_numeric::matrix::knn_match::ColumnDict;
use legume_numeric::matrix::traits::*;
use legume_numeric::param::dmatrix_gamma::*;
use legume_numeric::param::traits::Inference;
use legume_numeric::param::traits::*;
use log::{info, warn};
use nalgebra::DMatrix;
use rayon::prelude::*;
use std::ops::AddAssign;
use std::sync::{Arc, Mutex};

use crate::alg::random_projection::binary_sort_columns;

use rustc_hash::FxHashMap as HashMap;
type CscMat = nalgebra_sparse::CscMatrix<f32>;

/// Sparse pb-sample gene profile: `rows[pbsamp] = Vec<(gene_idx, sum)>` sorted
/// by `gene_idx`, as produced by `collect_pb_sample_gene_sums` and
/// consumed by `collect_matched_stat_coarse` and the refinement pass.
pub type GeneSums = Vec<Vec<(usize, f32)>>;

pub const DEFAULT_KNN: usize = 10;
pub const DEFAULT_OPT_ITER: usize = 100;

/// Per-level collapse output plus the partition hierarchy across levels.
///
/// Returned by [`collapse_columns_multilevel_with_hierarchy`] for
/// consumers (e.g. `graph-embedding-util`'s nested chain sampler) that
/// need parent/child maps between pb-samples at adjacent levels.
/// `cell_to_pb_per_level` is finest-first, parallel to `levels`:
/// `cell_to_pb_per_level[k][c] = pb_id at level k for cell c`.
mod pb_samples;
use pb_samples::{
    build_pb_sample_layout, build_pb_sample_to_cells, build_pb_samples,
    collect_pb_sample_gene_sums, per_batch_sc_neighbors,
};
pub use pb_samples::{PbSampleCollection, PbSampleLayout};
// Shared cross-batch pb-sample matching, reused by `refine_multilevel`.
pub(crate) use pb_samples::bbknn_match_one_pbsamp;
mod reassign_cells;
pub use reassign_cells::ReassignCellsParams;
mod pb_tree;
pub use pb_tree::{BelowEdge, ContrastGene, PbTree, PbTreeParams, RootRecord, SplitRecord};
mod refine;
pub use refine::StackCollapseOut;
use refine::{
    compute_level_sort_dims, fine_to_coarse_from_refined, pad_numeric_labels,
    refine_and_collect_single_layer, refine_and_collect_stack, split_anchored_finest_groups,
    RefineCollectCtx,
};
mod stats;
use stats::{
    collect_basic_stat_visitor, collect_batch_stat_visitor, collect_matched_stat_coarse,
    collect_matched_stat_visitor, merge_stat, optimize, KnnParams, DEFAULT_NUM_LEVELS,
};
pub use stats::{resample_and_optimize, CollapsedOut, CollapsedStat};
mod strata;
pub use strata::collapse_columns_multilevel_with_strata;

pub struct MultilevelCollapseOut {
    pub levels: Vec<CollapsedOut>,
    pub cell_to_pb_per_level: Vec<Vec<usize>>,
    /// The tree behind the finest partition, when it was grown (see
    /// `MultilevelParams::pb_tree`).
    pub pb_tree: Option<PbTree>,
}

/// Configuration for multi-level collapsing.
#[derive(Clone)]
pub struct MultilevelParams {
    pub knn_pb_samples: usize,
    pub num_levels: usize,
    pub sort_dim: usize,
    pub num_opt_iter: usize,
    /// BBKNN + Poisson DC-SBM refinement on top of the hash partition.
    /// Zero sweeps (`num_gibbs == 0 && num_greedy == 0`) keep the hash
    /// partition as is.
    pub refine: crate::alg::refine_multilevel::RefineParams,
    /// Which posterior planes the *output* `CollapsedOut` should carry.
    /// `MeanOnly` skips the sd / log_mean / log_sd allocations entirely —
    /// a big memory win for consumers that only read `posterior_mean()`
    /// (e.g. bge). Use `All` when the caller exports log-scale dictionaries.
    pub output_calibration: legume_numeric::param::traits::CalibrateTarget,
    /// Batch labels whose columns are ALREADY batch-corrected pseudobulks —
    /// a prior run's carried reference. `Some` switches the cross-batch
    /// counterfactual from pooled mutual adjustment to **greedy batch
    /// correction**: every pb-sample's counterfactual is drawn from the
    /// anchor batches only, so new batches are corrected *toward* the anchor
    /// frame while an anchor pb-sample self-matches (its δ settles at the
    /// prior ≈ 1 and is never re-adjusted). Anchored columns also skip
    /// pb-sample re-merging — each stays its own matching unit, at the
    /// resolution the prior run already paid for. `None` = pooled behavior,
    /// bit-identical to before this field existed.
    pub anchor_batches: Option<Vec<Box<str>>>,
    /// Batch labels whose columns are ALREADY summaries but are **not** a
    /// reference frame — bulk RNA-seq samples. Like anchors, each column
    /// bypasses the partition entirely (its own singleton pb-sample and its
    /// own singleton finest group — a bulk sample is never re-averaged), but
    /// unlike anchors it takes no part in cross-batch matching in EITHER
    /// direction: a bulk sample is a mixture over cell states, so matching it
    /// to a single-state pb-sample (as source or receiver) would let δ absorb
    /// composition rather than platform. Bulk therefore contributes observed
    /// sums only — the dictionary sees it, δ/γ never do, and δ for a bulk
    /// batch stays at its prior structurally. `None` = no bulk batches.
    pub bulk_batches: Option<Vec<Box<str>>>,
    /// Distinguish "unmeasured" from "measured as zero" when backends carry
    /// different feature panels (`RowAlignment::Union`): per-gene denominators
    /// count only the mass whose source measures the gene, and δ falls to its
    /// prior where a batch has no panel coverage. A no-op (bitwise) when every
    /// backend covers every row. Off for callers whose union is *intentional*
    /// structural disjointness — multiome stacks modalities on one axis, and
    /// re-weighting that is its own decision, not a side effect.
    pub observe_panels: bool,
    /// Retain the finest level's Gamma sufficient statistics even under
    /// `MeanOnly` calibration. `--emit-pb-reference` serializes
    /// `evidence_mean` — data sum over data denominator — which reads
    /// `a_stat`/`b_stat`; `MeanOnly` normally drops them per block to bound
    /// memory. Costs two extra `[D, S]` planes per finest-level parameter.
    pub keep_finest_stats: bool,
    /// Grow the finest partition as a tree from the marginal top nodes:
    /// reassign cells by likelihood, then bisect on residual components to
    /// one leaf target per level. `None` keeps the marginal hash, bit-identical
    /// to before this field existed.
    pub pb_tree: Option<PbTreeParams>,
    /// Per-cell CNV stratum (`0` = mixable residual). Crossed into finest
    /// codes so no pb-group mixes strata; BBKNN matches only within stratum;
    /// unmatched clone mass is excluded from the δ update so private CN
    /// stays in `mu_adjusted`. `None` is bit-identical to the pre-strata path.
    pub strata: Option<Vec<usize>>,
}

impl MultilevelParams {
    pub fn new(proj_dim: usize) -> Self {
        Self {
            knn_pb_samples: DEFAULT_KNN,
            num_levels: DEFAULT_NUM_LEVELS,
            sort_dim: proj_dim.min(12),
            num_opt_iter: DEFAULT_OPT_ITER,
            refine: crate::alg::refine_multilevel::RefineParams::default(),
            output_calibration: legume_numeric::param::traits::CalibrateTarget::All,
            anchor_batches: None,
            bulk_batches: None,
            observe_panels: true,
            keep_finest_stats: false,
            pb_tree: None,
            strata: None,
        }
    }

    /// Copy with [`Self::strata`] set (used by the thin `with_strata` wrapper).
    #[must_use]
    pub fn with_strata(&self, strata: Vec<usize>) -> Self {
        let mut p = self.clone();
        p.strata = Some(strata);
        p
    }
}

/// Bits needed to encode stratum labels in the low bits of a hash code.
/// A single occupied stratum needs no bits (identity with the unstratified
/// path); otherwise `ceil(log2(max_label + 1))`.
fn stratum_bits(strata: &[usize]) -> usize {
    let mut occupied: Vec<usize> = strata.to_vec();
    occupied.sort_unstable();
    occupied.dedup();
    if occupied.len() <= 1 {
        return 0;
    }
    let max = *occupied.last().unwrap_or(&0);
    let n = max + 1;
    (usize::BITS as usize - n.saturating_sub(1).leading_zeros() as usize).max(1)
}

/// Cross finest hash/tree codes with per-cell strata so no group mixes
/// strata. Stratum occupies the **low** bits. Level widths are **not**
/// bumped — expression budget (`1 << d` tree targets, refine bounds) stays
/// unchanged; hash/reproject masks add `s_bits` separately.
fn apply_strata_to_codes(
    codes: &[usize],
    level_dims: &[usize],
    strata: &[usize],
) -> anyhow::Result<(Vec<usize>, Vec<usize>, usize)> {
    anyhow::ensure!(
        codes.len() == strata.len(),
        "strata has {} entries, codes have {}",
        strata.len(),
        codes.len()
    );
    let s_bits = stratum_bits(strata);
    if s_bits == 0 {
        return Ok((codes.to_vec(), level_dims.to_vec(), 0));
    }
    let stratified: Vec<usize> = codes
        .iter()
        .zip(strata.iter())
        .map(|(&c, &s)| (c << s_bits) | s)
        .collect();
    Ok((stratified, level_dims.to_vec(), s_bits))
}

/// Finest cell codes plus what the per-level grouping needs to read them.
struct FinestCodes {
    /// Per-cell finest code (stratum bits in the low `strata_bits`).
    codes: Vec<usize>,
    /// Per-level expression widths, finest-first — un-bumped by strata.
    widths: Vec<usize>,
    /// The tree behind the finest partition, when one was grown.
    tree: Option<PbTree>,
    /// Width of the stratum field crossed into `codes` (`0` = no strata).
    strata_bits: usize,
}

/// Finest codes and their level widths (finest-first). Without residual
/// bits: the marginal sketch signs masked by `level_dims`. With them: the
/// marginal top nodes (optionally re-sorted by `reassign_cells`) are the
/// roots of a tree grown to one leaf target per level, `2^dim` for each of
/// `level_dims`, and the nested levels are packed into prefix codes whose
/// widths replace `level_dims`. Batch membership must already be
/// registered on `data_vec`.
///
/// When [`MultilevelParams::strata`] is set, stratum bits are crossed into
/// the codes after the tree/hash is grown (so leaf budgets stay expression-
/// only). Returned widths are the un-bumped expression dims; `strata_bits`
/// is carried separately for hash/reproject masks.
fn finest_codes(
    data_vec: &SparseIoVec,
    proj_kn: &DMatrix<f32>,
    level_dims: &[usize],
    params: &MultilevelParams,
) -> anyhow::Result<FinestCodes> {
    let finest_dim = level_dims[0];
    let nn = proj_kn.ncols();
    let kk = proj_kn.nrows().min(finest_dim).min(nn);
    let codes = binary_sort_columns(proj_kn, kk)?;
    let (codes, widths, tree) = match params.pb_tree.as_ref() {
        None => (codes, level_dims.to_vec(), None),
        Some(rb) => {
            let coarse_bits = if level_dims.len() >= 2 {
                *level_dims.last().expect("non-empty level dims")
            } else {
                stats::DEFAULT_COARSEST_SORT_DIM.min(kk)
            };
            let low_mask = (1usize << coarse_bits) - 1;
            let n = data_vec.num_columns();
            // Summary columns (carried reference, bulk) are not tree members; they
            // become singleton pb-samples downstream regardless of their code.
            let active: Vec<bool> = if data_vec.has_column_multiplicity() {
                (0..n)
                    .map(|c| (data_vec.column_multiplicity(c) - 1.0).abs() <= f32::EPSILON)
                    .collect()
            } else {
                vec![true; n]
            };
            let low: Vec<usize> = codes.iter().map(|&c| c & low_mask).collect();
            let (mut node, _) = crate::alg::dc_poisson::compact_labels(&low);
            let reassigned_cells = match rb.reassign_cells.as_ref() {
                Some(cr) => {
                    let col_to_batch = data_vec.get_batch_membership(0..n);
                    let csc = data_vec.read_columns_csc(0..n)?;
                    reassign_cells::reassign_cells_to_nodes(
                        &csc,
                        &col_to_batch,
                        data_vec.num_batches().max(1),
                        &active,
                        &mut node,
                        cr,
                    )
                }
                None => 0,
            };
            for (c, nd) in node.iter_mut().enumerate() {
                if !active[c] {
                    *nd = usize::MAX;
                }
            }
            // Leaf targets per level, coarse to fine — expression budget only.
            let targets: Vec<usize> = level_dims.iter().rev().map(|&d| 1usize << d).collect();
            let (codes, widths, mut tree) =
                pb_tree::build_tree(data_vec, &node, &codes, &targets, rb)?;
            tree.reassigned_cells = reassigned_cells;
            (codes, widths, Some(tree))
        }
    };
    maybe_stratify_codes(codes, widths, tree, params)
}

/// Apply [`MultilevelParams::strata`] to finest codes, or pass through.
fn maybe_stratify_codes(
    codes: Vec<usize>,
    widths: Vec<usize>,
    tree: Option<PbTree>,
    params: &MultilevelParams,
) -> anyhow::Result<FinestCodes> {
    let Some(strata) = params.strata.as_deref() else {
        return Ok(FinestCodes {
            codes,
            widths,
            tree,
            strata_bits: 0,
        });
    };
    anyhow::ensure!(
        strata.len() == codes.len(),
        "MultilevelParams.strata has {} entries, codes have {}",
        strata.len(),
        codes.len()
    );
    let (codes, widths, s_bits) = apply_strata_to_codes(&codes, &widths, strata)?;
    let n_occ = {
        let mut u = strata.to_vec();
        u.sort_unstable();
        u.dedup();
        u.len()
    };
    info!(
        "CNV strata: crossed finest codes with {} cell strata ({} occupied, {} stratum bits)",
        strata.len(),
        n_occ,
        s_bits
    );
    Ok(FinestCodes {
        codes,
        widths,
        tree,
        strata_bits: s_bits,
    })
}

/// Resolve [`MultilevelParams::anchor_batches`] / `bulk_batches` names to
/// batch indices.
///
/// A named batch that does not exist is an error, not a shrug: silently
/// dropping it would quietly fall back to pooled mutual adjustment — the
/// exact behavior anchoring (and the bulk exclusion) exists to prevent.
fn resolve_named_batches(
    data_vec: &SparseIoVec,
    role: &str,
    names: Option<&[Box<str>]>,
) -> anyhow::Result<Option<Vec<usize>>> {
    let Some(names) = names else { return Ok(None) };
    let map = data_vec
        .batch_name_map()
        .ok_or_else(|| anyhow::anyhow!("{role} batches given but no batches are registered"))?;
    let mut idx = Vec::with_capacity(names.len());
    for n in names {
        let Some(&b) = map.get(n) else {
            anyhow::bail!(
                "{role} batch `{n}` is not among the registered batches ({:?})",
                map.keys().collect::<Vec<_>>(),
            );
        };
        idx.push(b);
    }
    Ok(Some(idx))
}

/// A batch cannot be both an anchor (a counterfactual source) and bulk
/// (barred from matching): the two roles contradict each other, and which
/// one silently won would decide whether composition leaks into δ.
/// Greedy default: when bulk batches are named and no anchor is given, every
/// NON-bulk batch becomes the anchor frame.
///
/// This is the same discipline `senna update` applies to a carried
/// pb_reference — only the new samples are adjusted, the established frame
/// stays fixed. Concretely it makes bulk draw its counterfactual from the
/// cells and be corrected toward them, while the cells self-match (their δ
/// settles at the prior) and bulk, never being in the anchor set, cannot
/// serve as anyone's counterfactual. Pooled mutual adjustment — where the
/// cell frame drifts toward bulk — is what this avoids.
fn greedy_anchor_for_bulk(
    data_vec: &SparseIoVec,
    anchors: Option<Vec<usize>>,
    bulk: Option<&[usize]>,
) -> Option<Vec<usize>> {
    match (anchors, bulk) {
        (Some(a), _) => Some(a),
        (None, Some(b)) if !b.is_empty() => {
            let all = data_vec.num_batches();
            let frame: Vec<usize> = (0..all).filter(|i| !b.contains(i)).collect();
            (!frame.is_empty()).then(|| {
                info!(
                    "Greedy bulk correction: {} bulk batch(es) corrected toward {} cell batch(es); \
                     the cell frame is anchored and does not move",
                    b.len(),
                    frame.len()
                );
                frame
            })
        }
        (None, _) => None,
    }
}

fn ensure_disjoint_roles(anchors: Option<&[usize]>, bulk: Option<&[usize]>) -> anyhow::Result<()> {
    if let (Some(a), Some(b)) = (anchors, bulk) {
        if let Some(shared) = a.iter().find(|x| b.contains(x)) {
            anyhow::bail!(
                "batch index {shared} is named as both an anchor and a bulk batch; \
                 the roles are mutually exclusive"
            );
        }
    }
    Ok(())
}

/// Attach per-panel observability to a freshly collected fine-level stat.
///
/// `size_ds[g, s] = Σ_{sources b measuring g} count[b, s]` where `count[b, s]`
/// is the multiplicity-weighted mass batch/backends contribute to sample `s` —
/// the factorized form, so the cost is `O(D·S·B + N)`, never `O(D·N)`. The
/// per-(gene, batch) δ mask is 1 where any source in the batch measures the
/// gene. Uses whatever group assignment currently stands on `data_vec`, i.e.
/// exactly the one that filled the stat.
///
/// Leaves the stat untouched (the bitwise-historical path) when every backend
/// covers every row, or when columns merge several backends (column-union /
/// multiome), whose panel semantics are deliberate.
fn attach_observability(stat: &mut CollapsedStat, data_vec: &SparseIoVec) -> anyhow::Result<()> {
    let Some(coverage) = data_vec.row_coverage_by_backend() else {
        return Ok(());
    };
    let ncols = data_vec.num_columns();
    let mut sources = Vec::with_capacity(ncols);
    for c in 0..ncols {
        match data_vec.column_source(c) {
            Some(b) => sources.push(b),
            None => {
                warn!(
                    "panel observability skipped: column {c} merges several backends \
                     (column-union alignment)"
                );
                return Ok(());
            }
        }
    }

    let num_genes = stat.num_genes();
    let num_samples = stat.num_samples();
    let num_sources = data_vec.len();
    let cell_to_group = data_vec.get_group_membership(0..ncols)?;

    // Multiplicity-weighted mass per (source, sample) and per-batch source use.
    let mut count_bs = DMatrix::<f32>::zeros(num_sources, num_samples);
    let batch_of = data_vec.get_batch_membership(0..ncols);
    let num_batches = stat.num_batches();
    let mut source_in_batch = vec![vec![false; num_batches]; num_sources];
    for c in 0..ncols {
        let s = cell_to_group[c];
        if s < num_samples {
            count_bs[(sources[c], s)] += data_vec.column_multiplicity(c);
        }
        if let Some(&b) = batch_of.get(c) {
            if b < num_batches {
                source_in_batch[sources[c]][b] = true;
            }
        }
    }

    let mut size_ds = DMatrix::<f32>::zeros(num_genes, num_samples);
    for (src, cov) in coverage.iter().enumerate() {
        anyhow::ensure!(
            cov.len() == num_genes,
            "row coverage has {} rows but the stat has {num_genes} genes",
            cov.len(),
        );
        for (g, &covered) in cov.iter().enumerate() {
            if covered {
                for s in 0..num_samples {
                    size_ds[(g, s)] += count_bs[(src, s)];
                }
            }
        }
    }

    let mut mask_db = DMatrix::<f32>::zeros(num_genes, num_batches);
    for (src, cov) in coverage.iter().enumerate() {
        for (b, &used) in source_in_batch[src].iter().enumerate() {
            if used {
                for (g, &covered) in cov.iter().enumerate() {
                    if covered {
                        mask_db[(g, b)] = 1.0;
                    }
                }
            }
        }
    }

    info!(
        "panel observability attached: {} of {} (gene, sample) entries below full size",
        size_ds
            .column_iter()
            .enumerate()
            .map(|(s, col)| col.iter().filter(|&&v| v < stat.size_s[s]).count())
            .sum::<usize>(),
        num_genes * num_samples,
    );
    stat.size_ds = Some(size_ds);
    stat.obs_mask_db = (mask_db.iter().any(|&v| v == 0.0)).then_some(mask_db);
    Ok(())
}

pub struct EmptyArg {}

#[cfg(debug_assertions)]
use log::debug;

/// Given a feature/projection matrix (factor x cells), we assign each
/// cell to a sample and return pseudobulk (collapsed) matrices
///
/// (1) Register batches if needed (2) collapse columns/cells into samples
///
pub trait CollapsingOps {
    ///
    /// Collapse columns/cells into samples as allocated by
    /// `assign_columns_to_samples`
    ///
    /// # Arguments
    /// * `cells_per_group` - number of cells per sample (None: no down sampling)
    /// * `knn_batches` - number of nearest neighbour batches
    /// * `knn_cells` - number of nearest neighbors for building HNSW (default: 10)
    /// * `reference` - reference batch for counterfactual inference
    /// * `num_opt_iter` - number of optimization iterations (default: 100)
    ///
    fn collapse_columns(
        &self,
        knn_batches: Option<usize>,
        knn_cells: Option<usize>,
        reference_batch_names: Option<&[Box<str>]>,
        num_opt_iter: Option<usize>,
    ) -> anyhow::Result<CollapsedOut>;

    /// Register batch information and build a `HnswMap` object for
    /// each batch for fast nearest neighbor search within each batch
    /// and store them in the `SparseIoVec`
    ///
    /// # Arguments
    /// * `proj_kn` - random projection matrix
    /// * `col_to_batch` - map: cell -> batch
    fn build_hnsw_per_batch<T>(
        &mut self,
        proj_kn: &nalgebra::DMatrix<f32>,
        col_to_batch: &[T],
    ) -> anyhow::Result<()>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString;

    fn collect_basic_stat(&self, stat: &mut CollapsedStat) -> anyhow::Result<()>;

    fn collect_batch_stat(&self, stat: &mut CollapsedStat) -> anyhow::Result<()>;

    fn collect_matched_stat(
        &self,
        knn_batches: usize,
        knn_cols: usize,
        reference_indices: Option<&[usize]>,
        stat: &mut CollapsedStat,
    ) -> anyhow::Result<()>;
}

impl CollapsingOps for SparseIoVec {
    fn build_hnsw_per_batch<T>(
        &mut self,
        proj_kn: &nalgebra::DMatrix<f32>,
        col_to_batch: &[T],
    ) -> anyhow::Result<()>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        info!("creating batch-specific HNSW maps ...");
        self.register_batches_dmatrix(proj_kn, col_to_batch)?;

        info!(
            "found {} columns across {} batches",
            self.num_columns(),
            self.num_batches()
        );

        Ok(())
    }

    fn collapse_columns(
        &self,
        knn_batches: Option<usize>,
        knn_cells: Option<usize>,
        reference_batch_names: Option<&[Box<str>]>,
        num_opt_iter: Option<usize>,
    ) -> anyhow::Result<CollapsedOut> {
        let group_to_cols = self.take_grouped_columns().ok_or(anyhow::anyhow!(
            "The columns were not assigned before. Call `assign_columns_to_groups`"
        ))?;

        let num_features = self.num_rows();
        let num_groups = group_to_cols.len();
        let num_batches = self.num_batches();

        let mut stat = CollapsedStat::new(num_features, num_groups, num_batches);
        info!("basic statistics across {} groups", num_groups);
        self.collect_basic_stat(&mut stat)?;

        if num_batches > 1 {
            info!(
                "batch-specific statistics across {} batches over {} samples",
                num_batches, num_groups
            );

            let batch_name_map = self
                .batch_name_map()
                .ok_or(anyhow::anyhow!("unable to read batch names"))?;

            let reference_indices = reference_batch_names.map(|x| {
                x.iter()
                    .filter_map(|b| batch_name_map.get(b))
                    .copied()
                    .collect::<Vec<_>>()
            });

            if let Some(r) = reference_indices.as_ref() {
                if r.is_empty() {
                    let ref_names = reference_batch_names
                        .unwrap()
                        .iter()
                        .map(|x| x.to_string())
                        .collect::<Vec<_>>()
                        .join(",");

                    let bat_names = self
                        .batch_names()
                        .unwrap()
                        .iter()
                        .map(|x| x.to_string())
                        .collect::<Vec<_>>()
                        .join(",");

                    warn!("{} vs. {}", ref_names, bat_names);

                    return Err(anyhow::anyhow!("no reference batch names matched!"));
                }
            }

            self.collect_batch_stat(&mut stat)?;

            info!(
                "counterfactual inference across {} batches over {} samples",
                num_batches, num_groups,
            );

            let knn_batches = knn_batches.unwrap_or(2);
            let knn_cells = knn_cells.unwrap_or(DEFAULT_KNN);

            self.collect_matched_stat(
                knn_batches,
                knn_cells,
                reference_indices.as_deref(),
                &mut stat,
            )?;
        } // if num_batches > 1

        /////////////////////////////
        // Resolve mean parameters //
        /////////////////////////////

        info!("optimizing the collapsed parameters...");
        let (a0, b0) = (1_f32, 1_f32);
        optimize(
            &stat,
            (a0, b0),
            num_opt_iter.unwrap_or(DEFAULT_OPT_ITER),
            "Optimizing",
            CalibrateTarget::All,
            false,
        )
    }

    fn collect_basic_stat(&self, stat: &mut CollapsedStat) -> anyhow::Result<()> {
        self.visit_columns_by_group(&collect_basic_stat_visitor, &EmptyArg {}, stat)
    }

    fn collect_batch_stat(&self, stat: &mut CollapsedStat) -> anyhow::Result<()> {
        self.visit_columns_by_group(&collect_batch_stat_visitor, &EmptyArg {}, stat)
    }

    fn collect_matched_stat(
        &self,
        knn_batches: usize,
        knn_cells: usize,
        reference_indices: Option<&[usize]>,
        stat: &mut CollapsedStat,
    ) -> anyhow::Result<()> {
        // The reference batches source the counterfactual and pin δ.
        stat.anchor_batches = reference_indices.map(<[usize]>::to_vec).unwrap_or_default();
        self.visit_columns_by_group(
            &collect_matched_stat_visitor,
            &KnnParams {
                knn_batches,
                knn_cells,
                reference_indices,
            },
            stat,
        )
    }
}

pub trait MultilevelCollapsingOps {
    type LevelOutput;

    fn collapse_columns_multilevel<T>(
        &mut self,
        proj_kn: &DMatrix<f32>,
        batch_membership: &[T],
        params: &MultilevelParams,
    ) -> anyhow::Result<Self::LevelOutput>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString;

    fn collapse_columns_multilevel_vec<T>(
        &mut self,
        proj_kn: &DMatrix<f32>,
        batch_membership: &[T],
        params: &MultilevelParams,
    ) -> anyhow::Result<Vec<Self::LevelOutput>>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString;
}

/// Same as `SparseIoVec::collapse_columns_multilevel_vec`, but also returns
/// the per-level cell → pb mapping needed for hierarchical / nested
/// chain sampling in downstream consumers (currently
/// `graph-embedding-util`).
pub fn collapse_columns_multilevel_with_hierarchy<T>(
    data_vec: &mut SparseIoVec,
    proj_kn: &DMatrix<f32>,
    batch_membership: &[T],
    params: &MultilevelParams,
) -> anyhow::Result<MultilevelCollapseOut>
where
    T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
{
    let sort_dim = params.sort_dim;
    let knn = params.knn_pb_samples;
    let opt_iter = params.num_opt_iter;

    data_vec.register_batch_membership(batch_membership);
    let num_features = data_vec.num_rows();
    let num_batches = data_vec.num_batches();
    if num_batches >= 2 {
        data_vec.build_hnsw_per_batch(proj_kn, batch_membership)?;
    }

    let level_dims = compute_level_sort_dims(sort_dim, params.num_levels);
    let FinestCodes {
        codes: fine_codes,
        widths: level_dims,
        tree: pb_tree,
        strata_bits,
    } = finest_codes(data_vec, proj_kn, &level_dims, params)?;
    data_vec.assign_groups(&fine_codes, None);

    let group_to_cols = data_vec
        .take_grouped_columns()
        .ok_or_else(|| anyhow::anyhow!("columns not assigned"))?
        .clone();

    let refine_params = &params.refine;

    let anchor_batches =
        resolve_named_batches(data_vec, "anchor", params.anchor_batches.as_deref())?;
    let bulk_batches = resolve_named_batches(data_vec, "bulk", params.bulk_batches.as_deref())?;
    ensure_disjoint_roles(anchor_batches.as_deref(), bulk_batches.as_deref())?;
    // `summary_batches` are the batches whose columns are ALREADY summaries
    // (a carried pb_reference, or bulk) and so must never be re-averaged.
    // `anchor_batches` below is the MATCHING frame, which after the greedy
    // default is the cells — and cells must still collapse into pseudobulks.
    // Conflating the two turns every cell into its own pb-sample.
    let summary_batches = anchor_batches.clone();
    let anchor_batches = greedy_anchor_for_bulk(data_vec, anchor_batches, bulk_batches.as_deref());
    let ctx = RefineCollectCtx {
        fine_codes: &fine_codes,
        group_to_cols_finest: &group_to_cols,
        level_dims: &level_dims,
        num_features,
        num_batches,
        knn,
        opt_iter,
        refine_params,
        output_calibration: params.output_calibration,
        anchor_batches: anchor_batches.as_deref(),
        summary_batches: summary_batches.as_deref(),
        bulk_batches: bulk_batches.as_deref(),
        observe_panels: params.observe_panels,
        keep_finest_stats: params.keep_finest_stats,
        pb_tree: pb_tree.as_ref(),
        cell_to_stratum: params.strata.as_deref(),
        exclude_unmatched_from_delta: params.strata.is_some(),
        strata_bits,
    };
    refine_and_collect_single_layer(data_vec, proj_kn, &ctx)
}

/// Variant of [`collapse_columns_multilevel_with_hierarchy`] that **skips
/// the BBKNN + Poisson DC-SBM refinement**, instead synthesising the
/// per-pb-sample group assignment from a caller-supplied
/// `cell_to_pb_per_level` (finest-first; the same shape returned in
/// [`MultilevelCollapseOut.cell_to_pb_per_level`] by the refining
/// entry point). Typical use: `senna {topic, itopic, ce-topic} --from`
/// inheriting a prior run's partition.
///
/// Each level's pb-sample → group label is decided by majority vote
/// across the cells in that pb-sample. The downstream `optimize` step
/// still runs per level, so the returned `CollapsedOut` posteriors
/// reflect this run's batch model and priors — only the expensive
/// clustering work is bypassed.
pub fn collapse_columns_multilevel_with_partition<T>(
    data_vec: &mut SparseIoVec,
    proj_kn: &DMatrix<f32>,
    batch_membership: &[T],
    params: &MultilevelParams,
    cell_to_pb_per_level: &[Vec<usize>],
) -> anyhow::Result<MultilevelCollapseOut>
where
    T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
{
    let knn = params.knn_pb_samples;
    let opt_iter = params.num_opt_iter;

    data_vec.register_batch_membership(batch_membership);
    let num_features = data_vec.num_rows();
    let num_batches = data_vec.num_batches();
    if num_batches >= 2 {
        data_vec.build_hnsw_per_batch(proj_kn, batch_membership)?;
    }

    let level_dims = compute_level_sort_dims(params.sort_dim, params.num_levels);
    anyhow::ensure!(
        cell_to_pb_per_level.len() == level_dims.len(),
        "inherited cell_to_pb has {} levels but --num-levels is {}; \
         pass --num-levels to match the source run",
        cell_to_pb_per_level.len(),
        level_dims.len(),
    );
    // pb-samples are built from the INHERITED finest membership, so each
    // pb-sample is one (batch, source finest group) intersection and the
    // per-level vote below is unanimous by construction. Fresh marginal
    // leaves would straddle a source partition whose high bits came from
    // within-node residuals, and a modal vote over them would erode it.
    // pb-samples are still needed for the cross-batch matched-stat path on
    // multi-batch data; refinement is what we skip.
    let inherited_finest = &cell_to_pb_per_level[0];
    anyhow::ensure!(
        inherited_finest.len() == proj_kn.ncols(),
        "inherited cell_to_pb finest level has {} cells, data has {}",
        inherited_finest.len(),
        proj_kn.ncols()
    );
    let k_inherited = inherited_finest.iter().copied().max().map_or(0, |m| m + 1);
    data_vec.assign_groups(&pad_numeric_labels(inherited_finest, k_inherited), None);

    let anchor_batches =
        resolve_named_batches(data_vec, "anchor", params.anchor_batches.as_deref())?;
    let bulk_batches = resolve_named_batches(data_vec, "bulk", params.bulk_batches.as_deref())?;
    ensure_disjoint_roles(anchor_batches.as_deref(), bulk_batches.as_deref())?;
    // `summary_batches` are the batches whose columns are ALREADY summaries
    // (a carried pb_reference, or bulk) and so must never be re-averaged.
    // `anchor_batches` below is the MATCHING frame, which after the greedy
    // default is the cells — and cells must still collapse into pseudobulks.
    // Conflating the two turns every cell into its own pb-sample.
    let summary_batches = anchor_batches.clone();
    let anchor_batches = greedy_anchor_for_bulk(data_vec, anchor_batches, bulk_batches.as_deref());
    let pb_samples = build_pb_samples(
        data_vec,
        proj_kn,
        num_features,
        summary_batches.as_deref().unwrap_or(&[]),
        bulk_batches.as_deref().unwrap_or(&[]),
        params.strata.as_deref(),
    )?;
    let num_pb = pb_samples.layout.cell_counts.len();
    let ncols = proj_kn.ncols();
    let pb_sample_to_cells = build_pb_sample_to_cells(&pb_samples.layout);

    // Synthesize a RefinedAssignment from the inherited cell→pb
    // membership per pb-sample at each level. The vote is unanimous at the
    // finest level by construction; coarser levels vote in case the source
    // hierarchy was not strictly nested.
    let num_levels = level_dims.len();
    let mut pbsamp_to_group: Vec<Vec<usize>> = Vec::with_capacity(num_levels);
    let mut num_groups_per_level: Vec<usize> = Vec::with_capacity(num_levels);
    for (lvl_idx, lvl) in cell_to_pb_per_level.iter().enumerate() {
        anyhow::ensure!(
            lvl.len() == ncols,
            "inherited cell_to_pb level {} has {} cells, data has {}",
            lvl_idx,
            lvl.len(),
            ncols
        );
        let mut p2g: Vec<usize> = Vec::with_capacity(num_pb);
        for cells in &pb_sample_to_cells {
            p2g.push(modal_group(cells, lvl));
        }
        let (compact, k) = crate::alg::refine_multilevel::compact_labels(&p2g);
        num_groups_per_level.push(k);
        pbsamp_to_group.push(compact);
    }
    let refined = split_anchored_finest_groups(
        crate::alg::refine_multilevel::RefinedAssignment {
            pbsamp_to_group,
            num_groups_per_level,
        },
        &pb_samples.layout,
    );

    info!(
        "Inherited partition: {} cells, {} pb-samples, finest k={} (skipped BBKNN + DC-SBM refinement)",
        ncols, num_pb, refined.num_groups_per_level[0]
    );

    // From here the path matches refine_and_collect_single_layer's
    // post-refinement tail: assign groups, collect stats, fit Gamma
    // posteriors level-by-level via merge_stat.
    let k_finest = refined.num_groups_per_level[0];
    let mut cell_to_group_finest = vec![0usize; ncols];
    for (pbsamp, cells) in pb_sample_to_cells.iter().enumerate() {
        let g = refined.pbsamp_to_group[0][pbsamp];
        for &c in cells {
            cell_to_group_finest[c] = g;
        }
    }
    let finest_str = pad_numeric_labels(&cell_to_group_finest, k_finest);
    data_vec.assign_groups(&finest_str, None);
    debug_assert_eq!(data_vec.num_groups(), k_finest);

    let mut fine_stat = CollapsedStat::new(num_features, k_finest, num_batches);
    fine_stat.exclude_unmatched_from_delta = params.strata.is_some();
    info!("Collecting basic stats over {} groups ...", k_finest);
    data_vec.collect_basic_stat(&mut fine_stat)?;
    if num_batches >= 2 {
        info!(
            "Collecting per-batch stats over {} groups × {} batches ...",
            k_finest, num_batches
        );
        data_vec.collect_batch_stat(&mut fine_stat)?;
        let batch_knn = data_vec
            .batch_knn_lookup()
            .ok_or_else(|| anyhow::anyhow!("batch_knn_lookup not built"))?;
        info!(
            "Collecting cross-batch matched stats (knn={}) over {} pb-samples ...",
            knn, num_pb
        );
        collect_matched_stat_coarse(
            &pb_samples.layout,
            &pb_samples.gene_sums,
            &refined.pbsamp_to_group[0],
            batch_knn.as_slice(),
            knn,
            anchor_batches.as_deref(),
            &mut fine_stat,
        )?;
    }
    if params.observe_panels {
        attach_observability(&mut fine_stat, data_vec)?;
    }

    let mut results: Vec<CollapsedOut> = Vec::with_capacity(num_levels);
    info!(
        "Level 1/{}: inherited k={} (finest; {} cells)",
        num_levels, k_finest, ncols
    );
    let finest_out = optimize(
        &fine_stat,
        (1.0, 1.0),
        opt_iter,
        &format!("Inherit L1/{}", num_levels),
        CalibrateTarget::All,
        false,
    )?;
    results.push(finest_out);

    let mut prev_stat = fine_stat;
    for level in 1..num_levels {
        let k_prev = refined.num_groups_per_level[level - 1];
        let k_level = refined.num_groups_per_level[level];
        let fine_to_coarse = fine_to_coarse_from_refined(
            &refined.pbsamp_to_group[level - 1],
            &refined.pbsamp_to_group[level],
            k_prev,
        );
        let coarse_stat = merge_stat(&prev_stat, &fine_to_coarse, k_level);
        info!(
            "Level {}/{}: inherited k={} (merged from {})",
            level + 1,
            num_levels,
            k_level,
            k_prev
        );
        let level_opt_iter = (opt_iter / 2).max(10);
        let out = optimize(
            &coarse_stat,
            (1.0, 1.0),
            level_opt_iter,
            &format!("Inherit L{}/{}", level + 1, num_levels),
            CalibrateTarget::All,
            false,
        )?;
        results.push(out);
        prev_stat = coarse_stat;
    }
    info!(
        "Fitted pseudobulk posteriors for {} inherited levels: k = {:?} (finest first)",
        num_levels, refined.num_groups_per_level
    );

    // Re-derive cell_to_pb_per_level from the post-modal-vote groups so
    // the returned struct is self-consistent — small drift vs the
    // inherited values is expected when pb-samples spanned multiple
    // source groups (handled by majority).
    let mut cell_to_pb_per_level_out: Vec<Vec<usize>> = Vec::with_capacity(num_levels);
    for level in 0..num_levels {
        let mut c2g = vec![0usize; ncols];
        for (pbsamp, cells) in pb_sample_to_cells.iter().enumerate() {
            let g = refined.pbsamp_to_group[level][pbsamp];
            for &c in cells {
                c2g[c] = g;
            }
        }
        cell_to_pb_per_level_out.push(c2g);
    }

    Ok(MultilevelCollapseOut {
        levels: results,
        cell_to_pb_per_level: cell_to_pb_per_level_out,
        pb_tree: None,
    })
}

/// Modal `lvl[c]` over `c ∈ cells`. Returns 0 when `cells` is empty.
/// Fast small-Vec path for the common case where pb-samples contain
/// just a handful of cells.
fn modal_group(cells: &[usize], lvl: &[usize]) -> usize {
    match cells {
        [] => 0,
        [c] => lvl[*c],
        _ => {
            use rustc_hash::FxHashMap;
            let mut counts: FxHashMap<usize, usize> = FxHashMap::default();
            for &c in cells {
                *counts.entry(lvl[c]).or_insert(0) += 1;
            }
            counts
                .into_iter()
                .max_by_key(|&(_, n)| n)
                .map(|(g, _)| g)
                .unwrap_or(0)
        }
    }
}

impl MultilevelCollapsingOps for SparseIoVec {
    type LevelOutput = CollapsedOut;

    fn collapse_columns_multilevel<T>(
        &mut self,
        proj_kn: &DMatrix<f32>,
        batch_membership: &[T],
        params: &MultilevelParams,
    ) -> anyhow::Result<CollapsedOut>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        let mut results =
            self.collapse_columns_multilevel_vec(proj_kn, batch_membership, params)?;
        if results.is_empty() {
            return Err(anyhow::anyhow!("no levels processed"));
        }
        Ok(results.remove(0))
    }

    fn collapse_columns_multilevel_vec<T>(
        &mut self,
        proj_kn: &DMatrix<f32>,
        batch_membership: &[T],
        params: &MultilevelParams,
    ) -> anyhow::Result<Vec<CollapsedOut>>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        let sort_dim = params.sort_dim;
        let knn = params.knn_pb_samples;
        let opt_iter = params.num_opt_iter;

        self.register_batch_membership(batch_membership);
        let num_features = self.num_rows();
        let num_batches = self.num_batches();
        if num_batches >= 2 {
            self.build_hnsw_per_batch(proj_kn, batch_membership)?;
        }

        // Level dims: [finest, ..., coarsest]
        let level_dims = compute_level_sort_dims(sort_dim, params.num_levels);

        info!(
            "Multi-level collapsing (fine→coarse): {} levels, sort_dims={:?}, {} batches",
            level_dims.len(),
            level_dims,
            num_batches,
        );

        // Compute binary codes at finest resolution once
        let FinestCodes {
            codes: fine_codes,
            widths: level_dims,
            tree: pb_tree,
            strata_bits,
        } = finest_codes(self, proj_kn, &level_dims, params)?;

        // Partition at finest level
        self.assign_groups(&fine_codes, None);

        let group_to_cols = self
            .take_grouped_columns()
            .ok_or(anyhow::anyhow!("columns not assigned"))?
            .clone();

        ////////////////////////////////////////////////////////////////////
        // Opt-in refinement path: BBKNN + Poisson DC-SBM over pb-samples //
        ////////////////////////////////////////////////////////////////////

        let refine_params = &params.refine;
        let anchor_batches =
            resolve_named_batches(self, "anchor", params.anchor_batches.as_deref())?;
        let bulk_batches = resolve_named_batches(self, "bulk", params.bulk_batches.as_deref())?;
        ensure_disjoint_roles(anchor_batches.as_deref(), bulk_batches.as_deref())?;
        let summary_batches = anchor_batches.clone();
        let anchor_batches = greedy_anchor_for_bulk(self, anchor_batches, bulk_batches.as_deref());
        let ctx = RefineCollectCtx {
            fine_codes: &fine_codes,
            group_to_cols_finest: &group_to_cols,
            level_dims: &level_dims,
            num_features,
            num_batches,
            knn,
            opt_iter,
            refine_params,
            output_calibration: params.output_calibration,
            anchor_batches: anchor_batches.as_deref(),
            summary_batches: summary_batches.as_deref(),
            bulk_batches: bulk_batches.as_deref(),
            observe_panels: params.observe_panels,
            keep_finest_stats: params.keep_finest_stats,
            pb_tree: pb_tree.as_ref(),
            cell_to_stratum: params.strata.as_deref(),
            exclude_unmatched_from_delta: params.strata.is_some(),
            strata_bits,
        };
        refine_and_collect_single_layer(self, proj_kn, &ctx).map(|out| out.levels)
    }
}

impl MultilevelCollapsingOps for SparseIoStack {
    type LevelOutput = Vec<CollapsedOut>;

    fn collapse_columns_multilevel<T>(
        &mut self,
        proj_kn: &DMatrix<f32>,
        batch_membership: &[T],
        params: &MultilevelParams,
    ) -> anyhow::Result<Vec<CollapsedOut>>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        let mut results =
            self.collapse_columns_multilevel_vec(proj_kn, batch_membership, params)?;
        if results.is_empty() {
            return Err(anyhow::anyhow!("no levels processed"));
        }
        Ok(results.remove(0))
    }

    fn collapse_columns_multilevel_vec<T>(
        &mut self,
        proj_kn: &DMatrix<f32>,
        batch_membership: &[T],
        params: &MultilevelParams,
    ) -> anyhow::Result<Vec<Vec<CollapsedOut>>>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        Ok(
            collapse_stack_multilevel_with_hierarchy(self, proj_kn, batch_membership, params)?
                .levels,
        )
    }
}

/// The `SparseIoStack` collapse, also returning the per-level cell → pb
/// membership (finest first) the layers share — the stack counterpart of
/// [`collapse_columns_multilevel_with_hierarchy`]. With
/// `params.keep_finest_stats` the finest level keeps its sufficient statistics,
/// which `CollapsedOut::observed_counts` reads.
pub fn collapse_stack_multilevel_with_hierarchy<T>(
    stack: &mut SparseIoStack,
    proj_kn: &DMatrix<f32>,
    batch_membership: &[T],
    params: &MultilevelParams,
) -> anyhow::Result<StackCollapseOut>
where
    T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
{
    let num_layers = stack.num_types();
    if num_layers == 0 {
        return Err(anyhow::anyhow!("empty SparseIoStack"));
    }

    let sort_dim = params.sort_dim;
    let knn = params.knn_pb_samples;
    let opt_iter = params.num_opt_iter;

    stack.register_batch_membership(batch_membership);

    // Use first layer for num_batches (all layers share the same columns)
    let num_batches = stack.stack[0].num_batches();
    if num_batches >= 2 {
        for layer in stack.stack.iter_mut() {
            layer.build_hnsw_per_batch(proj_kn, batch_membership)?;
        }
    }

    let ncols = proj_kn.ncols();
    let refine_params = &params.refine;
    let level_dims = compute_level_sort_dims(sort_dim, params.num_levels);
    let finest_dim = level_dims[0];
    let kk = proj_kn.nrows().min(finest_dim).min(ncols);
    let codes = binary_sort_columns(proj_kn, kk)?;
    let (fine_codes, level_dims, strata_bits) = match params.strata.as_deref() {
        Some(strata) => {
            let (c, d, s) = apply_strata_to_codes(&codes, &level_dims, strata)?;
            (c, d, s)
        }
        None => (codes, level_dims, 0),
    };
    for layer in stack.stack.iter_mut() {
        layer.assign_groups(&fine_codes, None);
    }
    let group_to_cols = stack.stack[0]
        .take_grouped_columns()
        .ok_or(anyhow::anyhow!("columns not assigned"))?
        .clone();
    let num_features = stack.stack[0].num_rows();
    anyhow::ensure!(
        params.anchor_batches.is_none() && params.bulk_batches.is_none(),
        "anchor_batches / bulk_batches are not supported on the stack path — nothing \
         produces a carried reference or bulk input for stacked modalities"
    );
    let ctx = RefineCollectCtx {
        fine_codes: &fine_codes,
        group_to_cols_finest: &group_to_cols,
        level_dims: &level_dims,
        num_features,
        num_batches,
        knn,
        opt_iter,
        refine_params,
        output_calibration: params.output_calibration,
        anchor_batches: None,
        summary_batches: None,
        bulk_batches: None,
        observe_panels: false,
        keep_finest_stats: params.keep_finest_stats,
        pb_tree: None,
        cell_to_stratum: params.strata.as_deref(),
        exclude_unmatched_from_delta: params.strata.is_some(),
        strata_bits,
    };
    refine_and_collect_stack(stack, proj_kn, &ctx)
}