data-beans 0.12.0

//! Cell-axis quality control (library-only): robust (MAD) outlier detection
//! plus the near-empty floor. Built on the streaming stat collectors in the
//! sibling `qc` module and consumed by senna / pinto — NOT by the data-beans
//! CLI — so it lives in its own module that the binary never compiles (keeps
//! `qc` fully bin-used and avoids dead-code in the bin target).

use crate::qc::{collect_column_stat_across_vec, masked_column_stat};
use crate::sparse_data_visitors::VisitColumnsOps;
use crate::sparse_io_stack::SparseIoStack;
use crate::sparse_io_vector::SparseIoVec;
use legume_numeric::matrix::common_io::write_lines;
use legume_numeric::matrix::sparse_stat::{SparseColumnRunningStatistics, SparseRunningStatistics};
use legume_numeric::matrix::traits::RunningStatOps;
use log::warn;
use regex::Regex;
use std::sync::{Arc, Mutex};

///////////////////////////////////////////////////////////////////////////////////
// Cell-axis quality control: robust (MAD) outlier detection + near-empty floor. //
///////////////////////////////////////////////////////////////////////////////////
//
// This MAY DROP CELLS (columns): callers that enable QC end up with fewer
// cells in the working set and/or in the per-cell outputs than the input had,
// so downstream consumers must not assume a 1:1 positional mapping to the
// input barcodes — join by cell name. (Feature/row QC, off by default, may
// likewise drop genes/rows.)
//
// Shared by senna and pinto. Two-tier policy (see `compute_qc`):
//   * near-empty cells (`nnz < min_cell_nnz`) are kept in training but
//     dropped from the *output*;
//   * non-near-empty MAD outliers are dropped from *training* via
//     `SparseIoVec::mask_columns` (so they leave the outputs too).

/// Which side(s) of the robust band count as outliers.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Tail {
    /// Flag values far *below* the median (e.g. low counts / few genes).
    Lower,
    /// Flag values far *above* the median (e.g. high mito fraction).
    Upper,
    /// Flag both extremes.
    Both,
}

/// Configuration for cell QC. Plain struct (no clap) so non-clap crates
/// can construct it directly; the clap surface is [`QcArgs`].
#[derive(Clone, Debug)]
pub struct QcConfig {
    /// MAD multiplier for the robust band (larger = more permissive).
    pub n_mads: f32,
    /// Near-empty floor: cells with fewer than this many detected features
    /// (nnz across all rows) are masked at output, not dropped from
    /// training. 0 disables the floor.
    pub min_cell_nnz: usize,
    /// Hard floor on total counts per cell, in addition to MAD. 0 disables.
    pub min_counts_per_cell: f32,
    /// Regex over row (feature) names selecting mitochondrial genes; enables
    /// the per-cell mito-fraction metric. `None` disables it.
    pub mito_pattern: Option<String>,
    /// Optional hard max mitochondrial fraction (0..1).
    pub mito_max_frac: Option<f32>,
    /// Regex selecting ribosomal genes (enables ribo-fraction metric).
    pub ribo_pattern: Option<String>,
    /// Optional hard max ribosomal fraction (0..1).
    pub ribo_max_frac: Option<f32>,
    /// Per-cell metrics that drive MAD outlier flagging.
    pub mad_on_n_genes: bool,
    pub mad_on_counts: bool,
    pub mad_on_mito: bool,
    /// Feature-axis QC (off by default): drop genes expressed in fewer than
    /// this many cells. 0 disables.
    pub feature_min_cells: usize,
    /// Master switch for the MAD train-drop tier. When `false`, only the
    /// near-empty floor (output mask) is computed — used by inference
    /// (predict/impute) so query cells are never silently dropped.
    pub drop_outliers: bool,
    /// Automatic cell calling: pick the per-cell nnz cutoff at the trough
    /// between the ambient and the cell peak ([`crate::qc::suggest_nnz_cutoff`])
    /// and **train-drop** every cell below it. So called-out ambient is removed
    /// up front (via the caller's `mask_columns`), not just output-masked — it
    /// never shapes the model. The cutoff is authoritative; `min_cell_nnz` does
    /// NOT floor it (it only drives the separate near-empty output mask). No-op
    /// when there is no trough (unimodal) or when `drop_outliers` is false
    /// (inference). Unlike the loader's per-file empty-barcode gate, this runs
    /// on the pooled cell axis.
    pub auto_cell_cutoff: bool,
    /// Print the per-cell nnz histogram + the suggested/applied cutoff (the
    /// same ASCII summary as `data-beans squeeze --show-histogram`).
    pub qc_histogram: bool,
}

impl Default for QcConfig {
    fn default() -> Self {
        Self {
            n_mads: 5.0,
            min_cell_nnz: 2,
            min_counts_per_cell: 0.0,
            mito_pattern: None,
            mito_max_frac: None,
            ribo_pattern: None,
            ribo_max_frac: None,
            mad_on_n_genes: true,
            mad_on_counts: true,
            mad_on_mito: false,
            feature_min_cells: 0,
            drop_outliers: true,
            auto_cell_cutoff: false,
            qc_histogram: false,
        }
    }
}

/// Median of a non-empty slice via `O(n)` quickselect (reorders in place).
fn median_in_place(xs: &mut [f32]) -> f32 {
    let mid = xs.len() / 2;
    xs.select_nth_unstable_by(mid, |a, b| a.total_cmp(b));
    xs[mid]
}

/// The robust band [`robust_outlier_keep`] tests against: `median ± n_mads ·
/// MAD · 1.4826`, on the `log1p` scale when asked.
#[derive(Clone, Copy, Debug)]
pub struct RobustBand {
    pub median: f32,
    pub lo: f32,
    pub hi: f32,
}

