data-beans 0.6.13

use std::sync::Arc;

use log::{debug, info, warn};

use crate::convert::try_open_or_convert;
use crate::sparse_io_vector::SparseIoVec;
use legume_numeric::matrix::common_io::{self, basename, read_lines};
use rustc_hash::FxHashSet;

use crate::aux::feature_names::FeatureNameKind;
use crate::sparse_io_vector::{ColumnAlignment, RowAlignment};

/// When `RowAlignment::Disjoint` (the historical default) and multiple
/// files share most of their feature axis, the loader silently does the
/// horizontal-stack thing. When the axes are *mostly disjoint* AND raw
/// barcodes overlap, the inputs look multi-modal-shaped and the user
/// probably meant `--multiome`. This threshold (intersection size /
/// smallest per-backend row count) decides "mostly disjoint."
const MULTIMODAL_HINT_DISJOINTNESS_FRACTION: f64 = 0.5;

/// Arguments for loading multiple sparse data files with shared row names.
#[derive(Default)]
pub struct ReadSharedRowsArgs {
    pub data_files: Vec<Box<str>>,
    pub batch_files: Option<Vec<Box<str>>>,
    pub preload: bool,
    /// Cross-file row-name canonicalization rule. `None` = auto-detect
    /// via [`FeatureNameKind::auto_detect`] once row names are in hand.
    /// `Some(kind)` skips detection and uses the caller's choice.
    /// Default = `None` (auto).
    pub feature_kind: Option<FeatureNameKind>,
    /// How to align row names across input files. Default
    /// [`RowAlignment::Union`] keeps every row from any backend — the
    /// strictly more permissive option that reduces to single-modality
    /// semantics when all files share their row set, and supports
    /// multi-modal load (e.g. paired RNA + ATAC) when they don't.
    /// Switch to [`RowAlignment::Intersect`] for strict "common rows
    /// only" behavior.
    pub row_alignment: RowAlignment,
    /// How to align column (cell / barcode) names across input files.
    /// Default [`ColumnAlignment::Disjoint`] concatenates cells with
    /// `@<basename>` disambiguation, preserving single-modality
    /// semantics. Switch to [`ColumnAlignment::Union`] for **patchy
    /// multi-modal (multiome) load**: cells are glued by raw barcode
    /// across backends, a cell observed in only one modality
    /// contributes triplets only on that modality's row block, and no
    /// `@<basename>` suffix is added.
    pub column_alignment: ColumnAlignment,
    /// Optional shared cell QC. When `Some`, non-near-empty MAD outlier
    /// cells are dropped from the working set (via `mask_columns`), the
    /// returned `batch` Vec is filtered in lockstep, junk features are
    /// dropped (when `feature_min_cells > 0`), and the near-empty
    /// output keep-mask is returned in `output_keep_idx`. `None`
    /// (the default) = no QC, i.e. today's behavior.
    pub qc: Option<crate::qc_lib::QcConfig>,
    /// Opt out of the empty-barcode gate. By default each file's columns
    /// are cell-called on their own nnz distribution
    /// ([`crate::qc::suggest_nnz_cutoff`]: the trough between the ambient
    /// and the cell peak) and a column is dropped when it falls below the
    /// cut in every file that observes it — so an unfiltered barcode axis
    /// (e.g. ATAC from a fragments file) loses its empty droplets before any
    /// projection or collapse, while a cell in two modalities survives on
    /// either. A no-op on already-called data (no trough). Set `true` where
    /// every input column must come back, e.g. query cells at inference.
    /// Files flagged in `qc_exempt_files` are never gated.
    pub keep_empty_barcodes: bool,
    /// Per-`data_files` entry: `true` exempts that file's columns from QC —
    /// out of the band statistics AND out of every verdict (see
    /// `qc_from_metrics`). For inputs whose columns are not cells (a carried
    /// `pb_reference`, bulk samples): a pseudobulk standing for hundreds of
    /// cells is a legitimate depth outlier, and letting it into the MAD band
    /// either gets it dropped or — worse, as the mixture grows — recenters
    /// the band and guillotines the real cells. `None` = no exemption. Must
    /// match `data_files` length when `Some`.
    pub qc_exempt_files: Option<Vec<bool>>,
    /// Block size for the QC streaming stat passes (`None` = default).
    pub qc_block_size: Option<usize>,
    /// Optional path for a per-cell QC report TSV (`None` = don't write).
    pub qc_report_out: Option<Box<str>>,
    /// Optional per-file feature-name (row) modality suffix, one entry per
    /// `data_files` entry in order. When `Some`, backend `b`'s rows are
    /// renamed `{canon(row)}/{suffix[b]}` so files sharing raw feature
    /// names (e.g. spliced/unspliced) stay on separate rows, while the same
    /// name + suffix (same modality across samples) still merges. `None` =
    /// today's behavior (no suffixing). Must match `data_files` length.
    pub per_file_feature_suffix: Option<Vec<Box<str>>>,
    /// Optional per-file barcode (column) suffix, one entry per `data_files`
    /// entry in order. Under [`ColumnAlignment::Union`], backend `b`'s
    /// barcodes are tagged `{barcode}@{suffix[b]}` before the canonical merge,
    /// so callers can encode per-cell **sample identity**: the same barcode in
    /// two files merges into one cell only when both carry the same suffix
    /// (same sample across modalities), while different samples stay distinct.
    /// `None` (or a per-entry `None`) = no tag. Must match `data_files`
    /// length. No effect under `Disjoint` (which disambiguates via
    /// `@<basename>`).
    pub per_file_barcode_suffix: Option<Vec<Option<Box<str>>>>,
}

/// Sparse data with per-cell batch labels.
pub struct SparseDataWithBatch {
    pub data: SparseIoVec,
    pub batch: Vec<Box<str>>,
    /// Near-empty output keep-mask: indices (in post-`mask_columns`
    /// column order) of cells to emit at output. `None` when no QC ran
    /// (emit every cell). See [`crate::qc_lib::QcReport::output_keep_idx`].
    pub output_keep_idx: Option<Vec<usize>>,
}

/// Load multiple sparse data files, verify shared row names, and auto-detect batch.
///
/// Batch assignment priority:
/// 1. Explicit batch files (one label per cell per file)
/// 2. Embedded `@`-separated batch info in column names (e.g., `barcode@donor`)
/// 3. File name as batch label (one batch per input file)
pub fn read_data_on_shared_rows(args: ReadSharedRowsArgs) -> anyhow::Result<SparseDataWithBatch> {
    // to avoid duplicate barcodes in the column names
    let attach_data_name = args.data_files.len() > 1;

    // Open every backend first (preserving order). For LocusOverlap we
    // need to peek all row names before installing the canonicalizer.
    type OpenedBackend = Box<dyn crate::sparse_io::SparseIo<IndexIter = Vec<usize>>>;
    let mut opened: Vec<(Box<str>, OpenedBackend)> = Vec::with_capacity(args.data_files.len());
    for data_file in args.data_files.iter() {
        info!("Importing data file: {}", data_file);
        let mut data = try_open_or_convert(data_file)?;
        if args.preload {
            data.preload_columns()?;
        }
        opened.push((data_file.clone(), data));
    }

    // `Disjoint` keeps today's @<basename>-suffix semantics; `Union`
    // glues cells by raw barcode (no suffix). Affects `SparseIoVec::push`
    // *and* the batch-resolution branch below — per-file batch labels
    // become ambiguous under Union (one cell can be in two files) so we
    // require either a single unified batch file or per-cell embedded
    // tags that agree across backends.
    let attach_data_name = attach_data_name && args.column_alignment == ColumnAlignment::Disjoint;

    let mut data_vec = SparseIoVec::new()
        .with_row_alignment(args.row_alignment)
        .expect("with_row_alignment on empty SparseIoVec")
        .with_column_alignment(args.column_alignment)
        .expect("with_column_alignment on empty SparseIoVec");

    // Per-file feature-name modality suffix (e.g. `--multiome` modality
    // namespacing). Installed before any push so backend `b`'s rows become
    // `{canon(row)}/{suffix[b]}`.
    if let Some(suffix) = args.per_file_feature_suffix.clone() {
        anyhow::ensure!(
            suffix.len() == args.data_files.len(),
            "per_file_feature_suffix has {} entries but {} data files were given",
            suffix.len(),
            args.data_files.len(),
        );
        data_vec = data_vec
            .with_per_backend_row_suffix(suffix)
            .expect("with_per_backend_row_suffix on empty SparseIoVec");
    }

    // Per-file barcode (column) sample suffix — applied per push below under
    // Union. Validate length up front so a mismatch fails before any I/O.
    if let Some(sfx) = args.per_file_barcode_suffix.as_ref() {
        anyhow::ensure!(
            sfx.len() == args.data_files.len(),
            "per_file_barcode_suffix has {} entries but {} data files were given",
            sfx.len(),
            args.data_files.len(),
        );
    }

    use crate::aux::feature_names::FeatureNameKind;

    // Peek row names once if auto-detect needs them, or if the
    // caller-specified kind needs a global cross-name pass.
    let needs_names = args.feature_kind.is_none()
        || args
            .feature_kind
            .as_ref()
            .is_some_and(|k| k.needs_global_pass());
    // One flat list of every file's row names (the Mixed / locus-overlap
    // maps are built over all of them at once), plus where each file's
    // slice ends, so auto-detection can look at one file at a time without a
    // second read.
    let mut file_ends: Vec<usize> = Vec::with_capacity(opened.len());
    let all_names: Option<Vec<Box<str>>> = if needs_names {
        let mut acc: Vec<Box<str>> = Vec::new();
        for (_, d) in opened.iter() {
            acc.extend(d.row_names()?);
            file_ends.push(acc.len());
        }
        Some(acc)
    } else {
        None
    };

    let kind_was_auto = args.feature_kind.is_none();
    let resolved_kind: FeatureNameKind = match args.feature_kind.clone() {
        Some(k) => k,
        None => {
            // Per FILE, then reconciled — never over the pool. The naming
            // signature usually lives on one side only: a raw `ENSG_SYM`
            // cohort next to a reference already on the bare-symbol axis
            // (a carried `pb_reference`) left the pooled gene-like share
            // under half, so the pair sniffed as `Exact` and every gene
            // became two rows. See `FeatureNameKind::reconcile`.
            let names = all_names.as_ref().expect("peeked when auto");
            // Each file's slice of the flat list, from the cumulative ends: a
            // plain zip of starts and ends, so the slicing does not depend on
            // the closure running in order.
            let per_file: Vec<FeatureNameKind> = std::iter::once(0)
                .chain(file_ends.iter().copied())
                .zip(file_ends.iter().copied())
                .map(|(start, end)| FeatureNameKind::auto_detect(&names[start..end]))
                .collect();
            let k = FeatureNameKind::reconcile(&per_file);
            debug!(
                "Row alignment: auto-detected feature name kind → {:?} (per file: {:?}; {} rows)",
                k,
                per_file,
                names.len()
            );
            k
        }
    };

    // Install canonicalizer. Three paths:
    //   * Mixed                            → per-name dispatcher (needs names).
    //   * Locus { merge_overlapping: true } → overlap cluster map (needs names).
    //   * Anything else                    → pure per-name closure from `canonicalize`.
    match &resolved_kind {
        FeatureNameKind::Mixed => {
            let names = all_names.as_ref().expect("peeked for Mixed").clone();
            debug!(
                "Row alignment: building MIXED-kind canonical map over {} names \
                 across {} file(s)",
                names.len(),
                opened.len()
            );
            let canon = crate::aux::feature_names::build_mixed_kind_canonicalizer(&names);
            data_vec = data_vec
                .with_row_canonicalizer(move |name| canon(name))
                .expect("with_row_canonicalizer on empty SparseIoVec");
        }
        FeatureNameKind::Locus {
            merge_overlapping: true,
        } => {
            let names = all_names
                .as_ref()
                .expect("peeked for Locus merge_overlapping")
                .clone();
            debug!(
                "Row alignment: building locus-overlap canonical map over {} names \
                 across {} file(s)",
                names.len(),
                opened.len()
            );
            let canon = crate::aux::feature_names::build_locus_overlap_canonicalizer(&names);
            data_vec = data_vec
                .with_row_canonicalizer(move |name| canon(name))
                .expect("with_row_canonicalizer on empty SparseIoVec");
        }
        kind => {
            if let Some(canon) = kind.clone().into_canonicalizer() {
                debug!(
                    "Row alignment: applying {:?} canonicalizer across {} file(s)",
                    kind,
                    opened.len()
                );
                // SAFETY: data_vec is empty; with_row_canonicalizer only errors
                // if backends were already pushed.
                data_vec = data_vec
                    .with_row_canonicalizer(move |name| canon(name))
                    .expect("with_row_canonicalizer on empty SparseIoVec");
            }
        }
    }
    info!(
        "Row alignment: {:?} · {:?} canon{} · {} file(s)",
        args.row_alignment,
        resolved_kind,
        if kind_was_auto { " (auto)" } else { "" },
        opened.len(),
    );
    for (file_idx, (data_file, data)) in opened.into_iter().enumerate() {
        let data_name = attach_data_name.then(|| basename(&data_file)).transpose()?;
        // Per-file barcode sample tag (Union only; ignored under Disjoint).
        let barcode_suffix: Option<&str> = args
            .per_file_barcode_suffix
            .as_ref()
            .and_then(|v| v[file_idx].as_deref());
        data_vec.push_with_barcode_suffix(Arc::from(data), data_name, barcode_suffix)?;
    }

    // SparseIoVec already aligns rows to the intersection of row names
    // across all backends; warn if any backend introduced new rows that
    // had to be dropped.
    let intersection_size = data_vec.num_rows();
    for j in 0..data_vec.len() {
        let backend_rows = data_vec[j].num_rows().unwrap_or(0);
        if backend_rows != intersection_size {
            info!(
                "Backend {} has {} rows; using {} shared rows for fitting",
                j, backend_rows, intersection_size
            );
        }
    }

    // Soft guard: when running with the default Disjoint cell-stacking
    // and the feature axes are mostly disjoint AND raw barcodes overlap,
    // the inputs look multi-modal-shaped and the user probably meant
    // `--multiome`. Print a hint, don't error — a single-modality user
    // with quirky inputs should still get through.
    if args.column_alignment == ColumnAlignment::Disjoint && data_vec.len() >= 2 {
        maybe_warn_multimodal_pattern(&data_vec);
    }

    // check batch membership
    let n_cells = data_vec.num_columns();
    let mut batch_membership: Vec<Box<str>> = match args.column_alignment {
        ColumnAlignment::Disjoint => resolve_batch_disjoint(
            &args.data_files,
            &data_vec,
            args.batch_files.as_deref(),
            attach_data_name,
        )?,
        ColumnAlignment::Union => {
            resolve_batch_union(&data_vec, args.batch_files.as_deref(), n_cells)?
        }
    };

    if batch_membership.len() != data_vec.num_columns() {
        return Err(anyhow::anyhow!(
            "# batch membership {} != # of columns {}",
            batch_membership.len(),
            data_vec.num_columns()
        ));
    }

    // Empty-barcode gate: per-file cell calling, before QC and before any
    // batch/group registration (mask_columns renumbers the cells). Batch
    // labels were resolved on the full axis above (a unified batch file
    // lists every barcode), then filtered in lockstep.
    if !args.keep_empty_barcodes {
        if let Some(flags) = args.qc_exempt_files.as_ref() {
            anyhow::ensure!(
                flags.len() == args.data_files.len(),
                "qc_exempt_files has {} entries for {} data files",
                flags.len(),
                args.data_files.len(),
            );
        }
        if let Some(keep) = empty_barcode_keep(&data_vec, args.qc_exempt_files.as_deref()) {
            data_vec.mask_columns(&keep)?;
            batch_membership = crate::qc_lib::filter_by_keep(&batch_membership, &keep);
        }
    }

    // Optional shared cell QC — applied here (before any batch/group
    // registration, which happens later during projection) so all
    // downstream stages see the QC-reduced axes consistently.
    let output_keep_idx = if let Some(cfg) = args.qc.as_ref() {
        // Columns of exempt files stay out of the QC bands and verdicts;
        // resolved per column through the loaded vec's backend attribution so
        // it is correct under any stacking order.
        let mut exempt: Option<Vec<bool>> = None;
        if let Some(flags) = args.qc_exempt_files.as_ref() {
            anyhow::ensure!(
                flags.len() == args.data_files.len(),
                "qc_exempt_files has {} entries for {} data files",
                flags.len(),
                args.data_files.len(),
            );
            if flags.iter().any(|&f| f) {
                exempt = Some(
                    (0..data_vec.num_columns())
                        .map(|c| data_vec.column_source(c).is_some_and(|b| flags[b]))
                        .collect(),
                );
            }
        }
        let report = crate::qc_lib::compute_qc_exempting(
            &data_vec,
            cfg,
            args.qc_block_size,
            exempt.as_deref(),
        )?;
        if let Some(path) = args.qc_report_out.as_deref() {
            crate::qc_lib::write_qc_report(path, &data_vec.column_names()?, &report)?;
        }
        let n_near_empty = report.near_empty.iter().filter(|&&e| e).count();
        info!(
            "QC: dropped {}/{} cells from training, {} near-empty masked at output, {}/{} features dropped",
            report.n_cells_dropped,
            report.train_keep.len(),
            n_near_empty,
            report.n_features_dropped,
            report.feature_keep.len(),
        );
        // Feature axis first (compact-row space is still the original one).
        if report.n_features_dropped > 0 {
            data_vec.mask_rows(&report.feature_keep)?;
        }
        // Cell axis: indices computed against the original column order,
        // then mask_columns + lockstep batch filter.
        let keep_idx = report.output_keep_idx();
        if report.n_cells_dropped > 0 {
            data_vec.mask_columns(&report.train_keep)?;
            batch_membership = crate::qc_lib::filter_by_keep(&batch_membership, &report.train_keep);
        }
        Some(keep_idx)
    } else {
        None
    };

    Ok(SparseDataWithBatch {
        data: data_vec,
        batch: batch_membership,
        output_keep_idx,
    })
}