impl RobustBand {
    /// The band on the values' own scale, undoing `log1p`.
    pub fn raw(self, log1p: bool) -> Self {
        let back = |v: f32| if log1p { v.exp_m1() } else { v };
        Self {
            median: back(self.median),
            lo: back(self.lo),
            hi: back(self.hi),
        }
    }
}

/// A value on the scale its band is fitted on: `ln(1 + v)` for counts.
fn band_scale(v: f32, log1p: bool) -> f32 {
    if log1p {
        v.max(0.0).ln_1p()
    } else {
        v
    }
}

/// [`RobustBand`] of `values` (the `consider`ed ones, when given); `None`
/// when none is finite, or when they have no spread (a MAD of 0).
pub fn robust_band(
    values: &[f32],
    n_mads: f32,
    log1p: bool,
    consider: Option<&[bool]>,
) -> Option<RobustBand> {
    let xform = |v: f32| band_scale(v, log1p);

    // Band statistics computed over the considered subset only.
    let mut xs: Vec<f32> = values
        .iter()
        .enumerate()
        .filter(|(i, _)| consider.is_none_or(|c| c[*i]))
        .map(|(_, &v)| xform(v))
        .filter(|v| v.is_finite())
        .collect();
    if xs.is_empty() {
        return None;
    }
    // O(n) median via quickselect (no full sort).
    let median = median_in_place(&mut xs);

    let mut dev: Vec<f32> = xs.iter().map(|&v| (v - median).abs()).collect();
    let mad = median_in_place(&mut dev) * 1.4826;
    // No spread: more than half the values are one value (a mito fraction
    // of 0 in most nuclei, say). A band of width 0 would call every other
    // value an outlier, so there is none.
    if mad <= 0.0 {
        return None;
    }

    Some(RobustBand {
        median,
        lo: median - n_mads * mad,
        hi: median + n_mads * mad,
    })
}

/// Robust outlier keep-mask: `keep[i] = false` when `values[i]` falls
/// outside `median ± n_mads · MAD · 1.4826` on the requested `tail`.
/// Count-like metrics should pass `log1p = true` so the band is symmetric
/// on the multiplicative scale. Consolidates the median+MAD idiom used in
/// `cnv::hmm`.
/// `consider` (when `Some`) restricts the median/MAD band to the cells it
/// marks `true` — so e.g. near-empty cells don't contaminate the robust
/// center — while the keep decision is still returned for every cell.
pub fn robust_outlier_keep(
    values: &[f32],
    n_mads: f32,
    tail: Tail,
    log1p: bool,
    consider: Option<&[bool]>,
) -> Vec<bool> {
    let n = values.len();
    let Some(band) = robust_band(values, n_mads, log1p, consider) else {
        return vec![true; n];
    };
    let xform = |v: f32| band_scale(v, log1p);
    let (lo, hi) = (band.lo, band.hi);
    values
        .iter()
        .map(|&raw| {
            let v = xform(raw);
            if !v.is_finite() {
                return false;
            }
            match tail {
                Tail::Lower => v >= lo,
                Tail::Upper => v <= hi,
                Tail::Both => v >= lo && v <= hi,
            }
        })
        .collect()
}

/// Resolve a regex over row (feature) names into matching row indices, in
/// the same compact-row order as `SparseIoVec::row_names()`. Feeds
/// `collect_column_stat_across_vec(.., Some(&rows), ..)` for subset sums.
pub fn resolve_rows_by_regex(row_names: &[Box<str>], pattern: &str) -> anyhow::Result<Vec<usize>> {
    let re = Regex::new(pattern)?;
    Ok(row_names
        .iter()
        .enumerate()
        .filter(|(_, name)| re.is_match(name))
        .map(|(i, _)| i)
        .collect())
}

/// Per-cell / per-feature QC outcome.
pub struct QcReport {
    /// Cells to retain in training (len = num_columns). `false` only for
    /// non-near-empty MAD outliers — drives `SparseIoVec::mask_columns`.
    pub train_keep: Vec<bool>,
    /// Near-empty cells (len = num_columns): kept in training, masked at
    /// output via [`QcReport::output_keep_idx`].
    pub near_empty: Vec<bool>,
    /// Features to retain (len = num_rows). All `true` unless
    /// `feature_min_cells > 0`.
    pub feature_keep: Vec<bool>,
    pub n_genes: Vec<f32>,
    pub total_counts: Vec<f32>,
    pub mito_frac: Option<Vec<f32>>,
    pub ribo_frac: Option<Vec<f32>>,
    pub n_cells_dropped: usize,
    pub n_features_dropped: usize,
}

impl QcReport {
    /// Indices, in *post-`mask_columns`* column order, of the cells to emit
    /// at output: surviving (training-kept) cells that are not near-empty.
    /// Apply with `Mat::select_rows` on per-cell output matrices and filter
    /// the barcode slice by the same indices.
    pub fn output_keep_idx(&self) -> Vec<usize> {
        let mut idx = Vec::new();
        let mut new_pos = 0usize;
        for c in 0..self.train_keep.len() {
            if !self.train_keep[c] {
                continue; // dropped from training entirely
            }
            if !self.near_empty[c] {
                idx.push(new_pos); // c survives at compact index new_pos
            }
            new_pos += 1;
        }
        idx
    }

    /// Emit-keep indices in the **original** (unmasked) column order: cells
    /// that are neither MAD-dropped nor near-empty. For paths that keep every
    /// cell in training and instead `select_rows` the per-cell outputs
    /// directly (e.g. bge's `UnifiedData`, which has no `mask_columns`).
    pub fn emit_idx_unmasked(&self) -> Vec<usize> {
        (0..self.train_keep.len())
            .filter(|&c| self.train_keep[c] && !self.near_empty[c])
            .collect()
    }
}