/// Keep-mask of the empty-barcode gate, or `None` when nothing is dropped.
///
/// Each backend is cell-called on its own column nnz (read off the resident
/// indptr, no I/O): modalities count on different scales, so they are never
/// pooled. A global column survives when any backend observing it passes
/// (or is exempt, or has no indptr / no trough to call on).
fn empty_barcode_keep(data_vec: &SparseIoVec, exempt: Option<&[bool]>) -> Option<Vec<bool>> {
    let mut missing_indptr: Vec<usize> = Vec::new();
    let cutoffs: Vec<Option<u64>> = (0..data_vec.len())
        .map(|b| {
            if exempt.is_some_and(|e| e[b]) {
                return None;
            }
            let backend = &data_vec[b];
            let ncol = backend.num_columns().unwrap_or(0);
            let nnz: Option<Vec<f32>> = (0..ncol)
                .map(|c| backend.column_nnz(c).map(|x| x as f32))
                .collect();
            let Some(nnz) = nnz else {
                missing_indptr.push(b);
                return None;
            };
            crate::qc::suggest_nnz_cutoff(&nnz).map(|c| c as u64)
        })
        .collect();
    if !missing_indptr.is_empty() {
        warn!(
            "Empty-barcode gate: file index(es) {} have no resident column indptr; \
             skipping cell call for those backends",
            missing_indptr
                .iter()
                .map(ToString::to_string)
                .collect::<Vec<_>>()
                .join(", "),
        );
    }
    if cutoffs.iter().all(Option::is_none) {
        return None;
    }

    let keep: Vec<bool> = (0..data_vec.num_columns())
        .map(|c| {
            data_vec.column_locations(c).iter().any(|loc| {
                let b = loc.backend as usize;
                cutoffs[b].is_none_or(|cut| {
                    data_vec[b]
                        .column_nnz(loc.local_col as usize)
                        .is_none_or(|x| x >= cut)
                })
            })
        })
        .collect();
    let n_drop = keep.iter().filter(|&&k| !k).count();
    info!(
        "Empty-barcode gate: {} / {} columns called empty (per-file nnz cutoffs: {})",
        n_drop,
        keep.len(),
        cutoffs
            .iter()
            .map(|c| c.map_or("none".to_string(), |x| x.to_string()))
            .collect::<Vec<_>>()
            .join(", "),
    );
    // Never hand back an empty matrix.
    (n_drop > 0 && n_drop < keep.len()).then_some(keep)
}

/// Soft hint when running with `ColumnAlignment::Disjoint` and inputs
/// look like patchy multi-modal data. Logs once at WARN level; never
/// errors.
fn maybe_warn_multimodal_pattern(data_vec: &SparseIoVec) {
    let n_backends = data_vec.len();
    if n_backends < 2 {
        return;
    }
    let intersection = data_vec.num_rows_in_at_least(n_backends);
    let min_backend_rows = (0..n_backends)
        .map(|j| data_vec[j].num_rows().unwrap_or(0))
        .min()
        .unwrap_or(0);
    if min_backend_rows == 0 {
        return;
    }
    let disjointness = 1.0_f64 - (intersection as f64) / (min_backend_rows as f64);
    if disjointness < MULTIMODAL_HINT_DISJOINTNESS_FRACTION {
        return;
    }

    // Cheap raw-barcode set-intersection across backends. `retain`
    // shrinks the accumulator in place — no per-entry `Box<str>` clone,
    // and an empty intersection short-circuits the remaining backends.
    let mut shared: Option<FxHashSet<Box<str>>> = None;
    for j in 0..n_backends {
        let names = match data_vec[j].column_names() {
            Ok(n) => n,
            Err(_) => return, // backend can't list columns; skip the hint
        };
        let set: FxHashSet<Box<str>> = names.into_iter().collect();
        match shared.as_mut() {
            None => shared = Some(set),
            Some(prev) => {
                prev.retain(|k| set.contains(k));
                if prev.is_empty() {
                    return;
                }
            }
        }
    }
    let shared_count = shared.map(|s| s.len()).unwrap_or(0);
    if shared_count == 0 {
        return;
    }

    warn!(
        "Inputs look multi-modal-shaped (feature-axis disjointness {:.0}% across {} \
         backends) and {} barcode(s) overlap across files. To glue cells across \
         modalities, pass `--multiome` (or the equivalent ColumnAlignment::Union). \
         Continuing with default Disjoint stacking — cells with shared barcodes \
         will be treated as distinct.",
        disjointness * 100.0,
        n_backends,
        shared_count
    );
}