/// Compute the two-tier cell-QC report. `block_size` controls the
/// streaming stat passes (`None` = default chunking).
pub fn compute_qc(
    data: &SparseIoVec,
    cfg: &QcConfig,
    block_size: Option<usize>,
) -> anyhow::Result<QcReport> {
    compute_qc_exempting(data, cfg, block_size, None)
}

/// [`compute_qc`] with columns that are **not cells** exempted — a prior
/// run's carried pseudobulks. `exempt[c] = true` keeps column `c` out of all
/// band statistics and out of every verdict. See `qc_from_metrics`.
pub fn compute_qc_exempting(
    data: &SparseIoVec,
    cfg: &QcConfig,
    block_size: Option<usize>,
    exempt: Option<&[bool]>,
) -> anyhow::Result<QcReport> {
    let m = compute_qc_metrics(data, cfg, block_size, cfg.feature_min_cells > 0)?;
    Ok(qc_from_metrics(m, cfg, exempt))
}

/// The per-cell (and, with `feature_cells`, per-feature) metrics QC
/// decides on, measured once: detected features and total counts per
/// cell, the mito and ribo fractions where `cfg` names their patterns, and
/// how many cells detect each feature. [`cell_verdicts`] decides on them
/// under any config with those patterns, without reading the data again.
pub fn compute_qc_metrics(
    data: &SparseIoVec,
    cfg: &QcConfig,
    block_size: Option<usize>,
    feature_cells: bool,
) -> anyhow::Result<QcMetrics> {
    let row_names = data.row_names()?;

    // Optional mito / ribo row subsets via row-name regex.
    let rows_for = |pattern: &Option<String>| -> anyhow::Result<Option<Vec<usize>>> {
        let Some(pat) = pattern else {
            return Ok(None);
        };
        let rows = resolve_rows_by_regex(&row_names, pat)?;
        if rows.is_empty() {
            warn!(
                "QC: pattern `{}` matched no features — metric disabled",
                pat
            );
            return Ok(None);
        }
        Ok(Some(rows))
    };
    let mito_rows = rows_for(&cfg.mito_pattern)?;
    let ribo_rows = rows_for(&cfg.ribo_pattern)?;
    let subsets: Vec<&[usize]> = [&mito_rows, &ribo_rows]
        .into_iter()
        .filter_map(|rows| rows.as_deref())
        .collect();

    // Every metric from one pass over the data: per-cell totals across all
    // rows (nnz = n_genes, tot = total_counts), the subset totals, and the
    // feature axis stat (only when asked). Use the running-stat accessors
    // directly to skip the unused mean/std.
    let stat = collect_qc_stat(data, &subsets, feature_cells, block_size)?;
    let n_genes = stat.columns.count_positives();
    let total_counts = stat.columns.sum();

    // Fractions of the subsets measured, in the order they were asked for.
    let mut subset_totals = stat.subsets.iter().map(|s| s.sum());
    let mut frac_for = |rows: &Option<Vec<usize>>| -> Option<Vec<f32>> {
        rows.as_ref()?;
        let sub_tot = subset_totals.next()?;
        Some(
            sub_tot
                .iter()
                .zip(total_counts.iter())
                .map(|(&s, &t)| if t > 0.0 { s / t } else { 0.0 })
                .collect(),
        )
    };
    let mito_frac = frac_for(&mito_rows);
    let ribo_frac = frac_for(&ribo_rows);

    let feature_n_cells = stat.rows.map(|rows| rows.count_positives());

    Ok(QcMetrics {
        n_genes,
        total_counts,
        mito_frac,
        ribo_frac,
        feature_n_cells,
        n_rows: data.num_rows(),
    })
}

/// What [`collect_qc_stat`] collects in its one pass.
struct QcStat {
    /// Column statistics over all rows.
    columns: SparseColumnRunningStatistics<f32>,
    /// Column statistics over each row subset asked for, in that order.
    subsets: Vec<SparseColumnRunningStatistics<f32>>,
    /// Row statistics, when asked for.
    rows: Option<SparseRunningStatistics<f32>>,
}

/// Collect, reading every block once, the column statistics over all rows
/// (`collect_column_stat_across_vec` with no selection), those over each
/// of `row_subsets` (the same with each as its selection) and, with
/// `with_rows`, the row statistics (`collect_row_stat_across_vec`).
fn collect_qc_stat(
    data: &SparseIoVec,
    row_subsets: &[&[usize]],
    with_rows: bool,
    block_size: Option<usize>,
) -> anyhow::Result<QcStat> {
    let ncols = data.num_columns();
    let nrows = data.num_rows();

    let (masks, subsets): (Vec<Vec<bool>>, Vec<_>) = row_subsets
        .iter()
        .map(|sel| masked_column_stat(ncols, nrows, sel))
        .unzip();
    let mut stat = QcStat {
        columns: SparseColumnRunningStatistics::new(ncols, nrows),
        subsets,
        rows: with_rows.then(|| SparseRunningStatistics::new(nrows)),
    };
    data.visit_columns_by_block(&qc_stat_visitor, &masks, &mut stat, block_size)?;
    Ok(stat)
}

fn qc_stat_visitor(
    job: (usize, usize),
    data: &SparseIoVec,
    masks: &Vec<Vec<bool>>,
    arc_stat: Arc<Mutex<&mut QcStat>>,
) -> anyhow::Result<()> {
    let (lb, ub) = job;
    let csc = data.read_columns_csc(lb..ub)?;
    let mut stat = arc_stat.lock().expect("failed to lock qc stat");
    stat.columns.add_csc(&csc, lb);
    for (sub, mask) in stat.subsets.iter_mut().zip(masks) {
        sub.add_csc_masked(&csc, lb, mask);
    }
    if let Some(rows) = stat.rows.as_mut() {
        rows.add_csc(&csc);
    }
    Ok(())
}