/// Existing behavior: one batch label per cell, per-file slicing.
fn resolve_batch_disjoint(
    data_files: &[Box<str>],
    data_vec: &SparseIoVec,
    batch_files: Option<&[Box<str>]>,
    attach_data_name: bool,
) -> anyhow::Result<Vec<Box<str>>> {
    let mut batch_membership: Vec<Box<str>> = Vec::with_capacity(data_vec.num_columns());

    if let Some(batch_files) = batch_files {
        if batch_files.len() != data_files.len() {
            return Err(anyhow::anyhow!("# batch files != # of data files"));
        }
        for batch_file in batch_files.iter() {
            info!("Reading batch file: {}", batch_file);
            for s in read_lines(batch_file)? {
                batch_membership.push(s.to_string().into_boxed_str());
            }
        }
    } else {
        let column_counts = data_vec.num_columns_by_data()?;
        let column_names = data_vec.column_names()?;
        let mut col_start = 0usize;

        for (file_idx, &ncols) in column_counts.iter().enumerate() {
            let data_file = data_files[file_idx].clone();
            let (_dir, file_base, _ext) = common_io::dir_base_ext(&data_file)?;
            let col_end = col_start + ncols;
            let file_columns = &column_names[col_start..col_end];

            let appended_suffix =
                attach_data_name.then(|| format!("@{}", file_base).into_boxed_str());
            let (tags, used_embedded) = infer_batch_from_columns(
                file_columns,
                file_base.as_ref(),
                appended_suffix.as_deref(),
            );
            if used_embedded {
                info!(
                    "File {}: using embedded batch from column names (file '{}')",
                    file_idx, file_base
                );
            } else {
                info!(
                    "File {}: using file name '{}' as batch",
                    file_idx, file_base
                );
            }
            batch_membership.extend(tags);
            col_start = col_end;
        }
    }
    Ok(batch_membership)
}