/// Modality-agnostic cell QC for a [`SparseIoStack`]: per-cell `n_genes` /
/// `total_counts` are **summed across all member modalities** (a cell rich
/// in any one modality is kept). Mito/ribo and
/// feature-axis QC are skipped on stacks (row names are per-modality).
pub fn compute_qc_stack(
    stack: &SparseIoStack,
    cfg: &QcConfig,
    block_size: Option<usize>,
) -> anyhow::Result<QcReport> {
    let n_cols = stack.num_columns()?;
    let mut n_genes = vec![0f32; n_cols];
    let mut total_counts = vec![0f32; n_cols];
    for member in stack.stack.iter() {
        let cs = collect_column_stat_across_vec(member, None, block_size)?;
        let ng = cs.count_positives();
        let ct = cs.sum();
        for c in 0..n_cols {
            n_genes[c] += ng[c];
            total_counts[c] += ct[c];
        }
    }
    if cfg.mito_pattern.is_some() || cfg.ribo_pattern.is_some() || cfg.feature_min_cells > 0 {
        warn!("QC: mito/ribo/feature thresholds are ignored for stacked (multi-modal) data");
    }
    Ok(qc_from_metrics(
        QcMetrics {
            n_genes,
            total_counts,
            mito_frac: None,
            ribo_frac: None,
            feature_n_cells: None,
            n_rows: 0, // feature axis not masked on stacks
        },
        cfg,
        None,
    ))
}

/// Pre-computed per-cell / per-feature metrics fed to [`cell_verdicts`].
/// Lets the single-modality and stacked QC paths share one decision rule.
#[derive(Clone, Debug)]
pub struct QcMetrics {
    pub n_genes: Vec<f32>,
    pub total_counts: Vec<f32>,
    pub mito_frac: Option<Vec<f32>>,
    pub ribo_frac: Option<Vec<f32>>,
    /// Cells detecting each feature; `None` when not measured.
    pub feature_n_cells: Option<Vec<f32>>,
    pub n_rows: usize,
}

/// What [`cell_verdicts`] decides: the cells kept in training, the
/// near-empty ones, and the features kept.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Verdicts {
    pub train_keep: Vec<bool>,
    pub near_empty: Vec<bool>,
    pub feature_keep: Vec<bool>,
    pub n_cells_dropped: usize,
    pub n_features_dropped: usize,
    /// The rules would have dropped every cell, or every feature, and the
    /// guardrail kept them all.
    pub guarded: bool,
}

impl Verdicts {
    /// Whether cell `c` reaches the output: kept in training and not
    /// near-empty.
    pub fn emitted(&self, c: usize) -> bool {
        self.train_keep[c] && !self.near_empty[c]
    }
}

/// Apply the two-tier QC decision (near-empty floor + MAD outliers) to
/// pre-computed metrics. Shared by [`compute_qc`] and [`compute_qc_stack`].
///
/// `exempt`, when given, marks columns that are **not cells** — a prior run's
/// carried pseudobulks. They are excluded from every band statistic (a few
/// hundred smooth averages otherwise drag the MAD center and can guillotine
/// the real cells wholesale — measured: 400 of 400 real cells dropped as
/// "outliers" of the carried columns' band) and they receive no verdict:
/// never near-empty, never dropped. Their exclusion from OUTPUT is a separate
/// concern handled by the caller.
fn qc_from_metrics(m: QcMetrics, cfg: &QcConfig, exempt: Option<&[bool]>) -> QcReport {
    // The histogram + cutoff, optionally printed.
    if cfg.drop_outliers && (cfg.auto_cell_cutoff || cfg.qc_histogram) {
        // Of the cells alone, as the cutoff is found among them.
        let cells = cell_values(&m.n_genes, exempt);
        let suggested = crate::qc::suggest_nnz_cutoff(&cells);
        // Display cutoff: the trough suggestion if found, else the
        // near-empty floor (what the non-auto tier would use).
        let shown = suggested.unwrap_or(cfg.min_cell_nnz);
        crate::qc::eprint_nnz_summary("Cell", "nnz", &cells, shown, suggested);
    }
    let v = cell_verdicts(&m, cfg, exempt);
    let QcMetrics {
        n_genes,
        total_counts,
        mito_frac,
        ribo_frac,
        ..
    } = m;
    QcReport {
        train_keep: v.train_keep,
        near_empty: v.near_empty,
        feature_keep: v.feature_keep,
        n_genes,
        total_counts,
        mito_frac,
        ribo_frac,
        n_cells_dropped: v.n_cells_dropped,
        n_features_dropped: v.n_features_dropped,
    }
}

/// `values` of the columns that are cells: all of them, or those `exempt`
/// does not mark.
fn cell_values<'a>(values: &'a [f32], exempt: Option<&[bool]>) -> std::borrow::Cow<'a, [f32]> {
    match exempt {
        Some(e) => values
            .iter()
            .zip(e)
            .filter(|&(_, &x)| !x)
            .map(|(&v, _)| v)
            .collect(),
        None => values.into(),
    }
}

/// The [`QcReport`] of metrics measured once ([`compute_qc_metrics`]),
/// as [`compute_qc`] makes it: the histogram printed when `cfg` asks.
pub fn qc_report(m: QcMetrics, cfg: &QcConfig) -> QcReport {
    qc_from_metrics(m, cfg, None)
}