/// Union-mode batch resolution. A cell can be in multiple backends, so
/// per-file batch labels don't compose: a barcode shared across two
/// files cannot have two batch labels. Rules:
///
/// - `batch_files`: must have exactly one file, listing one label per
///   unified cell in `data_vec.column_names()` order.
/// - Embedded `@batch` tag in raw column names (no `@<basename>`
///   suffix is added under Union): each backend independently infers
///   tags; conflicts (same barcode, different tag in two backends)
///   are an error.
/// - Fallback: constant `"all"` for cells whose backends don't agree
///   on an embedded tag. (File-name fallback from the Disjoint path
///   doesn't apply — a cell can come from many files.)
fn resolve_batch_union(
    data_vec: &SparseIoVec,
    batch_files: Option<&[Box<str>]>,
    n_cells: usize,
) -> anyhow::Result<Vec<Box<str>>> {
    if let Some(batch_files) = batch_files {
        if batch_files.len() != 1 {
            return Err(anyhow::anyhow!(
                "Under ColumnAlignment::Union, --batch-files must have exactly one \
                 file listing one label per unified cell (got {} files for {} \
                 unified cells). A cell shared across modalities cannot carry two \
                 batch labels.",
                batch_files.len(),
                n_cells
            ));
        }
        info!("Reading unified batch file: {}", batch_files[0]);
        let labels: Vec<Box<str>> = read_lines(&batch_files[0])?;
        if labels.len() != n_cells {
            return Err(anyhow::anyhow!(
                "Unified batch file {} has {} lines but data has {} unified cells",
                batch_files[0],
                labels.len(),
                n_cells
            ));
        }
        return Ok(labels);
    }

    // No batch_files: derive the per-cell @batch tag from the UNIFIED column
    // names. Under Union the displayed name already carries any embedded
    // `@tag` — whether from data-prep (`barcode@donor`) or from a per-file
    // barcode suffix (`barcode@sample`, added by `push_with_barcode_suffix`).
    // Each unified cell has exactly one name, so no cross-backend
    // reconciliation is needed: two cells that disagreed on their tag would
    // have different merge keys and never folded into one cell.
    let unified_names = data_vec.column_names()?;
    let (per_cell_tags, used_embedded) = infer_batch_from_columns(&unified_names, "", None);
    if used_embedded {
        info!(
            "Union mode: per-cell @batch tag taken from unified barcodes ({} cells)",
            n_cells
        );
        return Ok(per_cell_tags);
    }

    info!(
        "No --batch-files and no embedded @batch tags — falling back to single \
         batch 'all' (Union mode: per-file batch fallback is ambiguous)."
    );
    Ok(vec!["all".to_string().into_boxed_str(); n_cells])
}