/// The two-tier QC decision (near-empty floor + MAD outliers) on `m`
/// under `cfg`, printing nothing: what [`compute_qc`] decides, for a view
/// that decides again as its settings change. `exempt` as in
/// [`compute_qc_exempting`].
pub fn cell_verdicts(m: &QcMetrics, cfg: &QcConfig, exempt: Option<&[bool]>) -> Verdicts {
    verdicts(m, cfg, exempt, true)
}

/// [`cell_verdicts`] without the guardrails' warnings, for a view that
/// decides again at every key and shows [`Verdicts::guarded`] itself.
pub(crate) fn cell_verdicts_quiet(m: &QcMetrics, cfg: &QcConfig) -> Verdicts {
    verdicts(m, cfg, None, false)
}

fn verdicts(m: &QcMetrics, cfg: &QcConfig, exempt: Option<&[bool]>, warn: bool) -> Verdicts {
    let QcMetrics {
        n_genes,
        total_counts,
        mito_frac,
        ribo_frac,
        feature_n_cells,
        n_rows,
    } = m;
    let n_rows = *n_rows;
    let n_cols = n_genes.len();

    let is_exempt = |c: usize| exempt.is_some_and(|e| e[c]);

    // Tier 1: near-empty floor. Exempt columns are never near-empty.
    let near_empty: Vec<bool> = n_genes
        .iter()
        .enumerate()
        .map(|(c, &g)| !is_exempt(c) && (g as usize) < cfg.min_cell_nnz)
        .collect();

    // Tier 2: MAD outliers among non-near-empty cells. The band statistics
    // are fit over the non-near-empty, non-exempt cells only, so neither
    // near-empty cells nor carried pseudobulks contaminate the robust center.
    let not_near_empty: Vec<bool> = near_empty
        .iter()
        .enumerate()
        .map(|(c, &e)| !e && !is_exempt(c))
        .collect();
    let consider = Some(not_near_empty.as_slice());
    let mut outlier = vec![false; n_cols];
    if cfg.drop_outliers {
        let mut bands: Vec<Vec<bool>> = Vec::new();
        if cfg.mad_on_n_genes {
            bands.push(robust_outlier_keep(
                n_genes,
                cfg.n_mads,
                Tail::Lower,
                true,
                consider,
            ));
        }
        if cfg.mad_on_counts {
            bands.push(robust_outlier_keep(
                total_counts,
                cfg.n_mads,
                Tail::Lower,
                true,
                consider,
            ));
        }
        if cfg.mad_on_mito {
            if let Some(mf) = mito_frac.as_ref() {
                bands.push(robust_outlier_keep(
                    mf,
                    cfg.n_mads,
                    Tail::Upper,
                    false,
                    consider,
                ));
            }
        }
        for c in 0..n_cols {
            if near_empty[c] || is_exempt(c) {
                continue; // floor takes precedence; exempt columns get no verdict
            }
            let mut fail = total_counts[c] < cfg.min_counts_per_cell;
            if let (Some(mf), Some(cap)) = (mito_frac.as_ref(), cfg.mito_max_frac) {
                fail |= mf[c] > cap;
            }
            if let (Some(rf), Some(cap)) = (ribo_frac.as_ref(), cfg.ribo_max_frac) {
                fail |= rf[c] > cap;
            }
            for band in bands.iter() {
                if !band[c] {
                    fail = true;
                    break;
                }
            }
            outlier[c] = fail;
        }

        // Automatic cell calling: the trough of n_genes picks the ambient↔real
        // boundary and train-drops every cell below it, so the caller's
        // `mask_columns(train_keep)` removes ambient up front. The cutoff
        // is authoritative — there is NO redundant `min_cell_nnz` floor on it.
        // When the data is unimodal the trough search returns `None`, so no auto
        // cutoff is applied and the near-empty floor / MAD tiers stand alone.
        // The trough is sought among the cells: exempt columns neither
        // shape it nor fall to it.
        if cfg.auto_cell_cutoff {
            if let Some(cut) = crate::qc::suggest_nnz_cutoff(&cell_values(n_genes, exempt)) {
                for (c, &g) in n_genes.iter().enumerate() {
                    if !is_exempt(c) && (g as usize) < cut {
                        outlier[c] = true;
                    }
                }
            }
        }
    }

    let mut train_keep: Vec<bool> = outlier.iter().map(|&o| !o).collect();
    let mut n_cells_dropped = outlier.iter().filter(|&&o| o).count();

    // Guardrail: never produce a matrix with no cell. Exempt columns are
    // never dropped, so the drops are held against the cells alone.
    let n_cells = (0..n_cols).filter(|&c| !is_exempt(c)).count();
    let mut guarded = false;
    if n_cells > 0 && n_cells_dropped >= n_cells {
        if warn {
            warn!(
                "QC would drop all {} cells — keeping all (check thresholds)",
                n_cells
            );
        }
        train_keep = vec![true; n_cols];
        n_cells_dropped = 0;
        guarded = true;
    }
    // Cells kept in training reach the output only above the near-empty
    // floor; where none does, the verdicts stand and the output is empty.
    let emits = |c: usize| !is_exempt(c) && train_keep[c] && !near_empty[c];
    if warn && n_cells > 0 && !(0..n_cols).any(emits) {
        warn!(
            "QC leaves none of the {} cells in the output: each is near-empty \
             (fewer than {} detected features) or dropped (check thresholds)",
            n_cells, cfg.min_cell_nnz
        );
    }

    // Feature axis (only when feature_n_cells was computed).
    let (feature_keep, n_features_dropped) = match feature_n_cells {
        Some(n_cells_expr) => {
            let keep: Vec<bool> = n_cells_expr
                .iter()
                .map(|&c| (c as usize) >= cfg.feature_min_cells)
                .collect();
            let dropped = keep.iter().filter(|&&k| !k).count();
            if !keep.is_empty() && dropped >= keep.len() {
                if warn {
                    warn!("QC would drop all features — keeping all");
                }
                guarded = true;
                (vec![true; keep.len()], 0)
            } else {
                (keep, dropped)
            }
        }
        None => (vec![true; n_rows], 0),
    };

    Verdicts {
        train_keep,
        near_empty,
        feature_keep,
        n_cells_dropped,
        n_features_dropped,
        guarded,
    }
}

/// Filter a per-cell `Vec<T>` in lockstep with a cell keep-mask. Used to
/// keep batch labels / coordinates aligned after `mask_columns`.
pub fn filter_by_keep<T: Clone>(items: &[T], keep: &[bool]) -> Vec<T> {
    items
        .iter()
        .zip(keep.iter())
        .filter(|&(_, &k)| k)
        .map(|(x, _)| x.clone())
        .collect()
}

/// Write a per-cell QC table (TSV): name, n_genes, total_counts,
/// [mito_frac], [ribo_frac], near_empty (0/1), train_keep (0/1).
pub fn write_qc_report(
    path: &str,
    cell_names: &[Box<str>],
    report: &QcReport,
) -> anyhow::Result<()> {
    use std::fmt::Write as _;
    let n = report.train_keep.len();
    anyhow::ensure!(
        cell_names.len() == n,
        "write_qc_report: {} names != {} cells",
        cell_names.len(),
        n
    );

    let mut header = String::from("#cell\tn_genes\ttotal_counts");
    if report.mito_frac.is_some() {
        header.push_str("\tmito_frac");
    }
    if report.ribo_frac.is_some() {
        header.push_str("\tribo_frac");
    }
    header.push_str("\tnear_empty\ttrain_keep");

    let mut lines: Vec<Box<str>> = Vec::with_capacity(n + 1);
    lines.push(header.into_boxed_str());
    for c in 0..n {
        let mut line = String::new();
        let _ = write!(
            line,
            "{}\t{}\t{}",
            cell_names[c], report.n_genes[c], report.total_counts[c]
        );
        if let Some(mf) = report.mito_frac.as_ref() {
            let _ = write!(line, "\t{}", mf[c]);
        }
        if let Some(rf) = report.ribo_frac.as_ref() {
            let _ = write!(line, "\t{}", rf[c]);
        }
        let _ = write!(
            line,
            "\t{}\t{}",
            report.near_empty[c] as u8, report.train_keep[c] as u8
        );
        lines.push(line.into_boxed_str());
    }
    write_lines(&lines, path)
}