/// Infer per-cell batch labels for one file's column names.
///
/// `appended_suffix` is the `@{file_base}` barcode disambiguator that
/// `SparseIoVec::push` tacks onto every column when multiple files are
/// loaded; it must be stripped *before* searching for a real embedded
/// `@batch` tag, otherwise `rsplit('@')` picks up the basename and every
/// cell in a file collapses to one wrong batch label.
///
/// Returns `(tags, used_embedded)` where `used_embedded=true` means the
/// raw column names already contained an `@batch` tag.
fn infer_batch_from_columns(
    file_columns: &[Box<str>],
    file_base: &str,
    appended_suffix: Option<&str>,
) -> (Vec<Box<str>>, bool) {
    fn raw_of<'a>(name: &'a str, suffix: Option<&str>) -> &'a str {
        match suffix {
            Some(sfx) => name.strip_suffix(sfx).unwrap_or(name),
            None => name,
        }
    }

    let has_embedded_batch = file_columns
        .first()
        .is_some_and(|name| raw_of(name.as_ref(), appended_suffix).contains('@'));

    if has_embedded_batch {
        let tags = file_columns
            .iter()
            .map(|col_name| {
                let raw = raw_of(col_name.as_ref(), appended_suffix);
                let embedded = raw.rsplit('@').next().unwrap_or(raw);
                embedded.to_string().into_boxed_str()
            })
            .collect();
        (tags, true)
    } else {
        let fallback: Box<str> = file_base.to_string().into_boxed_str();
        (vec![fallback; file_columns.len()], false)
    }
}

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

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

    fn cols(v: &[&str]) -> Vec<Box<str>> {
        v.iter()
            .map(|s| (*s).to_string().into_boxed_str())
            .collect()
    }

    #[test]
    fn embedded_donor_survives_push_suffix() {
        // Simulates `SparseIoVec::push` appending `@mix` to each column name
        // when multiple files are loaded. Raw names are `ACGT-1@donorA`,
        // `ACGT-2@donorB`.
        let names = cols(&[
            "ACGT-1@donorA@mix",
            "ACGT-2@donorB@mix",
            "ACGT-3@donorA@mix",
            "ACGT-4@donorB@mix",
        ]);
        let (tags, used_embedded) = infer_batch_from_columns(&names, "mix", Some("@mix"));
        assert!(used_embedded);
        assert_eq!(
            tags.iter().map(|b| b.as_ref()).collect::<Vec<_>>(),
            vec!["donorA", "donorB", "donorA", "donorB"]
        );
    }

    #[test]
    fn no_embedded_batch_falls_back_to_file_base() {
        // Barcodes without any embedded `@`.
        let names = cols(&["AAAA@s1", "CCCC@s1"]);
        let (tags, used_embedded) = infer_batch_from_columns(&names, "s1", Some("@s1"));
        assert!(!used_embedded);
        assert_eq!(
            tags.iter().map(|b| b.as_ref()).collect::<Vec<_>>(),
            vec!["s1", "s1"]
        );
    }

    #[test]
    fn single_file_embedded_batch() {
        // Single file → no `@file_base` suffix was appended.
        let names = cols(&["ACGT-1@donorA", "ACGT-2@donorB"]);
        let (tags, used_embedded) = infer_batch_from_columns(&names, "only", None);
        assert!(used_embedded);
        assert_eq!(
            tags.iter().map(|b| b.as_ref()).collect::<Vec<_>>(),
            vec!["donorA", "donorB"]
        );
    }

    #[test]
    fn single_file_no_embedded_batch() {
        let names = cols(&["AAAA", "CCCC"]);
        let (tags, used_embedded) = infer_batch_from_columns(&names, "only", None);
        assert!(!used_embedded);
        assert_eq!(
            tags.iter().map(|b| b.as_ref()).collect::<Vec<_>>(),
            vec!["only", "only"]
        );
    }

    #[test]
    fn empty_file_columns() {
        let names: Vec<Box<str>> = vec![];
        let (tags, used_embedded) = infer_batch_from_columns(&names, "x", Some("@x"));
        assert!(!used_embedded);
        assert!(tags.is_empty());
    }
}