/// Clap surface for cell QC, shared by senna and pinto subcommands.
///
/// **Cell QC is ON by default and CAN DROP CELLS (columns).** Low-quality
/// cells are removed: near-empty cells are omitted from the per-cell outputs,
/// and robust (MAD) outlier cells are excluded from training entirely. As a
/// result, the per-cell output files (`*.latent.parquet`, `*.cell_proj.parquet`,
/// `*.cell_to_pb.parquet`, `*.cell_embedding.parquet`, pinto propensity, etc.)
/// **may contain fewer rows (cells) than the input** — do not assume a 1:1,
/// positional correspondence with the input barcodes; always join by the cell
/// name/barcode column. Pass `--no-qc` to disable and keep every cell, or
/// `--qc-report <path>` to dump the per-cell keep/drop flags. (Feature/row QC
/// is OFF unless `--qc-feature-min-cells` is set.)
#[derive(clap::Args, Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(default = "legume_numeric::matrix::clap_defaults::clap_defaults")]
pub struct QcArgs {
    /// Disable cell QC entirely (keep every input cell).
    #[arg(
        long = "no-qc",
        default_value_t = false,
        long_help = "Disable cell quality control entirely and keep every input cell.\n\
                     \n\
                     By DEFAULT (without this flag) cell QC is:\n  \
                     - a near-empty nnz floor (--qc-min-cell-nnz), and\n  \
                     - MAD-outlier drops on detected features and total counts\n    \
                     (--qc-mad-on-genes / --qc-mad-on-counts, band --qc-mads).\n\
                     \n\
                     Empty barcodes are dropped earlier, per input file, by the loader.\n\
                     A pooled trough cut on top of that stays OFF unless --qc-auto-cutoff.\n\
                     \n\
                     Outputs may therefore have FEWER ROWS than the input.\n\
                     Join by the cell/barcode name column, never by position.\n\
                     \n\
                     Use --qc-report to see exactly what was dropped.\n\
                     For the older near-empty-floor-only gate,\n\
                     pass `--qc-mad-on-genes=false --qc-mad-on-counts=false`."
    )]
    pub no_qc: bool,

    /// MAD multiplier for the robust outlier band; smaller = drops more cells.
    #[arg(long = "qc-mads", default_value_t = 5.0)]
    pub qc_mads: f32,

    #[arg(
        long = "qc-min-cell-nnz",
        default_value_t = 2,
        help = "Near-empty floor on a cell's detected-feature count",
        long_help = "Near-empty floor on the detected-feature count.\n\
                     Cells below it are dropped from the per-cell outputs.\n\
                     They are still kept in training."
    )]
    pub qc_min_cell_nnz: usize,

    #[arg(
        long = "qc-min-counts",
        hide = true,
        default_value_t = 0.0,
        help = "Hard floor on total counts per cell",
        long_help = "Hard floor on total counts per cell.\n\
                     Cells below it are dropped from training. 0 disables the floor."
    )]
    pub qc_min_counts: f32,

    #[arg(
        long = "qc-mito-pattern",
        hide = true,
        help = "Regex over feature names selecting mitochondrial genes",
        long_help = "Regex over feature names selecting mitochondrial genes.\n\
                     It enables the mito-fraction outlier metric. An example is `(?i)^MT-`."
    )]
    pub qc_mito_pattern: Option<String>,

    /// Hard max mitochondrial fraction (0..1).
    #[arg(long = "qc-mito-max-frac", hide = true)]
    pub qc_mito_max_frac: Option<f32>,

    /// Regex over feature names selecting ribosomal genes.
    #[arg(long = "qc-ribo-pattern", hide = true)]
    pub qc_ribo_pattern: Option<String>,

    /// Hard max ribosomal fraction (0..1).
    #[arg(long = "qc-ribo-max-frac", hide = true)]
    pub qc_ribo_max_frac: Option<f32>,

    #[arg(
        long = "qc-feature-min-cells",
        hide = true,
        default_value_t = 0,
        help = "Feature/row QC: drop genes expressed in too few cells",
        long_help = "Feature/row QC; off by default.\n\
                     It DROPS gene rows expressed in fewer than this many cells.\n\
                     \n\
                     Not every consumer applies it.\n\
                     `bge` does not, since QC there is cell-only.\n\
                     `pinto` does not either: it reads only the cell verdict,\n\
                     so setting this costs a stats pass and changes nothing."
    )]
    pub qc_feature_min_cells: usize,

    #[arg(
        long = "qc-report",
        help = "Write a per-cell QC table (.tsv)",
        long_help = "Write a per-cell QC table, as .tsv.\n\
                     It carries the metrics plus near_empty and train_keep flags.\n\
                     You can then see exactly which cells were dropped."
    )]
    pub qc_report: Option<Box<str>>,

    #[arg(
        long = "qc-histogram",
        hide = true,
        default_value_t = false,
        help = "Print the per-cell nnz histogram + the (diagnostic) suggested trough cutoff",
        long_help = "Print an ASCII histogram of the per-cell nnz distribution.\n\
                     The suggested trough cutoff is marked.\n\
                     It is the same summary as `data-beans squeeze --show-histogram`.\n\
                     \n\
                     This is purely diagnostic. The cutoff is shown, not applied.\n\
                     The upfront gate is the conservative --qc-min-cell-nnz floor.\n\
                     Use the histogram to pick --qc-min-cell-nnz by hand."
    )]
    pub qc_histogram: bool,

    #[arg(
        long = "qc-mad-on-genes",
        hide = true,
        // `--qc-mad-on-genes=false` turns it off; a bare flag leaves it on.
        action = clap::ArgAction::Set,
        num_args = 0..=1,
        require_equals = true,
        default_missing_value = "true",
        default_value_t = true,
        help = "MAD-outlier drop on the per-cell detected-feature count",
        long_help = "Drop cells whose detected-feature count falls outside `median +/- --qc-mads * MAD * 1.4826`.\n\
                     \n\
                     ON by default. This and --qc-mad-on-counts were previously hardcoded OFF,\n\
                     with no way to enable them. That also made --qc-mads inert.\n\
                     Only a set --qc-mito-pattern revived it.\n\
                     \n\
                     Pass `--qc-mad-on-genes=false` for the old behaviour.\n\
                     That is the conservative near-empty nnz gate alone."
    )]
    pub qc_mad_on_genes: bool,

    #[arg(
        long = "qc-mad-on-counts",
        hide = true,
        // `--qc-mad-on-counts=false` turns it off; a bare flag leaves it on.
        action = clap::ArgAction::Set,
        num_args = 0..=1,
        require_equals = true,
        default_missing_value = "true",
        default_value_t = true,
        help = "MAD-outlier drop on per-cell total counts",
        long_help = "Drop cells whose total count falls outside `median +/- --qc-mads * MAD * 1.4826`.\n\
                     ON by default; see --qc-mad-on-genes."
    )]
    pub qc_mad_on_counts: bool,

    #[arg(
        long = "qc-mad-on-mito",
        hide = true,
        // `--qc-mad-on-mito=false` turns it off; a bare flag leaves it on.
        action = clap::ArgAction::Set,
        num_args = 0..=1,
        require_equals = true,
        default_missing_value = "true",
        default_value_t = true,
        help = "MAD-outlier drop on the mito fraction, given --qc-mito-pattern",
        long_help = "Drop cells whose mito fraction is above `median + --qc-mads * MAD * 1.4826`.\n\
                     On whenever --qc-mito-pattern is given. `--qc-mad-on-mito=false` leaves\n\
                     only the hard cap, --qc-mito-max-frac, when one is given."
    )]
    pub qc_mad_on_mito: bool,

    #[arg(
        long = "qc-auto-cutoff",
        hide = true,
        default_value_t = false,
        help = "Apply the nnz trough cell-calling cutoff on the pooled cell axis",
        long_help = "Apply the ambient/cell trough cutoff as a hard cell call.\n\
                     It runs on the pooled per-cell nnz distribution, after the loader's per-file gate.\n\
                     Without this flag it is only reported, via --qc-histogram.\n\
                     \n\
                     OFF by default. The intended gate is the conservative near-empty floor,\n\
                     plus the model's own empty-call.\n\
                     This flag was referenced in the docs before it existed."
    )]
    pub qc_auto_cutoff: bool,
}

impl QcArgs {
    /// QC config, or `None` under `--no-qc`.
    pub fn to_config(&self) -> Option<QcConfig> {
        (!self.no_qc).then(|| QcConfig {
            n_mads: self.qc_mads,
            min_cell_nnz: self.qc_min_cell_nnz,
            min_counts_per_cell: self.qc_min_counts,
            // An empty pattern names no genes: the metric is off, not a
            // regex that finds every gene.
            mito_pattern: self.qc_mito_pattern.clone().filter(|p| !p.is_empty()),
            mito_max_frac: self.qc_mito_max_frac,
            ribo_pattern: self.qc_ribo_pattern.clone().filter(|p| !p.is_empty()),
            ribo_max_frac: self.qc_ribo_max_frac,
            // MAD gates are ON by default and settable. They were hardcoded
            // `false` here with no CLI path, which also made `--qc-mads` a dead
            // flag unless `--qc-mito-pattern` was set (nothing else reads
            // `n_mads`). Mito stays implicit: it can only run when a pattern
            // selects the rows to measure.
            mad_on_n_genes: self.qc_mad_on_genes,
            mad_on_counts: self.qc_mad_on_counts,
            mad_on_mito: self
                .qc_mito_pattern
                .as_deref()
                .is_some_and(|p| !p.is_empty())
                && self.qc_mad_on_mito,
            feature_min_cells: self.qc_feature_min_cells,
            drop_outliers: true,
            // Off by default: the near-empty floor plus the model's own
            // empty-call (bge's embedding-norm two-step; topic/masked-topic's
            // flag-don't-drop) is the intended gate. `--qc-auto-cutoff` opts in.
            auto_cell_cutoff: self.qc_auto_cutoff,
            qc_histogram: self.qc_histogram,
        })
    }
}

#[cfg(test)]
mod qc_tests {
    use super::*;

    #[test]
    fn robust_lower_flags_low_outlier() {
        let v = vec![100.0, 110.0, 90.0, 105.0, 95.0, 1.0];
        let keep = robust_outlier_keep(&v, 3.0, Tail::Lower, true, None);
        assert!(!keep[5], "the value 1.0 should be a lower outlier");
        assert!(keep[..5].iter().all(|&k| k), "the bulk should be kept");
        // upper tail must not flag a low value
        let keep_up = robust_outlier_keep(&v, 3.0, Tail::Upper, true, None);
        assert!(keep_up[5], "lower outlier kept under Tail::Upper");
    }

    #[test]
    fn no_spread_flags_no_one() {
        // Most cells read no mito at all: the MAD is 0, and a cell with one
        // mito read is no outlier for it.
        let mut v = vec![0.0f32; 60];
        v.extend([0.01, 0.02, 0.05, 0.3]);
        assert!(robust_band(&v, 5.0, false, None).is_none());
        let keep = robust_outlier_keep(&v, 5.0, Tail::Upper, false, None);
        assert!(keep.iter().all(|&k| k));
    }

    #[test]
    fn robust_uniform_keeps_all() {
        let v = vec![7.0; 20];
        let keep = robust_outlier_keep(&v, 5.0, Tail::Both, true, None);
        assert!(keep.iter().all(|&k| k));
    }

    #[test]
    fn auto_cutoff_train_drops_ambient() {
        // Bimodal nnz: 30 ambient (~2) + 30 real (~100). With auto cell calling on
        // (MAD tiers off so the auto floor is the only decider), the trough
        // cutoff lands between the modes and the ambient cells are train-dropped.
        let n_genes: Vec<f32> = [vec![2.0; 30], vec![100.0; 30]].concat();
        let cfg = QcConfig {
            auto_cell_cutoff: true,
            qc_histogram: false,
            drop_outliers: true,
            mad_on_n_genes: false,
            mad_on_counts: false,
            mad_on_mito: false,
            min_cell_nnz: 0,
            ..QcConfig::default()
        };
        let report = qc_from_metrics(
            QcMetrics {
                n_genes: n_genes.clone(),
                total_counts: n_genes,
                mito_frac: None,
                ribo_frac: None,
                feature_n_cells: None,
                n_rows: 0,
            },
            &cfg,
            None,
        );
        assert_eq!(
            report.train_keep,
            [vec![false; 30], vec![true; 30]].concat()
        );
        assert_eq!(report.n_cells_dropped, 30);
    }

    #[test]
    fn robust_consider_excludes_contaminants_from_band() {
        // A real cluster around 100 plus a large block of near-empty 0s that
        // would drag a naive median/MAD down. With `consider` masking the
        // zeros out of the band, a genuine low real cell (40) is still flagged.
        let mut v = vec![0.0; 12];
        v.extend([100.0, 102.0, 98.0, 101.0, 99.0, 40.0]);
        let mut consider = vec![false; 12];
        consider.extend([true; 6]);
        let keep = robust_outlier_keep(&v, 2.0, Tail::Lower, true, Some(&consider));
        assert!(
            !keep[17],
            "40 is a lower outlier of the real cluster (~100)"
        );
        // Without `consider`, the zeros dominate the median and 40 survives.
        let keep_naive = robust_outlier_keep(&v, 2.0, Tail::Lower, true, None);
        assert!(
            keep_naive[17],
            "naive band (contaminated by zeros) keeps 40"
        );
    }

    #[test]
    fn output_keep_idx_skips_dropped_and_near_empty() {
        // cells: 0 keep, 1 near-empty (kept in training), 2 MAD-drop, 3 keep
        let report = QcReport {
            train_keep: vec![true, true, false, true],
            near_empty: vec![false, true, false, false],
            feature_keep: vec![],
            n_genes: vec![],
            total_counts: vec![],
            mito_frac: None,
            ribo_frac: None,
            n_cells_dropped: 1,
            n_features_dropped: 0,
        };
        // post-mask order: cell0->0, cell1->1, cell3->2 (cell2 dropped)
        // emit non-near-empty survivors: 0 and 2
        assert_eq!(report.output_keep_idx(), vec![0, 2]);
    }
}

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