bixverse-rs 0.4.6

//! SCENIC implementation from Aibar et al., Nat Methods, 2017, with the
//! following modifications:
//!
//! - **Quantisation and histogram-based splitting** to reduce predictor
//!   variable size.
//! - **Multi-output batching** for ExtraTrees and RandomForest learners,
//!   grouping genes to reduce the number of regression learners trained.
//! - **Gene batching strategies**: random assignment or SVD + k-means on
//!   gene loadings to group similar genes together.
//! - **GRNBoost2-style GBM** (Moerman et al., Bioinformatics, 2019) with
//!   histogram-based splits, parent-child subtraction, and OOB early
//!   stopping. Parallelism is exploited across targets rather than batching.

use ann_search_rs::prelude::*;
use ann_search_rs::utils::k_means_utils::{assign_all_parallel, train_centroids};
use faer::Mat;
use indexmap::IndexSet;
use rand::{Rng, SeedableRng, rngs::SmallRng};
use rayon::prelude::*;
use rustc_hash::FxHashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use thousands::Separable;

use crate::ml::clustering::k_means::*;
use crate::prelude::*;
use crate::single_cell::sc_processing::pca::{SingleCellPcaParams, pca_on_sc_streaming};
use crate::single_cell::sc_utils::simd::*;
use crate::single_cell::sc_utils::utils_tree::*;

/// How many genes to test for in one go
pub(crate) const SCENIC_GENE_CHUNK_SIZE: usize = 1024;

/// How many target genes to batch into a single multi-output tree ensemble.
/// 64 targets * 256 bins * 16 bytes ≈ 256KB per feature histogram, fits L2
/// comfortably.
pub const MULTI_OUTPUT_BATCH: usize = 64;

/// Threshold above which full-feature histogram subtraction outperforms
/// per-node feature sampling for single-target GBM.
const GBM_SUBTRACTION_CELL_THRESHOLD: usize = 10_000;

///////////
// Enums //
///////////

/// Which regression learner to use
#[derive(Clone, Debug)]
pub enum RegressionLearner {
    /// ExtraTree regression learner (RF's chaotic cousin)
    ExtraTrees(ExtraTreesConfig),
    /// RandomForest regression learner
    RandomForest(RandomForestConfig),
    /// GradientBoosted learner
    GradientBoosting(GradientBoostingConfig),
}

impl Default for RegressionLearner {
    fn default() -> Self {
        RegressionLearner::GradientBoosting(GradientBoostingConfig::default())
    }
}

/// Parse a regression learner from a string
///
/// ### Params
///
/// * `s` - String to parse. Accepts `"extratrees"`, `"rf"`, or
///   `"randomforest"` (case-insensitive).
///
/// ### Returns
///
/// `Some(RegressionLearner)` with default parameters, or `None` if unrecognised.
pub fn parse_regression_learner(s: &str) -> Option<RegressionLearner> {
    match s.to_lowercase().as_str() {
        "extratrees" => Some(RegressionLearner::ExtraTrees(ExtraTreesConfig::default())),
        "rf" | "randomforest" => Some(RegressionLearner::RandomForest(
            RandomForestConfig::default(),
        )),
        "boosted" | "grnboost2" | "lightgbm" => Some(RegressionLearner::GradientBoosting(
            GradientBoostingConfig::default(),
        )),
        _ => None,
    }
}

//////////////////
// Shared trait //
//////////////////

/// Shared configuration interface for tree-based regressors
pub trait TreeRegressorConfig: Sync {
    /// Number of trees in the ensemble
    fn n_trees(&self) -> usize;

    /// Minimum number of samples required in a leaf node
    fn min_samples_leaf(&self) -> usize;

    /// Number of features to consider at each split.
    /// A value of `0` means use `sqrt(n_features)`.
    fn n_features_split(&self) -> usize;

    /// Whether to use random thresholds (ExtraTrees) rather than optimal
    /// thresholds (RF)
    fn random_threshold(&self) -> bool;

    /// Fraction of samples to use per tree
    ///
    /// ### Returns
    ///
    /// Proportion in `[0, 1]`. Defaults to `1.0` (use all samples).
    fn subsample_rate(&self) -> f32 {
        1.0
    }

    /// Whether to sample with replacement (bootstrapping)
    fn bootstrap(&self) -> bool {
        false
    }

    /// Maximum tree depth. `None` means grow until other stopping criteria
    /// are met.
    fn max_depth(&self) -> Option<usize> {
        None
    }

    /// Minimum node variance below which no split is attempted
    fn min_variance(&self) -> f32 {
        1e-10
    }

    /// Number of random thresholds to evaluate per feature per node
    /// (ExtraTrees only)
    fn n_thresholds(&self) -> usize {
        1
    }

    /// If `Some(frac)`, subsample this fraction of cells per tree.
    /// Overrides `subsample_rate` when set.
    fn subsample_frac(&self) -> Option<f32> {
        None
    }
}

////////////
// Params //
////////////

/// Configuration for an ExtraTrees ensemble
///
/// ### Fields
///
/// * `n_trees` - Number of trees to build.
/// * `min_samples_leaf` - Minimum samples per leaf node.
/// * `n_features_split` - Features considered per split; `0` means
///   `sqrt(n_features)`.
/// * `n_thresholds` - Random thresholds to test per feature per node.
/// * `max_depth` - Maximum tree depth.
/// * `subsample_frac` - Optional fraction of cells to subsample per tree.
#[derive(Clone, Debug)]
pub struct ExtraTreesConfig {
    /// Number of trees to build.
    pub n_trees: usize,
    /// Minimum samples per leaf node.
    pub min_samples_leaf: usize,
    /// Features considered per split; `0` means `sqrt(n_features)`.
    pub n_features_split: usize,
    /// Random thresholds to test per feature per node.
    pub n_thresholds: usize,
    /// Maximum tree depth.
    pub max_depth: Option<usize>,
    /// Optional fraction of cells to subsample per tree.
    pub subsample_frac: Option<f32>,
}

/// Default implementation for `ExtraTreesConfig`
impl Default for ExtraTreesConfig {
    fn default() -> Self {
        Self {
            n_trees: 500,
            min_samples_leaf: 50,
            n_features_split: 0,
            n_thresholds: 1,
            max_depth: Some(10),
            subsample_frac: None,
        }
    }
}

/// Implementation of TreeRegressorConfig for `ExtraTreesConfig`
impl TreeRegressorConfig for ExtraTreesConfig {
    fn n_trees(&self) -> usize {
        self.n_trees
    }
    fn min_samples_leaf(&self) -> usize {
        self.min_samples_leaf
    }
    fn n_features_split(&self) -> usize {
        self.n_features_split
    }
    fn random_threshold(&self) -> bool {
        true
    }
    fn n_thresholds(&self) -> usize {
        self.n_thresholds
    }
    fn max_depth(&self) -> Option<usize> {
        self.max_depth
    }
    fn subsample_frac(&self) -> Option<f32> {
        self.subsample_frac
    }
}

/// Configuration for a RandomForest ensemble
///
/// ### Fields
///
/// * `n_trees` - Number of trees to build.
/// * `min_samples_leaf` - Minimum samples per leaf node.
/// * `n_features_split` - Features considered per split; `0` means
///   `sqrt(n_features)`.
/// * `subsample_rate` - Fraction of samples to draw per tree (without
///   replacement unless `bootstrap` is set).
/// * `bootstrap` - Whether to sample with replacement.
/// * `max_depth` - Maximum tree depth.
/// * `subsample_frac` - Optional fraction of cells to subsample per tree
///   (overrides `subsample_rate`).
#[derive(Clone, Debug)]
pub struct RandomForestConfig {
    /// Number of trees to build.
    pub n_trees: usize,
    /// Minimum samples per leaf node.
    pub min_samples_leaf: usize,
    /// Features considered per split; `0` means `sqrt(n_features)`.
    pub n_features_split: usize,
    /// Fraction of samples to draw per tree (without replacement unless
    /// `bootstrap` is set).
    pub subsample_rate: f32,
    /// Whether to sample with replacement.
    pub bootstrap: bool,
    /// Maximum tree depth.
    pub max_depth: Option<usize>,
    /// Optional fraction of cells to subsample per tree (overrides
    /// `subsample_rate`).
    pub subsample_frac: Option<f32>,
}

/// Default implementation for `RandomForestConfig`
impl Default for RandomForestConfig {
    fn default() -> Self {
        Self {
            n_trees: 250,
            min_samples_leaf: 50,
            n_features_split: 0,
            subsample_rate: 0.632,
            bootstrap: false,
            max_depth: Some(10),
            subsample_frac: None,
        }
    }
}

/// Implementation of TreeRegressorConfig for `RandomForestConfig`
impl TreeRegressorConfig for RandomForestConfig {
    fn n_trees(&self) -> usize {
        self.n_trees
    }
    fn min_samples_leaf(&self) -> usize {
        self.min_samples_leaf
    }
    fn n_features_split(&self) -> usize {
        self.n_features_split
    }
    fn random_threshold(&self) -> bool {
        false
    }
    fn subsample_rate(&self) -> f32 {
        self.subsample_rate
    }
    fn bootstrap(&self) -> bool {
        self.bootstrap
    }
    fn max_depth(&self) -> Option<usize> {
        self.max_depth
    }
    fn subsample_frac(&self) -> Option<f32> {
        self.subsample_frac
    }
}

/// Configuration for a GRNBoost2-style gradient boosted tree ensemble.
///
/// Single-target only. Trees are built sequentially, each fitting the
/// residuals from the prior ensemble. Early stopping via OOB improvement
/// prevents overfitting and is the main source of speedup over RF/ET.
#[derive(Clone, Debug)]
pub struct GradientBoostingConfig {
    /// Maximum number of boosting rounds (early stopping usually triggers
    /// well before this).
    pub n_trees_max: usize,
    /// Shrinkage applied to each tree's predictions.
    pub learning_rate: f32,
    /// Maximum tree depth. Shallow trees (3-5) work best for GBM.
    pub max_depth: usize,
    /// Minimum training samples per leaf node.
    pub min_samples_leaf: usize,
    /// Number of recent OOB improvements to average for the early stopping
    /// criterion. Stops when the rolling average drops to zero or below.
    pub early_stop_window: usize,
    /// Fraction of samples used for training each tree. The complement
    /// forms the OOB set used for early stopping evaluation.
    pub subsample_rate: f32,
    /// Features to evaluate per split; `0` means all features (recommended
    /// with histogram subtraction).
    pub n_features_split: usize,
}

impl Default for GradientBoostingConfig {
    fn default() -> Self {
        Self {
            n_trees_max: 1000,
            learning_rate: 0.01,
            max_depth: 3,
            min_samples_leaf: 50,
            early_stop_window: 25,
            subsample_rate: 0.9,
            n_features_split: 0,
        }
    }
}

//////////////////////
// Sparse Y builder //
//////////////////////

/// Static sparse Y representation for a batch of targets.
///
/// Indexed by original cell ID - remains constant across all trees and nodes.
/// Only sample index arrays are partitioned during tree building.
pub struct SparseYBatch {
    /// Start offset into `target_indices` and `values` for each cell. Length:
    /// n_cells + 1. Cell `c` has entries at `offsets[c]..offsets[c + 1]`.
    pub offsets: Vec<u32>,
    /// Which target within the batch is non-zero (0..n_targets). Since
    /// n_targets <= 64, u8 suffices.
    pub target_indices: Vec<u8>,
    /// Corresponding expression values.
    pub values: Vec<f32>,
}

impl SparseYBatch {
    /// Build from the sparse target columns for one batch.
    ///
    /// ### Params
    ///
    /// * `targets` - Sparse columns for each target in this batch.
    /// * `n_cells` - Total number of cells (determines offsets length).
    ///
    /// ### Returns
    ///
    /// A `SparseYBatch` where entries are sorted by cell, then by target
    /// index within each cell.
    pub fn from_targets(
        targets: &[SparseAxis<u32, f32>],
        n_cells: usize,
    ) -> Result<Self, BixverseErrors> {
        // count non-zeros per cell across all targets
        let mut counts_per_cell = vec![0u32; n_cells];
        for target in targets {
            let (indices, _) = target.get_indices_data_2()?;
            for &idx in indices {
                counts_per_cell[idx] += 1;
            }
        }

        // build offsets via prefix sum
        let mut offsets = Vec::with_capacity(n_cells + 1);
        offsets.push(0u32);
        let mut running = 0u32;
        for &c in &counts_per_cell {
            running += c;
            offsets.push(running);
        }
        let total_nnz = running as usize;

        let mut target_indices = vec![0u8; total_nnz];
        let mut values = vec![0.0f32; total_nnz];

        // fill using a write cursor per cell
        let mut cursor = vec![0u32; n_cells];
        for (k, target) in targets.iter().enumerate() {
            let (indices, data) = target.get_indices_data_2()?;
            for (i, &cell) in indices.iter().enumerate() {
                let pos = (offsets[cell] + cursor[cell]) as usize;
                target_indices[pos] = k as u8;
                values[pos] = data[i];
                cursor[cell] += 1;
            }
        }

        Ok(Self {
            offsets,
            target_indices,
            values,
        })
    }

    /// Iterate non-zero entries for a given cell.
    ///
    /// ### Params
    ///
    /// * `cell` - Cell index
    ///
    /// ### Returns
    ///
    /// Returns the target indices and values
    #[inline(always)]
    fn cell_entries(&self, cell: usize) -> (&[u8], &[f32]) {
        let s = self.offsets[cell] as usize;
        let e = self.offsets[cell + 1] as usize;
        (&self.target_indices[s..e], &self.values[s..e])
    }
}

/////////////////////
// Dense Y builder //
/////////////////////

/// Build an interleaved dense Y matrix from multiple sparse target columns
///
/// ### Params
///
/// * `targets` - Sparse target gene expression columns.
/// * `n_samples` - Total number of cells.
///
/// ### Returns
///
/// Dense vector of length `n_samples * n_targets` with layout
/// `y[sample * n_targets + target_idx]`. Absent entries default to `0.0`.
#[allow(dead_code)]
fn build_y_dense_multi(
    targets: &[SparseAxis<u32, f32>],
    n_samples: usize,
) -> Result<Vec<f32>, BixverseErrors> {
    let n_targets = targets.len();
    let mut y_dense = vec![0.0f32; n_samples * n_targets];
    for (k, target) in targets.iter().enumerate() {
        let (y_indices, y_data) = target.get_indices_data_2()?;
        for (i, &idx) in y_indices.iter().enumerate() {
            y_dense[idx * n_targets + k] = y_data[i];
        }
    }

    Ok(y_dense)
}

///////////////////////////////
// Multi-output tree buffers //
///////////////////////////////

/// Reusable scratch buffers for building multi-output trees
///
/// All allocations are done once at construction and reused across trees and
/// nodes to avoid repeated heap allocation in the hot path.
#[allow(dead_code)]
struct TreeBuffers {
    /// Feature index permutation buffer for partial Fisher-Yates shuffle.
    feat_buf: Vec<usize>,
    /// Temporary sample indices for the left child during partitioning.
    left_buf: Vec<u32>,
    /// Temporary sample indices for the right child during partitioning.
    right_buf: Vec<u32>,
    /// Temporary interleaved Y values for the left child; layout
    /// `[sample * n_targets + target]`.
    left_y_buf: Vec<f32>,
    /// Temporary interleaved Y values for the right child; same layout as
    /// `left_y_buf`.
    right_y_buf: Vec<f32>,
    /// Per-bin sample counts; layout `counts[bin]`.
    counts: [u32; 256],
    /// Per-bin, per-target Y sums; layout
    y_sums: Vec<f32>,
    /// Per-bin, per-target Y sum-of-squares; same layout as `y_sums`.
    y_sum_sqs: Vec<f32>,
    /// Prefix-sum of `counts` over bins.
    cum_counts: [u32; 256],
    /// Prefix-sum of `y_sums` over bins; same layout as `y_sums`.
    cum_y_sums: Vec<f32>,
    /// Prefix-sum of `y_sum_sqs` over bins; same layout as `y_sums`.
    cum_y_sum_sqs: Vec<f32>,
    /// Left-child Y sums captured at the best split.
    best_y_sums_l: Vec<f32>,
    /// Left-child Y sum-of-squares captured at the best split.
    best_y_sum_sqs_l: Vec<f32>,
    /// Per-target parent variance scratch space.
    parent_vars: Vec<f32>,
}

impl TreeBuffers {
    /// Allocate all scratch buffers for the given problem dimensions
    ///
    /// ### Params
    ///
    /// * `n_features` - Total number of features (TFs).
    /// * `n_samples` - Maximum number of samples per tree.
    /// * `n_targets` - Number of target genes in the current batch.
    #[allow(dead_code)]
    fn new(n_features: usize, n_samples: usize, n_targets: usize) -> Self {
        Self {
            feat_buf: (0..n_features).collect(),
            left_buf: vec![0; n_samples],
            right_buf: vec![0; n_samples],
            left_y_buf: vec![0.0; n_samples * n_targets],
            right_y_buf: vec![0.0; n_samples * n_targets],
            counts: [0; 256],
            y_sums: vec![0_f32; 256 * n_targets],
            y_sum_sqs: vec![0_f32; 256 * n_targets],
            cum_counts: [0; 256],
            cum_y_sums: vec![0_f32; 256 * n_targets],
            cum_y_sum_sqs: vec![0_f32; 256 * n_targets],
            best_y_sums_l: vec![0_f32; n_targets],
            best_y_sum_sqs_l: vec![0_f32; n_targets],
            parent_vars: vec![0_f32; n_targets],
        }
    }

    /// Generate a tree buffer for SparseYBatch
    ///
    /// ### Params
    ///
    /// * `n_features` - Number of features
    /// * `n_samples` - Number of samples
    /// * `n_targets` - Number of target variables (do not go over 64 here!)
    ///
    /// ### Returns
    ///
    /// Initialised self
    fn new_sparse(n_features: usize, n_samples: usize, n_targets: usize) -> Self {
        Self {
            feat_buf: (0..n_features).collect(),
            left_buf: vec![0; n_samples],
            right_buf: vec![0; n_samples],
            left_y_buf: Vec::new(),
            right_y_buf: Vec::new(),
            counts: [0; 256],
            y_sums: vec![0_f32; 256 * n_targets],
            y_sum_sqs: vec![0_f32; 256 * n_targets],
            cum_counts: [0; 256],
            cum_y_sums: vec![0_f32; 256 * n_targets],
            cum_y_sum_sqs: vec![0_f32; 256 * n_targets],
            best_y_sums_l: vec![0_f32; n_targets],
            best_y_sum_sqs_l: vec![0_f32; n_targets],
            parent_vars: vec![0_f32; n_targets],
        }
    }

    /// Build bin histograms and their prefix-sum (cumulative) counterparts
    ///
    /// ### Params
    ///
    /// * `tf_col` - Quantised feature column for all cells.
    /// * `sample_slice` - Indices of the active samples in this node.
    /// * `y_slice` - Interleaved target values; layout
    ///   `[sample * n_targets + target]`.
    /// * `n_targets` - Number of targets in the current batch.
    ///
    /// ### Implementation details
    ///
    /// Iterates over `sample_slice` and accumulates each sample into its bin
    /// (`tf_col[sample_idx]`). After the pass, a single forward scan builds the
    /// cumulative counts, sums, and sum-of-squares needed to evaluate any split
    /// threshold in O(1) during the split search.
    #[inline]
    #[allow(dead_code)]
    fn build_histograms(
        &mut self,
        tf_col: &[u8],
        sample_slice: &[u32],
        y_slice: &[f32],
        n_targets: usize,
    ) -> (usize, usize) {
        if sample_slice.is_empty() {
            return (0, 0);
        }

        let mut min_bin = 255u8;
        let mut max_bin = 0u8;
        for &s in sample_slice {
            let bin = tf_col[s as usize];
            if bin < min_bin {
                min_bin = bin;
            }
            if bin > max_bin {
                max_bin = bin;
            }
        }

        let min_b = min_bin as usize;
        let max_b = max_bin as usize;

        self.counts[min_b..=max_b].fill(0);
        let start_idx = min_b * n_targets;
        let end_idx = (max_b + 1) * n_targets;
        self.y_sums[start_idx..end_idx].fill(0.0);
        self.y_sum_sqs[start_idx..end_idx].fill(0.0);

        for i in 0..sample_slice.len() {
            let bin = tf_col[sample_slice[i] as usize] as usize;
            self.counts[bin] += 1;
            let y_base = i * n_targets;
            let h_base = bin * n_targets;
            for k in 0..n_targets {
                let y = y_slice[y_base + k];
                self.y_sums[h_base + k] += y;
                self.y_sum_sqs[h_base + k] += y * y;
            }
        }

        self.cum_counts[min_b] = self.counts[min_b];
        let start_h = min_b * n_targets;
        let next_h = (min_b + 1) * n_targets;
        self.cum_y_sums[start_h..next_h].copy_from_slice(&self.y_sums[start_h..next_h]);
        self.cum_y_sum_sqs[start_h..next_h].copy_from_slice(&self.y_sum_sqs[start_h..next_h]);

        for b in (min_b + 1)..=max_b {
            self.cum_counts[b] = self.cum_counts[b - 1] + self.counts[b];
            let prev = (b - 1) * n_targets;
            let curr = b * n_targets;
            for k in 0..n_targets {
                self.cum_y_sums[curr + k] = self.cum_y_sums[prev + k] + self.y_sums[curr + k];
                self.cum_y_sum_sqs[curr + k] =
                    self.cum_y_sum_sqs[prev + k] + self.y_sum_sqs[curr + k];
            }
        }

        (min_b, max_b)
    }

    /// Build histograms from sparse Y -- only touches non-zero target entries.
    ///
    /// Counts are accumulated for all samples (target-independent). Y sums and
    /// sum-of-squares are accumulated only for non-zero entries; zero targets
    /// contribute nothing and are skipped.
    ///
    /// * `tf_col` - Quantised feature column for all cells.
    /// * `sample_slice` - Indices of the active samples in this node.
    /// * `sparse_y` - Sparse representation of the targets
    /// * `n_targets` - Number of targets in the current batch.
    #[inline]
    fn build_histograms_sparse(
        &mut self,
        tf_col: &[u8],
        sample_slice: &[u32],
        sparse_y: &SparseYBatch,
        n_targets: usize,
    ) -> (usize, usize) {
        if sample_slice.is_empty() {
            return (0, 0);
        }

        // 1. Blazingly fast O(N) pre-pass to find the active bin range
        let mut min_bin = 255u8;
        let mut max_bin = 0u8;
        for &s in sample_slice {
            let bin = tf_col[s as usize];
            if bin < min_bin {
                min_bin = bin;
            }
            if bin > max_bin {
                max_bin = bin;
            }
        }

        let min_b = min_bin as usize;
        let max_b = max_bin as usize;

        // 2. Only zero out the dirtied slice (The Bandwidth Fix)
        self.counts[min_b..=max_b].fill(0);
        let start_idx = min_b * n_targets;
        let end_idx = (max_b + 1) * n_targets;
        self.y_sums[start_idx..end_idx].fill(0.0);
        self.y_sum_sqs[start_idx..end_idx].fill(0.0);

        // 3. Accumulate only within the active bounds
        for &s in sample_slice {
            let bin = tf_col[s as usize] as usize;
            self.counts[bin] += 1;

            let (tgt_indices, tgt_values) = sparse_y.cell_entries(s as usize);
            let h_base = bin * n_targets;
            for i in 0..tgt_indices.len() {
                let k = tgt_indices[i] as usize;
                let y = tgt_values[i];
                self.y_sums[h_base + k] += y;
                self.y_sum_sqs[h_base + k] += y * y;
            }
        }

        // 4. Prefix sums strictly over the active range
        self.cum_counts[min_b] = self.counts[min_b];
        let start_h = min_b * n_targets;
        let next_h = (min_b + 1) * n_targets;
        self.cum_y_sums[start_h..next_h].copy_from_slice(&self.y_sums[start_h..next_h]);
        self.cum_y_sum_sqs[start_h..next_h].copy_from_slice(&self.y_sum_sqs[start_h..next_h]);

        for b in (min_b + 1)..=max_b {
            self.cum_counts[b] = self.cum_counts[b - 1] + self.counts[b];
            let prev = (b - 1) * n_targets;
            let curr = b * n_targets;

            // copy current bin values into cumulative position
            self.cum_y_sums[curr..curr + n_targets]
                .copy_from_slice(&self.y_sums[curr..curr + n_targets]);
            self.cum_y_sum_sqs[curr..curr + n_targets]
                .copy_from_slice(&self.y_sum_sqs[curr..curr + n_targets]);

            // for n_targets <= 64 this fits on the stack comfortably.
            let mut prev_sums = [0_f32; MULTI_OUTPUT_BATCH];
            let mut prev_sum_sqs = [0_f32; MULTI_OUTPUT_BATCH];
            prev_sums[..n_targets].copy_from_slice(&self.cum_y_sums[prev..prev + n_targets]);
            prev_sum_sqs[..n_targets].copy_from_slice(&self.cum_y_sum_sqs[prev..prev + n_targets]);

            accumulate_f32_simd(
                &mut self.cum_y_sums[curr..curr + n_targets],
                &prev_sums[..n_targets],
                n_targets,
            );
            accumulate_f32_simd(
                &mut self.cum_y_sum_sqs[curr..curr + n_targets],
                &prev_sum_sqs[..n_targets],
                n_targets,
            );
        }

        (min_b, max_b)
    }
}

//////////////////
// Tree helpers //
//////////////////

/// Compute the variance of a node from sufficient statistics
///
/// ### Params
///
/// * `sum` - Sum of values in the node.
/// * `sum_sq` - Sum of squared values in the node.
/// * `n` - Number of samples in the node.
///
/// ### Returns
///
/// The variance, clamped to `0.0`. Returns `0.0` for nodes with fewer than 2
/// samples.
#[inline]
fn node_variance_f32(sum: f32, sum_sq: f32, n: usize) -> f32 {
    if n < 2 {
        return 0.0;
    }
    let nf = n as f32;
    f32::max(0.0, sum_sq / nf - (sum / nf) * (sum / nf))
}

/// Extract per-target importance vectors from the interleaved
/// accumulator and normalise each to sum to 1.0.
///
/// The interleaved layout is `importances[feature * n_targets + target]`.
///
/// ### Params
///
/// * `importances` - Flat interleaved importance accumulator.
/// * `n_features` - Number of features.
/// * `n_targets` - Number of targets in the batch.
///
/// ### Returns
///
/// `result[target_idx]` is a `Vec<f32>` of length `n_features`,
/// normalised to sum to 1.0.
fn extract_and_normalise_multi_importances(
    importances: &[f32],
    n_features: usize,
    n_targets: usize,
) -> Vec<Vec<f32>> {
    let mut result = Vec::with_capacity(n_targets);
    for k in 0..n_targets {
        let mut target_imp = Vec::with_capacity(n_features);
        for f in 0..n_features {
            target_imp.push(importances[f * n_targets + k]);
        }
        normalise_importances(&mut target_imp);
        result.push(target_imp);
    }
    result
}

/// Evaluate one candidate split threshold across all targets and update the
/// best-split state
///
/// ### Params
///
/// * `threshold` - Bin index to split at (samples with bin ≤ threshold go left).
/// * `feat` - Feature index being evaluated.
/// * `parent_vars` - Per-target parent variance.
/// * `n` - Total samples in this node.
/// * `y_sums_total` - Per-target sum of Y in this node.
/// * `y_sum_sqs_total` - Per-target sum of Y² in this node.
/// * `cum_counts` - Cumulative bin sample counts.
/// * `cum_y_sums` - Cumulative per-target Y sums over bins.
/// * `cum_y_sum_sqs` - Cumulative per-target Y² sums over bins.
/// * `n_targets` - Number of targets in the current batch.
/// * `min_samples_leaf` - Minimum samples required per child.
/// * `best_score` - Current best aggregate variance reduction (updated in
///   place).
/// * `best_feature` - Feature index of the current best split (updated in
///   place).
/// * `best_threshold_u8` - Threshold of the current best split (updated in
///   place).
/// * `best_n_left` - Left child size at the current best split (updated in
///   place).
/// * `best_y_sums_l` - Left-child Y sums at the current best split (updated in
///   place).
/// * `best_y_sum_sqs_l` - Left-child Y² sums at the current best split (updated
///   in place).
///
/// ### Implementation details
///
/// The split score is the sum of variance reductions across all targets,
/// weighted by node size relative to the root. Splits that would produce a
/// child smaller than `min_samples_leaf` are rejected immediately. The
/// left-child sufficient statistics are read directly from the cumulative
/// histograms; right-child statistics are derived by subtraction.
#[inline(always)]
#[allow(clippy::too_many_arguments)]
fn evaluate_split_multi(
    threshold: usize,
    feat: usize,
    parent_vars: &[f32],
    n: usize,
    y_sums_total: &[f32],
    y_sum_sqs_total: &[f32],
    cum_counts: &[u32; 256],
    cum_y_sums: &[f32],
    cum_y_sum_sqs: &[f32],
    n_targets: usize,
    min_samples_leaf: u32,
    best_score: &mut f32,
    best_feature: &mut usize,
    best_threshold_u8: &mut u8,
    best_n_left: &mut u32,
    best_y_sums_l: &mut [f32],
    best_y_sum_sqs_l: &mut [f32],
) {
    let n_left = cum_counts[threshold];
    let n_right = n as u32 - n_left;

    if n_left < min_samples_leaf || n_right < min_samples_leaf {
        return;
    }

    let nl = n_left as f32;
    let nr = n_right as f32;
    let nf = n as f32;
    let h_base = threshold * n_targets;

    let score = evaluate_split_score_f32_simd(
        parent_vars,
        y_sums_total,
        y_sum_sqs_total,
        &cum_y_sums[h_base..h_base + n_targets],
        &cum_y_sum_sqs[h_base..h_base + n_targets],
        n_targets,
        1.0 / nl,
        1.0 / nr,
        nl / nf,
        nr / nf,
    );

    if score > *best_score {
        *best_score = score;
        *best_feature = feat;
        *best_threshold_u8 = threshold as u8;
        *best_n_left = n_left;
        best_y_sums_l.copy_from_slice(&cum_y_sums[h_base..h_base + n_targets]);
        best_y_sum_sqs_l.copy_from_slice(&cum_y_sum_sqs[h_base..h_base + n_targets]);
    }
}

///////////////////
// Tree building //
///////////////////

/// Recursively build a single tree node using sparse Y.
///
/// Only sample_slice is partitioned in place. Y lookups go through the static
/// SparseYBatch. Histograms are built once per candidate feature; both the
/// ExtraTrees (random threshold) and RF (exhaustive) paths evaluate splits via
/// the cumulative histogram in O(1) per threshold.
///
/// ### Params
///
/// * `sparse_y` - SparseYBatch structure
/// * `x` - Quantised feature store.
/// * `sample_slice` - Indices of the active samples; partitioned in place
///   around the best split.
/// * `y_sums` - Per-target Y sums for this node.
/// * `y_sum_sqs` - Per-target Y squared sums for this node.
/// * `n_total` - Total samples at the tree root (used to weight importance
///   contributions).
/// * `n_targets` - Number of targets in the current batch.
/// * `n_features_split` - Number of features to sample at each split.
/// * `config` - Learner configuration.
/// * `depth` - Current depth in the tree.
/// * `importances` - Accumulated importance array; layout
///   `[feature * n_targets + target]`.
/// * `bufs` - Reusable scratch buffers.
/// * `rng` - Per-tree RNG.
#[allow(clippy::too_many_arguments)]
fn build_node_multi_sparse(
    sparse_y: &SparseYBatch,
    x: &QuantisedStore,
    sample_slice: &mut [u32],
    y_sums: &[f32],
    y_sum_sqs: &[f32],
    n_total: usize,
    n_targets: usize,
    n_features_split: usize,
    config: &dyn TreeRegressorConfig,
    depth: usize,
    importances: &mut [f32],
    bufs: &mut TreeBuffers,
    rng: &mut SmallRng,
) {
    let n = sample_slice.len();

    let mut total_parent_var = 0.0f32;
    for k in 0..n_targets {
        let v = node_variance_f32(y_sums[k], y_sum_sqs[k], n);
        bufs.parent_vars[k] = v;
        total_parent_var += v;
    }

    let max_depth_reached = config.max_depth().is_some_and(|d| depth >= d);
    if n < 2 * config.min_samples_leaf()
        || (total_parent_var) < config.min_variance()
        || max_depth_reached
    {
        return;
    }

    let n_features = x.n_features;
    let k_feats = n_features_split.min(n_features);

    for i in 0..k_feats {
        let j = rng.random_range(i..n_features);
        bufs.feat_buf.swap(i, j);
    }

    let mut best_score = 0.0f32;
    let mut best_feature = usize::MAX;
    let mut best_threshold_u8 = 0u8;
    let mut best_n_left = 0u32;

    for fi_idx in 0..k_feats {
        let feat = bufs.feat_buf[fi_idx];
        let tf_col = x.get_col(feat);

        let (min_bin, max_bin) =
            bufs.build_histograms_sparse(tf_col, sample_slice, sparse_y, n_targets);

        if min_bin == max_bin {
            continue;
        }

        if config.random_threshold() {
            for _ in 0..config.n_thresholds() {
                let threshold = rng.random_range(min_bin..max_bin);
                evaluate_split_multi(
                    threshold,
                    feat,
                    &bufs.parent_vars,
                    n,
                    y_sums,
                    y_sum_sqs,
                    &bufs.cum_counts,
                    &bufs.cum_y_sums,
                    &bufs.cum_y_sum_sqs,
                    n_targets,
                    config.min_samples_leaf() as u32,
                    &mut best_score,
                    &mut best_feature,
                    &mut best_threshold_u8,
                    &mut best_n_left,
                    &mut bufs.best_y_sums_l,
                    &mut bufs.best_y_sum_sqs_l,
                );
            }
        } else {
            for threshold in min_bin..max_bin {
                evaluate_split_multi(
                    threshold,
                    feat,
                    &bufs.parent_vars,
                    n,
                    y_sums,
                    y_sum_sqs,
                    &bufs.cum_counts,
                    &bufs.cum_y_sums,
                    &bufs.cum_y_sum_sqs,
                    n_targets,
                    config.min_samples_leaf() as u32,
                    &mut best_score,
                    &mut best_feature,
                    &mut best_threshold_u8,
                    &mut best_n_left,
                    &mut bufs.best_y_sums_l,
                    &mut bufs.best_y_sum_sqs_l,
                );
            }
        }
    }

    if best_feature == usize::MAX {
        return;
    }

    // Accumulate importance
    let nl = best_n_left as f32;
    let nr = (n as u32 - best_n_left) as f32;
    let nf = n as f32;
    let weight = nf / n_total as f32;
    let imp_base = best_feature * n_targets;

    for k in 0..n_targets {
        let y_sum_l = bufs.best_y_sums_l[k];
        let y_sum_sq_l = bufs.best_y_sum_sqs_l[k];
        let y_sum_r = y_sums[k] - y_sum_l;
        let y_sum_sq_r = y_sum_sqs[k] - y_sum_sq_l;
        let mean_l = y_sum_l / nl;
        let var_l = f32::max(0.0, y_sum_sq_l / nl - mean_l * mean_l);
        let mean_r = y_sum_r / nr;
        let var_r = f32::max(0.0, y_sum_sq_r / nr - mean_r * mean_r);
        let reduction = bufs.parent_vars[k] - (nl / nf) * var_l - (nr / nf) * var_r;
        importances[imp_base + k] += weight * f32::max(0.0, reduction);
    }

    // Partition
    let mut left_y_sums = [0.0f32; MULTI_OUTPUT_BATCH];
    let mut left_y_sum_sqs = [0.0f32; MULTI_OUTPUT_BATCH];
    left_y_sums[..n_targets].copy_from_slice(&bufs.best_y_sums_l[..n_targets]);
    left_y_sum_sqs[..n_targets].copy_from_slice(&bufs.best_y_sum_sqs_l[..n_targets]);

    let mut right_y_sums = [0.0f32; MULTI_OUTPUT_BATCH];
    let mut right_y_sum_sqs = [0.0f32; MULTI_OUTPUT_BATCH];
    for k in 0..n_targets {
        right_y_sums[k] = y_sums[k] - left_y_sums[k];
        right_y_sum_sqs[k] = y_sum_sqs[k] - left_y_sum_sqs[k];
    }

    let tf_col = x.get_col(best_feature);
    let mut l_idx = 0usize;
    let mut r_idx = 0usize;

    for i in 0..n {
        let s = sample_slice[i];
        if tf_col[s as usize] <= best_threshold_u8 {
            bufs.left_buf[l_idx] = s;
            l_idx += 1;
        } else {
            bufs.right_buf[r_idx] = s;
            r_idx += 1;
        }
    }

    sample_slice[..l_idx].copy_from_slice(&bufs.left_buf[..l_idx]);
    sample_slice[l_idx..].copy_from_slice(&bufs.right_buf[..r_idx]);

    let (left_samples, right_samples) = sample_slice.split_at_mut(l_idx);

    build_node_multi_sparse(
        sparse_y,
        x,
        left_samples,
        &left_y_sums[..n_targets],
        &left_y_sum_sqs[..n_targets],
        n_total,
        n_targets,
        n_features_split,
        config,
        depth + 1,
        importances,
        bufs,
        rng,
    );
    build_node_multi_sparse(
        sparse_y,
        x,
        right_samples,
        &right_y_sums[..n_targets],
        &right_y_sum_sqs[..n_targets],
        n_total,
        n_targets,
        n_features_split,
        config,
        depth + 1,
        importances,
        bufs,
        rng,
    );
}

/// Recursively build a single tree node, accumulating feature importance in
/// place
///
/// No tree structure is stored -- importance is the sole output.
///
/// ### Params
///
/// * `y_slice` - Interleaved target values for the active samples; layout
///   `[sample * n_targets + target]`.
/// * `x` - Quantised feature store.
/// * `sample_slice` - Indices of the active samples; partitioned in place
///   around the best split.
/// * `y_sums` - Per-target Y sums for this node.
/// * `y_sum_sqs` - Per-target Y² sums for this node.
/// * `n_total` - Total samples at the tree root (used to weight importance
///   contributions).
/// * `n_targets` - Number of targets in the current batch.
/// * `n_features_split` - Number of features to sample at each split.
/// * `config` - Learner configuration.
/// * `depth` - Current depth in the tree.
/// * `importances` - Accumulated importance array; layout
///   `[feature * n_targets + target]`.
/// * `bufs` - Reusable scratch buffers.
/// * `rng` - Per-tree RNG.
///
/// ### Implementation details
///
/// Stopping criteria are checked first (min samples, min variance, max depth).
/// Features are selected via a partial Fisher-Yates shuffle of `bufs.feat_buf`.
/// For each selected feature, bin histograms are built once and then all
/// candidate thresholds are evaluated in a single pass over the 256 bins. For
/// ExtraTrees, only `n_thresholds` random thresholds are tested per feature
/// instead of all bins.
///
/// After the best split is found, its variance reduction is accumulated into
/// `importances` weighted by `n / n_total`. Samples and `y_slice` are
/// partitioned in place using `bufs.left_buf` / `bufs.right_buf` as temporary
/// storage, and the function recurses into both children.
#[allow(clippy::too_many_arguments)]
#[allow(dead_code)]
fn build_node_multi(
    y_slice: &mut [f32],
    x: &QuantisedStore,
    sample_slice: &mut [u32],
    y_sums: &[f32],
    y_sum_sqs: &[f32],
    n_total: usize,
    n_targets: usize,
    n_features_split: usize,
    config: &dyn TreeRegressorConfig,
    depth: usize,
    importances: &mut [f32],
    bufs: &mut TreeBuffers,
    rng: &mut SmallRng,
) {
    let n = sample_slice.len();

    let mut total_parent_var = 0.0f32;
    for k in 0..n_targets {
        let v = node_variance_f32(y_sums[k], y_sum_sqs[k], n);
        bufs.parent_vars[k] = v;
        total_parent_var += v;
    }

    let max_depth_reached = config.max_depth().is_some_and(|d| depth >= d);

    if n < 2 * config.min_samples_leaf()
        || total_parent_var < config.min_variance()
        || max_depth_reached
    {
        return;
    }

    let n_features = x.n_features;
    let k_feats = n_features_split.min(n_features);

    for i in 0..k_feats {
        let j = rng.random_range(i..n_features);
        bufs.feat_buf.swap(i, j);
    }

    let mut best_score = 0.0f32;
    let mut best_feature = usize::MAX;
    let mut best_threshold_u8 = 0u8;
    let mut best_n_left = 0u32;

    for fi_idx in 0..k_feats {
        let feat = bufs.feat_buf[fi_idx];
        let tf_col = x.get_col(feat);

        let (min_bin, max_bin) = bufs.build_histograms(tf_col, sample_slice, y_slice, n_targets);

        if min_bin == max_bin {
            continue;
        }

        if config.random_threshold() {
            for _ in 0..config.n_thresholds() {
                let threshold = rng.random_range(min_bin..max_bin);
                evaluate_split_multi(
                    threshold,
                    feat,
                    &bufs.parent_vars,
                    n,
                    y_sums,
                    y_sum_sqs,
                    &bufs.cum_counts,
                    &bufs.cum_y_sums,
                    &bufs.cum_y_sum_sqs,
                    n_targets,
                    config.min_samples_leaf() as u32,
                    &mut best_score,
                    &mut best_feature,
                    &mut best_threshold_u8,
                    &mut best_n_left,
                    &mut bufs.best_y_sums_l,
                    &mut bufs.best_y_sum_sqs_l,
                );
            }
        } else {
            for threshold in min_bin..max_bin {
                evaluate_split_multi(
                    threshold,
                    feat,
                    &bufs.parent_vars,
                    n,
                    y_sums,
                    y_sum_sqs,
                    &bufs.cum_counts,
                    &bufs.cum_y_sums,
                    &bufs.cum_y_sum_sqs,
                    n_targets,
                    config.min_samples_leaf() as u32,
                    &mut best_score,
                    &mut best_feature,
                    &mut best_threshold_u8,
                    &mut best_n_left,
                    &mut bufs.best_y_sums_l,
                    &mut bufs.best_y_sum_sqs_l,
                );
            }
        }
    }

    if best_feature == usize::MAX {
        return;
    }

    let nl = best_n_left as f32;
    let nr = (n as u32 - best_n_left) as f32;
    let nf = n as f32;
    let weight = nf / n_total as f32;
    let imp_base = best_feature * n_targets;

    for k in 0..n_targets {
        let y_sum_l = bufs.best_y_sums_l[k];
        let y_sum_sq_l = bufs.best_y_sum_sqs_l[k];
        let y_sum_r = y_sums[k] - y_sum_l;
        let y_sum_sq_r = y_sum_sqs[k] - y_sum_sq_l;
        let mean_l = y_sum_l / nl;
        let var_l = f32::max(0.0, y_sum_sq_l / nl - mean_l * mean_l);
        let mean_r = y_sum_r / nr;
        let var_r = f32::max(0.0, y_sum_sq_r / nr - mean_r * mean_r);
        let reduction = bufs.parent_vars[k] - (nl / nf) * var_l - (nr / nf) * var_r;
        importances[imp_base + k] += weight * f32::max(0.0, reduction);
    }

    let mut left_y_sums = [0.0f32; MULTI_OUTPUT_BATCH];
    let mut left_y_sum_sqs = [0.0f32; MULTI_OUTPUT_BATCH];
    let mut right_y_sums = [0.0f32; MULTI_OUTPUT_BATCH];
    let mut right_y_sum_sqs = [0.0f32; MULTI_OUTPUT_BATCH];

    left_y_sums[..n_targets].copy_from_slice(&bufs.best_y_sums_l[..n_targets]);
    left_y_sum_sqs[..n_targets].copy_from_slice(&bufs.best_y_sum_sqs_l[..n_targets]);
    for k in 0..n_targets {
        right_y_sums[k] = y_sums[k] - left_y_sums[k];
        right_y_sum_sqs[k] = y_sum_sqs[k] - left_y_sum_sqs[k];
    }

    let tf_col = x.get_col(best_feature);
    let mut l_idx = 0usize;
    let mut r_idx = 0usize;

    for i in 0..n {
        let s = sample_slice[i];
        let val = tf_col[s as usize];
        let src_base = i * n_targets;

        if val <= best_threshold_u8 {
            bufs.left_buf[l_idx] = s;
            let dst = l_idx * n_targets;
            bufs.left_y_buf[dst..dst + n_targets]
                .copy_from_slice(&y_slice[src_base..src_base + n_targets]);
            l_idx += 1;
        } else {
            bufs.right_buf[r_idx] = s;
            let dst = r_idx * n_targets;
            bufs.right_y_buf[dst..dst + n_targets]
                .copy_from_slice(&y_slice[src_base..src_base + n_targets]);
            r_idx += 1;
        }
    }

    sample_slice[..l_idx].copy_from_slice(&bufs.left_buf[..l_idx]);
    sample_slice[l_idx..].copy_from_slice(&bufs.right_buf[..r_idx]);

    let y_left_len = l_idx * n_targets;
    let y_right_len = r_idx * n_targets;
    y_slice[..y_left_len].copy_from_slice(&bufs.left_y_buf[..y_left_len]);
    y_slice[y_left_len..y_left_len + y_right_len].copy_from_slice(&bufs.right_y_buf[..y_right_len]);

    let (left_samples, right_samples) = sample_slice.split_at_mut(l_idx);
    let (left_y, right_y) = y_slice.split_at_mut(y_left_len);

    build_node_multi(
        left_y,
        x,
        left_samples,
        &left_y_sums[..n_targets],
        &left_y_sum_sqs[..n_targets],
        n_total,
        n_targets,
        n_features_split,
        config,
        depth + 1,
        importances,
        bufs,
        rng,
    );
    build_node_multi(
        right_y,
        x,
        right_samples,
        &right_y_sums[..n_targets],
        &right_y_sum_sqs[..n_targets],
        n_total,
        n_targets,
        n_features_split,
        config,
        depth + 1,
        importances,
        bufs,
        rng,
    );
}

//////////////////
// Core fitting //
//////////////////

/// Fit a tree ensemble for a batch of target genes simultaneously.
///
/// ### Params
///
/// * `targets` - Sparse target gene expression columns for this batch.
/// * `feature_matrix` - Quantised TF feature store shared across all targets.
/// * `n_samples` - Number of cells.
/// * `config` - Learner configuration.
/// * `seed` - Base seed for reproducibility; each tree gets a derived seed.
///
/// ### Returns
///
/// One importance vector per target: `result[target_idx][feature_idx]`.
/// Importance values are normalised to sum to 1.0 per target.
#[allow(dead_code)]
fn fit_multi_trees(
    targets: &[SparseAxis<u32, f32>],
    feature_matrix: &QuantisedStore,
    n_samples: usize,
    config: &dyn TreeRegressorConfig,
    seed: usize,
) -> Result<Vec<Vec<f32>>, BixverseErrors> {
    let n_features = feature_matrix.n_features;
    let n_targets = targets.len();
    let n_features_split = resolve_n_features_split(config.n_features_split(), n_features);

    let n_sub = if let Some(frac) = config.subsample_frac() {
        ((n_samples as f32 * frac).round() as usize).max(2 * config.min_samples_leaf())
    } else if config.subsample_rate() >= 1.0 {
        n_samples
    } else {
        ((n_samples as f32 * config.subsample_rate()).round() as usize)
            .max(2 * config.min_samples_leaf())
    };

    let y_dense = build_y_dense_multi(targets, n_samples)?;

    let mut sample_indices: Vec<u32> = vec![0; n_samples];
    let mut root_y_buf: Vec<f32> = vec![0.0; n_samples * n_targets];
    let mut bufs = TreeBuffers::new(n_features, n_samples, n_targets);
    let mut importances = vec![0.0f32; n_features * n_targets];
    let mut y_sums_root = vec![0.0f32; n_targets];
    let mut y_sum_sqs_root = vec![0.0f32; n_targets];

    for tree_idx in 0..config.n_trees() {
        let mut rng = SmallRng::seed_from_u64(tree_seed(seed, tree_idx));

        let active_len = if n_sub < n_samples {
            if config.bootstrap() {
                for i in 0..n_sub {
                    sample_indices[i] = rng.random_range(0..n_samples as u32);
                }
                n_sub
            } else {
                init_and_split(&mut sample_indices, n_samples, n_sub, &mut rng);
                n_sub
            }
        } else {
            sample_indices
                .iter_mut()
                .enumerate()
                .for_each(|(i, v)| *v = i as u32);
            n_samples
        };

        let active = &mut sample_indices[..active_len];
        let root_y = &mut root_y_buf[..active_len * n_targets];

        y_sums_root.fill(0.0);
        y_sum_sqs_root.fill(0.0);

        for i in 0..active_len {
            let s = active[i] as usize;
            let src_base = s * n_targets;
            let dst_base = i * n_targets;
            root_y[dst_base..dst_base + n_targets]
                .copy_from_slice(&y_dense[src_base..src_base + n_targets]);
            for k in 0..n_targets {
                let y = root_y[dst_base + k];
                y_sums_root[k] += y;
                y_sum_sqs_root[k] += y * y;
            }
        }

        build_node_multi(
            root_y,
            feature_matrix,
            active,
            &y_sums_root,
            &y_sum_sqs_root,
            active_len,
            n_targets,
            n_features_split,
            config,
            0,
            &mut importances,
            &mut bufs,
            &mut rng,
        );
    }

    Ok(extract_and_normalise_multi_importances(
        &importances,
        n_features,
        n_targets,
    ))
}

/// Fit a tree ensemble for a batch of target genes simultaneously
/// (sparse Y variant).
///
/// ### Params
///
/// * `targets` - Sparse target gene expression columns for this batch.
/// * `feature_matrix` - Quantised TF feature store shared across all targets.
/// * `n_samples` - Number of cells.
/// * `config` - Learner configuration.
/// * `seed` - Base seed for reproducibility; each tree gets a derived seed.
///
/// ### Returns
///
/// One importance vector per target: `result[target_idx][feature_idx]`.
/// Importance values are normalised to sum to 1.0 per target.
pub fn fit_multi_trees_sparse(
    targets: &[SparseAxis<u32, f32>],
    feature_matrix: &QuantisedStore,
    n_samples: usize,
    config: &dyn TreeRegressorConfig,
    seed: usize,
) -> Result<Vec<Vec<f32>>, BixverseErrors> {
    let n_features = feature_matrix.n_features;
    let n_targets = targets.len();
    let n_features_split = resolve_n_features_split(config.n_features_split(), n_features);

    let n_sub = if let Some(frac) = config.subsample_frac() {
        ((n_samples as f32 * frac).round() as usize).max(2 * config.min_samples_leaf())
    } else if config.subsample_rate() >= 1.0 {
        n_samples
    } else {
        ((n_samples as f32 * config.subsample_rate()).round() as usize)
            .max(2 * config.min_samples_leaf())
    };

    let sparse_y = SparseYBatch::from_targets(targets, n_samples)?;

    let mut sample_indices: Vec<u32> = vec![0; n_samples];
    let mut bufs = TreeBuffers::new_sparse(n_features, n_samples, n_targets);
    let mut importances = vec![0.0f32; n_features * n_targets];
    let mut y_sums_root = vec![0.0f32; n_targets];
    let mut y_sum_sqs_root = vec![0.0f32; n_targets];

    for tree_idx in 0..config.n_trees() {
        let mut rng = SmallRng::seed_from_u64(tree_seed(seed, tree_idx));

        let active_len = if n_sub < n_samples {
            if config.bootstrap() {
                for i in 0..n_sub {
                    sample_indices[i] = rng.random_range(0..n_samples as u32);
                }
                n_sub
            } else {
                init_and_split(&mut sample_indices, n_samples, n_sub, &mut rng);
                n_sub
            }
        } else {
            sample_indices
                .iter_mut()
                .enumerate()
                .for_each(|(i, v)| *v = i as u32);
            n_samples
        };

        let active = &mut sample_indices[..active_len];

        y_sums_root.fill(0.0);
        y_sum_sqs_root.fill(0.0);
        for &s in active.iter() {
            let (tgt_idx, tgt_val) = sparse_y.cell_entries(s as usize);
            for i in 0..tgt_idx.len() {
                let k = tgt_idx[i] as usize;
                let y = tgt_val[i];
                y_sums_root[k] += y;
                y_sum_sqs_root[k] += y * y;
            }
        }

        build_node_multi_sparse(
            &sparse_y,
            feature_matrix,
            active,
            &y_sums_root,
            &y_sum_sqs_root,
            active_len,
            n_targets,
            n_features_split,
            config,
            0,
            &mut importances,
            &mut bufs,
            &mut rng,
        );
    }

    Ok(extract_and_normalise_multi_importances(
        &importances,
        n_features,
        n_targets,
    ))
}

/////////
// GBM //
/////////

//////////////////////
// GBM Split Result //
//////////////////////

/// Information about a split found in a node's histograms.
#[derive(Clone, Debug)]
pub struct GbmSplitInfo {
    /// Feature index of the best split.
    pub feature: usize,
    /// Bin threshold; samples with bin <= threshold go left.
    pub threshold: u8,
    /// Weighted variance reduction score.
    pub score: f32,
    /// Number of training samples in the left child.
    pub n_left: u32,
    /// Sum of residuals in the left child.
    pub y_sum_left: f32,
    /// Sum of squared residuals in the left child.
    pub y_sum_sq_left: f32,
}

/////////////////////
// Node Histograms //
/////////////////////

/// Full-feature bin histograms for a single GBM tree node (single target).
///
/// Stores per-bin sample counts, residual sums, and residual sum-of-squares
/// for every feature simultaneously. This enables the histogram subtraction
/// trick: build the smaller child's histogram from scratch, derive the larger
/// child as `parent - smaller` in O(n_features * 256) rather than scanning
/// the larger child's samples.
///
/// All arrays use layout `[feature * 256 + bin]`.
pub struct NodeHistograms {
    /// Bin counts.
    pub counts: Vec<u32>,
    /// Residual sums per bin.
    pub y_sums: Vec<f32>,
    /// Residual sum-of-squares per bin.
    pub y_sum_sqs: Vec<f32>,
    /// Number of features.
    pub n_features: usize,
}

impl NodeHistograms {
    /// Allocate zeroed histograms for the given number of features.
    ///
    /// ### Params
    ///
    /// * `n_features` - Number of features
    ///
    /// ### Returns
    ///
    /// Initialised self
    pub fn new(n_features: usize) -> Self {
        let n = n_features * 256;
        Self {
            counts: vec![0u32; n],
            y_sums: vec![0.0f32; n],
            y_sum_sqs: vec![0.0f32; n],
            n_features,
        }
    }

    /// Zero all bins.
    pub fn reset(&mut self) {
        let n = self.n_features * 256;
        self.counts[..n].fill(0);
        self.y_sums[..n].fill(0.0);
        self.y_sum_sqs[..n].fill(0.0);
    }

    /// Build histograms from a set of sample indices.
    ///
    /// Iterates feature-major for cache-friendly access to the quantised
    /// feature columns. The histogram accumulation region (256 bins * 12
    /// bytes) fits comfortably in L1.
    ///
    /// ### Params
    ///
    /// * `x` - Quantised feature store.
    /// * `samples` - Indices of the samples to include.
    /// * `residuals` - Dense residual array indexed by cell id.
    pub fn build_from_samples(&mut self, x: &QuantisedStore, samples: &[u32], residuals: &[f32]) {
        let n = self.n_features * 256;
        self.counts[..n].fill(0);
        self.y_sums[..n].fill(0.0);
        self.y_sum_sqs[..n].fill(0.0);

        for f in 0..self.n_features {
            let col = x.get_col(f);
            let base = f * 256;
            for &s in samples {
                let bin = col[s as usize] as usize;
                let idx = base + bin;
                self.counts[idx] += 1;
                let r = residuals[s as usize];
                self.y_sums[idx] += r;
                self.y_sum_sqs[idx] += r * r;
            }
        }
    }

    /// Find the best split across all (or a random subset of) features.
    ///
    /// For each candidate feature, computes prefix sums over the 256 bins
    /// and evaluates every valid threshold. Returns `None` if no improving
    /// split exists.
    ///
    /// ### Params
    ///
    /// * `total_sum` - Sum of residuals in this node.
    /// * `total_sum_sq` - Sum of squared residuals in this node.
    /// * `n_samples` - Number of training samples in this node.
    /// * `min_samples_leaf` - Minimum samples per child.
    /// * `n_features_split` - Features to evaluate; `0` means all.
    /// * `feat_buf` - Scratch buffer for feature permutation (length >=
    ///   n_features).
    /// * `rng` - RNG for feature subsampling.
    ///
    /// ### Returns
    ///
    /// Option of the GbmSplitInfo
    #[allow(clippy::too_many_arguments)]
    pub fn find_best_split(
        &self,
        total_sum: f32,
        total_sum_sq: f32,
        n_samples: u32,
        min_samples_leaf: u32,
        n_features_split: usize,
        feat_buf: &mut [usize],
        rng: &mut SmallRng,
    ) -> Option<GbmSplitInfo> {
        let n_features = self.n_features;
        let k_feats = if n_features_split == 0 || n_features_split >= n_features {
            n_features
        } else {
            n_features_split
        };

        // partial Fisher-Yates to select k_feats features
        for i in 0..k_feats {
            let j = rng.random_range(i..n_features);
            feat_buf.swap(i, j);
        }

        let parent_var = node_variance_f32(total_sum, total_sum_sq, n_samples as usize);
        if parent_var < 1e-10 {
            return None;
        }

        let nf = n_samples as f32;
        let mut best = GbmSplitInfo {
            feature: usize::MAX,
            threshold: 0,
            score: 0.0,
            n_left: 0,
            y_sum_left: 0.0,
            y_sum_sq_left: 0.0,
        };

        for fi in 0..k_feats {
            let feat = feat_buf[fi];
            let base = feat * 256;

            let mut cum_count = 0u32;
            let mut cum_sum = 0.0f32;
            let mut cum_sum_sq = 0.0f32;

            // scan bins 0..254; threshold = bin means val <= bin goes left
            for bin in 0..255usize {
                cum_count += self.counts[base + bin];
                cum_sum += self.y_sums[base + bin];
                cum_sum_sq += self.y_sum_sqs[base + bin];

                let nl = cum_count;
                let nr = n_samples - nl;
                if nl < min_samples_leaf || nr < min_samples_leaf {
                    continue;
                }

                let nlf = nl as f32;
                let nrf = nr as f32;
                let mean_l = cum_sum / nlf;
                let var_l = f32::max(0.0, cum_sum_sq / nlf - mean_l * mean_l);
                let sum_r = total_sum - cum_sum;
                let sum_sq_r = total_sum_sq - cum_sum_sq;
                let mean_r = sum_r / nrf;
                let var_r = f32::max(0.0, sum_sq_r / nrf - mean_r * mean_r);

                let score = parent_var - (nlf / nf) * var_l - (nrf / nf) * var_r;
                if score > best.score {
                    best.score = score;
                    best.feature = feat;
                    best.threshold = bin as u8;
                    best.n_left = nl;
                    best.y_sum_left = cum_sum;
                    best.y_sum_sq_left = cum_sum_sq;
                }
            }
        }

        if best.feature == usize::MAX {
            None
        } else {
            Some(best)
        }
    }
}

/////////////////////
// Histogram Pool  //
/////////////////////

/// Pre-allocated pool of `NodeHistograms` with acquire/release semantics.
///
/// Avoids repeated heap allocation during tree building. The pool is sized
/// for the maximum recursion depth at construction time.
pub struct HistogramPool {
    /// Vector of NodeHistograms
    pub histograms: Vec<NodeHistograms>,
    /// Free pools
    free: Vec<usize>,
}

impl HistogramPool {
    /// Create a pool with `capacity` histograms, each sized for number of
    /// features
    ///
    /// ### Params
    ///
    /// * `capacity` - The capacity of the histogram pool
    /// * `n_features` - Number of features
    ///
    /// ### Returns
    ///
    /// Initialised self
    pub fn new(capacity: usize, n_features: usize) -> Self {
        let histograms = (0..capacity)
            .map(|_| NodeHistograms::new(n_features))
            .collect();
        let free = (0..capacity).rev().collect();
        Self { histograms, free }
    }

    /// Acquire a histogram index from the pool. Panics if the pool is exhausted.
    ///
    /// ### Returns
    ///
    /// Index into `self.histograms` for the acquired histogram.
    pub fn acquire(&mut self) -> usize {
        self.free.pop().expect("histogram pool exhausted")
    }

    /// Return a histogram to the pool.
    ///
    /// ### Params
    ///
    /// * `idx` - Releases the pool
    fn release(&mut self, idx: usize) {
        self.free.push(idx);
    }

    /// Compute `out = parent - child` element-wise for all bins.
    ///
    /// All three indices must be distinct.
    ///
    /// ### Params
    ///
    /// * `parent` - Index of the parent node histogram.
    /// * `child` - Index of the smaller child histogram (built from scratch).
    /// * `out` - Index of the histogram to write the result into.
    fn subtract(&mut self, parent: usize, child: usize, out: usize) {
        debug_assert_ne!(parent, out);
        debug_assert_ne!(child, out);
        let n = self.histograms[0].n_features * 256;
        for i in 0..n {
            let pc = self.histograms[parent].counts[i];
            let cc = self.histograms[child].counts[i];
            let py = self.histograms[parent].y_sums[i];
            let cy = self.histograms[child].y_sums[i];
            let pys = self.histograms[parent].y_sum_sqs[i];
            let cys = self.histograms[child].y_sum_sqs[i];
            self.histograms[out].counts[i] = pc - cc;
            self.histograms[out].y_sums[i] = py - cy;
            self.histograms[out].y_sum_sqs[i] = pys - cys;
        }
    }
}

////////////////
// GBM buffer //
////////////////

/// Reusable scratch buffers for the sampled (per-feature) GBM path.
///
/// Combines partition buffers, feature permutation, and single-feature
/// histogram arrays (256 bins on the stack). Used by
/// `build_gbm_node_sampled` where histograms are built one feature at
/// a time rather than all features simultaneously.
struct SampledGbmScratch {
    /// Feature permutation buffer for partial Fisher-Yates.
    feat_buf: Vec<usize>,
    /// Partition buffers for train and OOB sample sets.
    partitions: PartitionBuffers,
    /// Per-bin sample counts (single feature).
    counts: [u32; 256],
    /// Per-bin residual sums (single feature).
    y_sums: [f32; 256],
    /// Per-bin residual sum-of-squares (single feature).
    y_sum_sqs: [f32; 256],
}

impl SampledGbmScratch {
    /// Allocate scratch buffers for the given problem size.
    ///
    /// ### Params
    ///
    /// * `n_features` - Number of features.
    /// * `n_samples` - Maximum number of samples.
    ///
    /// ### Returns
    ///
    /// Initialised `SampledGbmScratch`.
    fn new(n_features: usize, n_samples: usize) -> Self {
        Self {
            feat_buf: (0..n_features).collect(),
            partitions: PartitionBuffers::new(n_samples),
            counts: [0u32; 256],
            y_sums: [0.0f32; 256],
            y_sum_sqs: [0.0f32; 256],
        }
    }

    /// Build a single-feature histogram over the given samples.
    ///
    /// Only the active bin range is zeroed and accumulated, avoiding
    /// a full 256-bin clear on every call.
    ///
    /// ### Params
    ///
    /// * `tf_col` - Quantised feature column for all cells.
    /// * `samples` - Active sample indices.
    /// * `residuals` - Dense residual array indexed by cell id.
    ///
    /// ### Returns
    ///
    /// `(min_bin, max_bin)` of the active range.
    #[inline]
    fn build_histogram(
        &mut self,
        tf_col: &[u8],
        samples: &[u32],
        residuals: &[f32],
    ) -> (usize, usize) {
        if samples.is_empty() {
            return (0, 0);
        }

        let mut min_bin = 255u8;
        let mut max_bin = 0u8;
        for &s in samples {
            let bin = tf_col[s as usize];
            if bin < min_bin {
                min_bin = bin;
            }
            if bin > max_bin {
                max_bin = bin;
            }
        }

        let min_b = min_bin as usize;
        let max_b = max_bin as usize;

        self.counts[min_b..=max_b].fill(0);
        self.y_sums[min_b..=max_b].fill(0.0);
        self.y_sum_sqs[min_b..=max_b].fill(0.0);

        for &s in samples {
            let bin = tf_col[s as usize] as usize;
            let r = residuals[s as usize];
            self.counts[bin] += 1;
            self.y_sums[bin] += r;
            self.y_sum_sqs[bin] += r * r;
        }

        (min_b, max_b)
    }

    /// Scan prefix sums over the active bin range and return the
    /// best split threshold.
    ///
    /// ### Params
    ///
    /// * `min_b` - Minimum active bin.
    /// * `max_b` - Maximum active bin.
    /// * `total_sum` - Sum of residuals in this node.
    /// * `total_sum_sq` - Sum of squared residuals in this node.
    /// * `n_samples` - Number of training samples in this node.
    /// * `min_samples_leaf` - Minimum samples per child.
    /// * `parent_var` - Pre-computed parent variance.
    ///
    /// ### Returns
    ///
    /// `Some((threshold, score, n_left, y_sum_left, y_sum_sq_left))`
    /// if an improving split was found, `None` otherwise.
    #[allow(clippy::too_many_arguments)]
    fn find_best_threshold(
        &self,
        min_b: usize,
        max_b: usize,
        total_sum: f32,
        total_sum_sq: f32,
        n_samples: u32,
        min_samples_leaf: u32,
        parent_var: f32,
    ) -> Option<(u8, f32, u32, f32, f32)> {
        if min_b == max_b {
            return None;
        }

        let nf = n_samples as f32;
        let mut best_score = 0.0f32;
        let mut best_threshold = 0u8;
        let mut best_n_left = 0u32;
        let mut best_sum_l = 0.0f32;
        let mut best_sum_sq_l = 0.0f32;

        let mut cum_count = 0u32;
        let mut cum_sum = 0.0f32;
        let mut cum_sum_sq = 0.0f32;

        for bin in min_b..max_b {
            cum_count += self.counts[bin];
            cum_sum += self.y_sums[bin];
            cum_sum_sq += self.y_sum_sqs[bin];

            let nl = cum_count;
            let nr = n_samples - nl;
            if nl < min_samples_leaf || nr < min_samples_leaf {
                continue;
            }

            let nlf = nl as f32;
            let nrf = nr as f32;
            let mean_l = cum_sum / nlf;
            let var_l = f32::max(0.0, cum_sum_sq / nlf - mean_l * mean_l);
            let sum_r = total_sum - cum_sum;
            let sum_sq_r = total_sum_sq - cum_sum_sq;
            let mean_r = sum_r / nrf;
            let var_r = f32::max(0.0, sum_sq_r / nrf - mean_r * mean_r);

            let score = parent_var - (nlf / nf) * var_l - (nrf / nf) * var_r;
            if score > best_score {
                best_score = score;
                best_threshold = bin as u8;
                best_n_left = nl;
                best_sum_l = cum_sum;
                best_sum_sq_l = cum_sum_sq;
            }
        }

        if best_score > 0.0 {
            Some((
                best_threshold,
                best_score,
                best_n_left,
                best_sum_l,
                best_sum_sq_l,
            ))
        } else {
            None
        }
    }
}

////////////////////
// GBM Node Build //
////////////////////

///////////////
// Full hist //
///////////////

/// Recursively build a single GBM tree node using histogram
/// subtraction.
///
/// At each internal node: find the best split from the precomputed
/// histogram, partition both training and OOB samples, build the
/// smaller child's histogram from scratch, derive the larger via
/// subtraction, and recurse. At leaves: compute the mean residual
/// as the leaf prediction, accumulate OOB improvement, then update
/// residuals for all samples.
///
/// ### Params
///
/// * `pool` - Histogram pool for allocation and subtraction.
/// * `node_hist_idx` - Index of this node's histogram in the pool.
/// * `x` - Quantised feature store.
/// * `residuals` - Dense residual array (updated in place at leaves).
/// * `train_samples` - Training sample indices for this node
///   (partitioned in place).
/// * `oob_samples` - OOB sample indices for this node (partitioned
///   in place).
/// * `y_sum` - Sum of training residuals in this node.
/// * `y_sum_sq` - Sum of squared training residuals in this node.
/// * `n_total_train` - Total training samples at the tree root (for
///   importance weighting).
/// * `config` - GBM configuration.
/// * `depth` - Current tree depth.
/// * `learning_rate` - Shrinkage factor applied to leaf predictions.
/// * `importances` - Per-feature importance accumulator.
/// * `oob_improvement` - Accumulated OOB squared-error improvement.
/// * `scratch` - Reusable non-histogram buffers.
/// * `rng` - Per-tree RNG.
#[allow(clippy::too_many_arguments)]
pub fn build_gbm_node(
    pool: &mut HistogramPool,
    node_hist_idx: usize,
    x: &QuantisedStore,
    residuals: &mut [f32],
    train_samples: &mut [u32],
    oob_samples: &mut [u32],
    y_sum: f32,
    y_sum_sq: f32,
    n_total_train: usize,
    config: &GradientBoostingConfig,
    depth: usize,
    learning_rate: f32,
    importances: &mut [f32],
    oob_improvement: &mut f32,
    scratch: &mut GbmScratch,
    rng: &mut SmallRng,
) {
    let n_train = train_samples.len();

    if n_train < 2 * config.min_samples_leaf || depth >= config.max_depth {
        apply_regression_leaf(
            residuals,
            train_samples,
            oob_samples,
            y_sum,
            n_train,
            learning_rate,
            oob_improvement,
        );
        pool.release(node_hist_idx);
        return;
    }

    let parent_var = node_variance(y_sum, y_sum_sq, n_train);
    if parent_var < 1e-10 {
        apply_regression_leaf(
            residuals,
            train_samples,
            oob_samples,
            y_sum,
            n_train,
            learning_rate,
            oob_improvement,
        );
        pool.release(node_hist_idx);
        return;
    }

    let split = {
        let hist = &pool.histograms[node_hist_idx];
        hist.find_best_split(
            y_sum,
            y_sum_sq,
            n_train as u32,
            config.min_samples_leaf as u32,
            config.n_features_split,
            &mut scratch.feat_buf,
            rng,
        )
    };

    let split = match split {
        Some(s) => s,
        None => {
            apply_regression_leaf(
                residuals,
                train_samples,
                oob_samples,
                y_sum,
                n_train,
                learning_rate,
                oob_improvement,
            );
            pool.release(node_hist_idx);
            return;
        }
    };

    // Accumulate importance
    let nl = split.n_left as f32;
    let nr = (n_train as u32 - split.n_left) as f32;
    let nf = n_train as f32;
    let weight = nf / n_total_train as f32;

    let mean_l = split.y_sum_left / nl;
    let var_l = f32::max(0.0, split.y_sum_sq_left / nl - mean_l * mean_l);
    let y_sum_right = y_sum - split.y_sum_left;
    let y_sum_sq_right = y_sum_sq - split.y_sum_sq_left;
    let mean_r = y_sum_right / nr;
    let var_r = f32::max(0.0, y_sum_sq_right / nr - mean_r * mean_r);
    let reduction = parent_var - (nl / nf) * var_l - (nr / nf) * var_r;
    importances[split.feature] += weight * f32::max(0.0, reduction);

    // Partition
    let tf_col = x.get_col(split.feature);
    let tl = scratch
        .partitions
        .partition_train(train_samples, tf_col, split.threshold);
    let ol = scratch
        .partitions
        .partition_oob(oob_samples, tf_col, split.threshold);

    let (left_train, right_train) = train_samples.split_at_mut(tl);
    let (left_oob, right_oob) = oob_samples.split_at_mut(ol);

    // Histogram subtraction: build smaller child, derive larger
    let left_is_smaller = tl <= (n_train - tl);
    let smaller_idx = pool.acquire();
    let larger_idx = pool.acquire();

    if left_is_smaller {
        pool.histograms[smaller_idx].build_from_samples(x, left_train, residuals);
        pool.subtract(node_hist_idx, smaller_idx, larger_idx);
    } else {
        pool.histograms[smaller_idx].build_from_samples(x, right_train, residuals);
        pool.subtract(node_hist_idx, smaller_idx, larger_idx);
    }

    pool.release(node_hist_idx);

    let (left_hist_idx, right_hist_idx) = if left_is_smaller {
        (smaller_idx, larger_idx)
    } else {
        (larger_idx, smaller_idx)
    };

    build_gbm_node(
        pool,
        left_hist_idx,
        x,
        residuals,
        left_train,
        left_oob,
        split.y_sum_left,
        split.y_sum_sq_left,
        n_total_train,
        config,
        depth + 1,
        learning_rate,
        importances,
        oob_improvement,
        scratch,
        rng,
    );
    build_gbm_node(
        pool,
        right_hist_idx,
        x,
        residuals,
        right_train,
        right_oob,
        y_sum_right,
        y_sum_sq_right,
        n_total_train,
        config,
        depth + 1,
        learning_rate,
        importances,
        oob_improvement,
        scratch,
        rng,
    );
}

/////////////
// Sampled //
/////////////

/// Recursively build a single GBM tree node (single target,
/// per-feature histogram scanning).
///
/// Uses per-node feature-sampled histograms built one feature at a
/// time (stack-sized arrays), with in-place residual updates at
/// leaves. No tree structure is stored; importance and OOB
/// improvement are the sole outputs.
///
/// ### Params
///
/// * `x` - Quantised feature store.
/// * `residuals` - Dense residual array (updated in place at leaves).
/// * `train_samples` - Training sample indices for this node
///   (partitioned in place).
/// * `oob_samples` - OOB sample indices for this node (partitioned
///   in place).
/// * `y_sum` - Sum of training residuals in this node.
/// * `y_sum_sq` - Sum of squared training residuals in this node.
/// * `n_total_train` - Total training samples at the tree root (for
///   importance weighting).
/// * `n_features_split` - Number of features to sample per node.
/// * `config` - GBM configuration.
/// * `depth` - Current tree depth.
/// * `learning_rate` - Shrinkage factor.
/// * `importances` - Per-feature importance accumulator.
/// * `oob_improvement` - Accumulated OOB improvement (updated at
///   leaves).
/// * `bufs` - Reusable scratch buffers.
/// * `rng` - Per-tree RNG.
#[allow(clippy::too_many_arguments)]
fn build_gbm_node_sampled(
    x: &QuantisedStore,
    residuals: &mut [f32],
    train_samples: &mut [u32],
    oob_samples: &mut [u32],
    y_sum: f32,
    y_sum_sq: f32,
    n_total_train: usize,
    n_features_split: usize,
    config: &GradientBoostingConfig,
    depth: usize,
    learning_rate: f32,
    importances: &mut [f32],
    oob_improvement: &mut f32,
    bufs: &mut SampledGbmScratch,
    rng: &mut SmallRng,
) {
    let n_train = train_samples.len();
    let parent_var = node_variance(y_sum, y_sum_sq, n_train);

    if n_train < 2 * config.min_samples_leaf || parent_var < 1e-10 || depth >= config.max_depth {
        apply_regression_leaf(
            residuals,
            train_samples,
            oob_samples,
            y_sum,
            n_train,
            learning_rate,
            oob_improvement,
        );
        return;
    }

    let n_features = x.n_features;
    let k_feats = sample_features(&mut bufs.feat_buf, n_features, n_features_split, rng);

    let mut best_score = 0.0f32;
    let mut best_feature = usize::MAX;
    let mut best_threshold = 0u8;
    let mut best_n_left = 0u32;
    let mut best_sum_l = 0.0f32;
    let mut best_sum_sq_l = 0.0f32;

    for fi in 0..k_feats {
        let feat = bufs.feat_buf[fi];
        let tf_col = x.get_col(feat);

        let (min_b, max_b) = bufs.build_histogram(tf_col, train_samples, residuals);

        if let Some((threshold, score, n_left, sum_l, sum_sq_l)) = bufs.find_best_threshold(
            min_b,
            max_b,
            y_sum,
            y_sum_sq,
            n_train as u32,
            config.min_samples_leaf as u32,
            parent_var,
        ) && score > best_score
        {
            best_score = score;
            best_feature = feat;
            best_threshold = threshold;
            best_n_left = n_left;
            best_sum_l = sum_l;
            best_sum_sq_l = sum_sq_l;
        }
    }

    if best_feature == usize::MAX {
        apply_regression_leaf(
            residuals,
            train_samples,
            oob_samples,
            y_sum,
            n_train,
            learning_rate,
            oob_improvement,
        );
        return;
    }

    // Accumulate importance
    let nl = best_n_left as f32;
    let nr = (n_train as u32 - best_n_left) as f32;
    let nf = n_train as f32;
    let weight = nf / n_total_train as f32;

    let mean_l = best_sum_l / nl;
    let var_l = f32::max(0.0, best_sum_sq_l / nl - mean_l * mean_l);
    let sum_r = y_sum - best_sum_l;
    let sum_sq_r = y_sum_sq - best_sum_sq_l;
    let mean_r = sum_r / nr;
    let var_r = f32::max(0.0, sum_sq_r / nr - mean_r * mean_r);
    let reduction = parent_var - (nl / nf) * var_l - (nr / nf) * var_r;
    importances[best_feature] += weight * f32::max(0.0, reduction);

    // Partition
    let tf_col = x.get_col(best_feature);
    let tl = bufs
        .partitions
        .partition_train(train_samples, tf_col, best_threshold);
    let ol = bufs
        .partitions
        .partition_oob(oob_samples, tf_col, best_threshold);

    let (left_train, right_train) = train_samples.split_at_mut(tl);
    let (left_oob, right_oob) = oob_samples.split_at_mut(ol);

    let right_sum = y_sum - best_sum_l;
    let right_sum_sq = y_sum_sq - best_sum_sq_l;

    build_gbm_node_sampled(
        x,
        residuals,
        left_train,
        left_oob,
        best_sum_l,
        best_sum_sq_l,
        n_total_train,
        n_features_split,
        config,
        depth + 1,
        learning_rate,
        importances,
        oob_improvement,
        bufs,
        rng,
    );
    build_gbm_node_sampled(
        x,
        residuals,
        right_train,
        right_oob,
        right_sum,
        right_sum_sq,
        n_total_train,
        n_features_split,
        config,
        depth + 1,
        learning_rate,
        importances,
        oob_improvement,
        bufs,
        rng,
    );
}

//////////////////
// Core fitting //
//////////////////

/// Fit a GRNBoost2-style gradient boosted ensemble for a single
/// target gene using full-feature histogram subtraction.
///
/// Builds trees sequentially, each fitting the current residuals.
/// Each tree uses a random train/OOB split controlled by
/// `config.subsample_rate`. Early stopping triggers when the
/// rolling average of OOB improvements over the last
/// `config.early_stop_window` trees drops to zero or below.
///
/// ### Params
///
/// * `target` - Sparse target gene expression column.
/// * `feature_matrix` - Quantised TF feature store.
/// * `n_samples` - Number of cells.
/// * `config` - GBM configuration.
/// * `seed` - Base seed for reproducibility.
///
/// ### Returns
///
/// Normalised importance vector of length `n_features`, summing
/// to 1.0.
pub fn fit_grnboost2_full_hist(
    target: &SparseAxis<u32, f32>,
    feature_matrix: &QuantisedStore,
    n_samples: usize,
    config: &GradientBoostingConfig,
    seed: usize,
) -> Result<Vec<f32>, BixverseErrors> {
    let n_features = feature_matrix.n_features;
    let pool_capacity = 2 * config.max_depth + 3;

    let mut residuals = vec![0.0f32; n_samples];
    let (indices, values) = target.get_indices_data_2()?;
    for (i, &idx) in indices.iter().enumerate() {
        residuals[idx] = values[i];
    }

    let n_train = ((n_samples as f32 * config.subsample_rate).round() as usize)
        .max(2 * config.min_samples_leaf);
    let n_oob = n_samples - n_train;

    let mut sample_indices: Vec<u32> = (0..n_samples as u32).collect();
    let mut importances = vec![0.0f32; n_features];
    let mut early_stop = EarlyStopMonitor::new(config.early_stop_window);
    let mut scratch = GbmScratch::new(n_features, n_samples);
    let mut pool = HistogramPool::new(pool_capacity, n_features);

    for tree_idx in 0..config.n_trees_max {
        let mut rng = SmallRng::seed_from_u64(tree_seed(seed, tree_idx));

        init_and_split(&mut sample_indices, n_samples, n_train, &mut rng);
        let (train, oob) = sample_indices.split_at_mut(n_train);

        let root_idx = pool.acquire();
        pool.histograms[root_idx].build_from_samples(feature_matrix, train, &residuals);

        let mut y_sum = 0.0f32;
        let mut y_sum_sq = 0.0f32;
        for &s in train.iter() {
            let r = residuals[s as usize];
            y_sum += r;
            y_sum_sq += r * r;
        }

        let mut oob_improvement = 0.0f32;

        build_gbm_node(
            &mut pool,
            root_idx,
            feature_matrix,
            &mut residuals,
            train,
            oob,
            y_sum,
            y_sum_sq,
            n_train,
            config,
            0,
            config.learning_rate,
            &mut importances,
            &mut oob_improvement,
            &mut scratch,
            &mut rng,
        );

        let oob_mse = if n_oob > 0 {
            oob_improvement / n_oob as f32
        } else {
            0.0
        };

        if early_stop.push(oob_mse) {
            break;
        }
    }

    normalise_importances(&mut importances);

    Ok(importances)
}

/// Fit a GRNBoost2-style gradient boosted ensemble for a single
/// target gene using per-node feature-sampled histograms.
///
/// Per-node histogram construction uses `sqrt(n_features)` sampled
/// features (or `config.n_features_split` if explicitly set). No
/// full-feature histograms or subtraction tricks are used.
///
/// ### Params
///
/// * `target` - Sparse target gene expression column.
/// * `feature_matrix` - Quantised TF feature store.
/// * `n_samples` - Number of cells.
/// * `config` - GBM configuration.
/// * `seed` - Base seed for reproducibility.
///
/// ### Returns
///
/// Normalised importance vector of length `n_features`, summing
/// to 1.0.
pub fn fit_grnboost2_sampled(
    target: &SparseAxis<u32, f32>,
    feature_matrix: &QuantisedStore,
    n_samples: usize,
    config: &GradientBoostingConfig,
    seed: usize,
) -> Result<Vec<f32>, BixverseErrors> {
    let n_features = feature_matrix.n_features;
    let n_features_split = resolve_n_features_split(config.n_features_split, n_features);

    let mut residuals = vec![0.0f32; n_samples];
    let (indices, values) = target.get_indices_data_2()?;
    for (i, &idx) in indices.iter().enumerate() {
        residuals[idx] = values[i];
    }

    let n_train = ((n_samples as f32 * config.subsample_rate).round() as usize)
        .max(2 * config.min_samples_leaf);
    let n_oob = n_samples - n_train;

    let mut sample_indices: Vec<u32> = (0..n_samples as u32).collect();
    let mut importances = vec![0.0f32; n_features];
    let mut early_stop = EarlyStopMonitor::new(config.early_stop_window);
    let mut bufs = SampledGbmScratch::new(n_features, n_samples);

    for tree_idx in 0..config.n_trees_max {
        let mut rng = SmallRng::seed_from_u64(tree_seed(seed, tree_idx));

        init_and_split(&mut sample_indices, n_samples, n_train, &mut rng);
        let (train, oob) = sample_indices.split_at_mut(n_train);

        let mut y_sum = 0.0f32;
        let mut y_sum_sq = 0.0f32;
        for &s in train.iter() {
            let r = residuals[s as usize];
            y_sum += r;
            y_sum_sq += r * r;
        }

        let mut oob_improvement = 0.0f32;

        build_gbm_node_sampled(
            feature_matrix,
            &mut residuals,
            train,
            oob,
            y_sum,
            y_sum_sq,
            n_train,
            n_features_split,
            config,
            0,
            config.learning_rate,
            &mut importances,
            &mut oob_improvement,
            &mut bufs,
            &mut rng,
        );

        let oob_mse = if n_oob > 0 {
            oob_improvement / n_oob as f32
        } else {
            0.0
        };

        if early_stop.push(oob_mse) {
            break;
        }
    }

    normalise_importances(&mut importances);

    Ok(importances)
}

/// Fit a GRNBoost2-style gradient boosted ensemble for a single
/// target gene.
///
/// Dispatches to `fit_grnboost2_sampled` for small cell counts and
/// `fit_grnboost2_full_hist` for large cell counts where histogram
/// subtraction amortises the cost of evaluating all features at
/// every node.
///
/// ### Params
///
/// * `target` - Sparse target gene expression column.
/// * `feature_matrix` - Quantised TF feature store.
/// * `n_samples` - Number of cells.
/// * `config` - GBM configuration.
/// * `seed` - Base seed for reproducibility.
///
/// ### Returns
///
/// Normalised importance vector of length `n_features`, summing
/// to 1.0.
pub fn fit_grnboost2_sparse(
    target: &SparseAxis<u32, f32>,
    feature_matrix: &QuantisedStore,
    n_samples: usize,
    config: &GradientBoostingConfig,
    seed: usize,
) -> Result<Vec<f32>, BixverseErrors> {
    if n_samples >= GBM_SUBTRACTION_CELL_THRESHOLD {
        fit_grnboost2_full_hist(target, feature_matrix, n_samples, config, seed)
    } else {
        fit_grnboost2_sampled(target, feature_matrix, n_samples, config, seed)
    }
}

///////////////////
// Gene batching //
///////////////////

/// Strategy for grouping target genes into multi-output batches
#[derive(Clone, Debug, Default)]
pub enum GeneBatchStrategy {
    /// Shuffle genes randomly before chunking
    #[default]
    Random,
    /// Group genes by expression correlation via SVD + k-means
    /// Fields: (n_svd_components, n_subsampled_cells)
    Correlated {
        /// Number of n_svd
        n_comp: usize,
        /// Number of cells to subsample
        n_cells_subsample: usize,
    },
}

/// Parse the gene batching strategy
///
/// ### Params
///
/// * `s` - String defining the gene batch strategy
/// * `n_comp` - Number of PCs to use for the (randomised) SVD
/// * `n_cells_subsample` - Number of cells to subsample for very large data
///   sets
///
/// ### Returns
///
/// The GeneBatchStrategy option
pub fn parse_gene_batch_strategy(
    s: &str,
    n_comp: usize,
    n_cells_subsample: usize,
) -> Option<GeneBatchStrategy> {
    match s.to_lowercase().as_str() {
        "random" => Some(GeneBatchStrategy::Random),
        "correlated" => Some(GeneBatchStrategy::Correlated {
            n_comp,
            n_cells_subsample,
        }),
        _ => None,
    }
}

/// Chunk a flat gene list into batches of at most `batch_size`.
///
/// ### Params
///
/// * `genes` - Indices of the genes
/// * `batch_size` - Maximum batch size
///
/// ### Returns
///
/// A vector of batches, each containing at most `batch_size` gene indices
fn chunk_into_batches(genes: &[usize], batch_size: usize) -> Vec<Vec<usize>> {
    genes.chunks(batch_size).map(|c| c.to_vec()).collect()
}

/// Group consecutive batches into I/O reads of at most `max_genes` genes,
/// never splitting a batch across reads. Returns [start, end) batch indices.
///
/// ### Params
///
/// * `batches` - Slice of gene batches to group
/// * `max_genes` - Maximum number of genes per I/O group
///
/// ### Returns
///
/// Half-open `[start, end)` index ranges into `batches`, each covering at
/// most `max_genes` genes without splitting a batch across groups
pub(crate) fn group_batches_for_io(
    batches: &[Vec<usize>],
    max_genes: usize,
) -> Vec<(usize, usize)> {
    let mut groups = Vec::new();
    let mut start = 0usize;
    let mut count = 0usize;
    for (i, b) in batches.iter().enumerate() {
        if count > 0 && count + b.len() > max_genes {
            groups.push((start, i));
            start = i;
            count = 0;
        }
        count += b.len();
    }
    if start < batches.len() {
        groups.push((start, batches.len()));
    }
    groups
}

/// Batch genes randomly
///
/// ### Params
///
/// * `gene_indices` - The genes to include
/// * `seed` - Random seed
///
/// ### Returns
///
/// Returns the shuffled indices
pub fn batch_genes_random(gene_indices: &[usize], seed: usize) -> Vec<usize> {
    let mut indices = gene_indices.to_vec();
    let mut rng = SmallRng::seed_from_u64(seed as u64);
    // Fisher-Yates
    for i in (1..indices.len()).rev() {
        let j = rng.random_range(0..=i);
        indices.swap(i, j);
    }
    indices
}

/// Subsampling of cells
///
/// ### Params
///
/// * `cell_indices` - The indices to take forward
/// * `n_target` - Number of cells to return
/// * `seed` - For reproducibility
///
/// ### Returns
///
/// Indices of cells to keep here
pub fn subsample_cells(cell_indices: &[usize], n_target: usize, seed: usize) -> Vec<usize> {
    if cell_indices.len() <= n_target {
        return cell_indices.to_vec();
    }
    let mut indices = cell_indices.to_vec();
    let mut rng = SmallRng::seed_from_u64(seed as u64);
    // Partial Fisher-Yates for n_target elements
    for i in 0..n_target {
        let j = rng.random_range(i..indices.len());
        indices.swap(i, j);
    }
    indices.truncate(n_target);
    indices
}

/// Batch genes by correlation structure
///
/// Runs randomised SVD on a subset of the cells, uses the gene loadings
/// to cluster the genes together into sensible batches of genes.
///
/// ### Params
///
/// * `reader` - Reader over the gene-based count store
/// * `gene_indices` - The genes to include
/// * `batch_size` - Size of the batches
/// * `n_components` - Number of components to use for the clustering
/// * `n_cells_subsample` - Number of cells to subsample
/// * `seed` - Seed for reproducibility
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// Returns gene indices reordered so that co-expressed genes are contiguous for
/// subsequent batching
#[allow(clippy::too_many_arguments)]
fn batch_genes_correlated<S: SingleCellReading>(
    reader: &S,
    gene_indices: &[usize],
    cell_indices: &[usize],
    batch_size: usize,
    n_components: usize,
    n_cells_subsample: usize,
    seed: usize,
    verbose: usize,
) -> Result<Vec<Vec<usize>>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let n_genes = gene_indices.len();
    let n_centroids = n_genes.div_ceil(batch_size);

    let sub_cells = subsample_cells(cell_indices, n_cells_subsample, seed);

    if verbosity.normal_verbosity() {
        println!(
            "Computing gene loadings: {} genes, {} subsampled cells, {} components",
            n_genes,
            sub_cells.len(),
            n_components
        );
    }

    let pca_params = SingleCellPcaParams::new(true, true, true, false, 1e4);

    // loadings is (n_genes, n_components)
    // using a streaming version here to avoid memory blowing up
    let (_, loadings, _, _) = pca_on_sc_streaming(
        reader,
        &sub_cells,
        gene_indices,
        n_components,
        &pca_params,
        None,
        seed,
        false,
        SCENIC_GENE_CHUNK_SIZE,
        verbose,
    )?;

    // flatten to row-major for k-means: gene g, component c -> [g * dim + c]
    let dim = loadings.ncols();
    let mut gene_loadings = vec![0.0f32; n_genes * dim];
    for g in 0..n_genes {
        for c in 0..dim {
            gene_loadings[g * dim + c] = loadings[(g, c)];
        }
    }

    if verbosity.normal_verbosity() {
        println!("Clustering {} genes into {} groups", n_genes, n_centroids);
    }

    let k_means_params = KMeansParamsWrappers::new(50, None, None);

    let centroids = train_centroids(
        &gene_loadings,
        dim,
        n_genes,
        n_centroids,
        &Dist::SquaredEuclidean,
        Some(k_means_params.get_data()),
        seed,
        verbosity.detailed_verbosity(),
    )?;

    let centroid_norms: Vec<f32> = (0..n_centroids)
        .map(|i| {
            let c = &centroids[i * dim..(i + 1) * dim];
            f32::dot_simd(c, c)
        })
        .collect();

    let data_norms: Vec<f32> = (0..n_genes)
        .map(|i| {
            let v = &gene_loadings[i * dim..(i + 1) * dim];
            f32::dot_simd(v, v)
        })
        .collect();

    let assignments = assign_all_parallel(
        &gene_loadings,
        &data_norms,
        dim,
        n_genes,
        &centroids,
        &centroid_norms,
        n_centroids,
        &Dist::SquaredEuclidean,
    );

    let mut clusters: Vec<Vec<usize>> = vec![Vec::new(); n_centroids];
    for (i, &cluster_id) in assignments.iter().enumerate() {
        clusters[cluster_id].push(gene_indices[i]);
    }

    let mut rng = SmallRng::seed_from_u64(seed.wrapping_add(1) as u64);
    let mut batches: Vec<Vec<usize>> = Vec::with_capacity(n_centroids);
    for cluster in &mut clusters {
        if cluster.is_empty() {
            continue;
        }
        for i in (1..cluster.len()).rev() {
            let j = rng.random_range(0..=i);
            cluster.swap(i, j);
        }
        // split into <= batch_size pieces so a batch never straddles a cluster
        for chunk in cluster.chunks(batch_size) {
            batches.push(chunk.to_vec());
        }
    }

    Ok(batches)
}

/// Reorder gene indices into batches of `batch_size` according to the chosen
/// strategy.
///
/// ### Params
///
/// * `reader` - Reader over the sparse expression store.
/// * `gene_indices` - Target gene indices to batch.
/// * `cell_indices` - Cell indices in use.
/// * `batch_size` - Target batch size (MULTI_OUTPUT_BATCH).
/// * `strategy` - Batching strategy.
/// * `seed` - RNG seed.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// The gene indices partitioned into batches of at most `batch_size`,
/// ordered according to the chosen strategy
pub fn batch_genes<S: SingleCellReading>(
    reader: &S,
    gene_indices: &[usize],
    cell_indices: &[usize],
    batch_size: usize,
    strategy: &GeneBatchStrategy,
    seed: usize,
    verbose: usize,
) -> Result<Vec<Vec<usize>>, BixverseErrors> {
    match strategy {
        GeneBatchStrategy::Random => Ok(chunk_into_batches(
            &batch_genes_random(gene_indices, seed),
            batch_size,
        )),
        GeneBatchStrategy::Correlated {
            n_comp,
            n_cells_subsample,
        } => {
            if gene_indices.len() <= batch_size {
                return Ok(chunk_into_batches(
                    &batch_genes_random(gene_indices, seed),
                    batch_size,
                ));
            }
            batch_genes_correlated(
                reader,
                gene_indices,
                cell_indices,
                batch_size,
                *n_comp,
                (*n_cells_subsample).min(cell_indices.len()),
                seed,
                verbose,
            )
        }
    }
}

/////////////////////
// Fit type choice //
/////////////////////

/// Multi-output RF/ET path: genes are batched and fitted with shared tree
/// structure across targets within each batch.
///
/// ### Params
///
/// * `reader` - Reader for target gene I/O.
/// * `cell_set` - Set of active cell IDs for filtering.
/// * `cell_indices` - Cell indices as a slice (forwarded to `batch_genes`).
/// * `gene_indices` - Target gene indices.
/// * `tf_data` - Quantised TF feature store.
/// * `n_cells` - Number of active cells.
/// * `n_tfs` - Number of TFs (features).
/// * `n_genes` - Number of target genes.
/// * `scenic_params` - SCENIC configuration.
/// * `seed` - Base random seed.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
/// * `start_total` - Timer from the top-level call for elapsed reporting.
///
/// ### Returns
///
/// A `Mat<f32>` of shape `(n_genes, n_tfs)` containing normalised
/// per-target feature importances.
#[allow(clippy::too_many_arguments)]
fn run_scenic_multi_output<S: SingleCellReading>(
    reader: &S,
    cell_set: &IndexSet<u32>,
    cell_indices: &[usize],
    gene_indices: &[usize],
    tf_data: &QuantisedStore,
    n_cells: usize,
    n_tfs: usize,
    n_genes: usize,
    scenic_params: &ScenicParams,
    seed: usize,
    verbose: usize,
    start_total: Instant,
) -> Result<Mat<f32>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let n_multi_output = scenic_params
        .gene_batch_size
        .unwrap_or(MULTI_OUTPUT_BATCH)
        .min(MULTI_OUTPUT_BATCH);

    let strategy = parse_gene_batch_strategy(
        &scenic_params.gene_batch_strategy,
        scenic_params.n_pcs,
        scenic_params.n_subsample,
    )
    .unwrap_or(GeneBatchStrategy::Random);

    let batches = batch_genes(
        reader,
        gene_indices,
        cell_indices,
        n_multi_output,
        &strategy,
        seed,
        verbose,
    )?;
    let ordered_genes: Vec<usize> = batches.iter().flatten().copied().collect();

    let gene_id_to_pos: FxHashMap<usize, usize> = gene_indices
        .iter()
        .enumerate()
        .map(|(pos, &gid)| (gid, pos))
        .collect();

    let start_gene_read = Instant::now();
    let mut all_sparse_cols: Vec<SparseAxis<u32, f32>> = Vec::with_capacity(n_genes);

    for (iter, chunk) in ordered_genes.chunks(SCENIC_GENE_CHUNK_SIZE).enumerate() {
        let mut gene_chunks: Vec<CscGeneChunk> = reader.read_gene_parallel(chunk)?;
        gene_chunks.par_iter_mut().for_each(|c| {
            c.filter_selected_cells(cell_set);
        });

        for gc in gene_chunks.iter() {
            all_sparse_cols.push(gc.to_sparse_axis(n_cells));
        }

        if verbosity.detailed_verbosity() {
            println!(
                "  Read gene chunk {}/{} ({} genes)",
                iter + 1,
                ordered_genes.len().div_ceil(SCENIC_GENE_CHUNK_SIZE),
                all_sparse_cols.len(),
            );
        }
    }

    if verbosity.normal_verbosity() {
        println!(
            "Read and filtered {} target genes in {:.2?}",
            n_genes,
            start_gene_read.elapsed()
        );
    }

    // Slice by cluster-bounded batches, not by fixed n_multi_output.
    let mut col_batches: Vec<&[SparseAxis<u32, f32>]> = Vec::with_capacity(batches.len());
    let mut offset = 0usize;
    for b in &batches {
        col_batches.push(&all_sparse_cols[offset..offset + b.len()]);
        offset += b.len();
    }
    let total_batches = col_batches.len();

    let config: &dyn TreeRegressorConfig = match &scenic_params.regression_learner {
        RegressionLearner::ExtraTrees(cfg) => cfg,
        RegressionLearner::RandomForest(cfg) => cfg,
        RegressionLearner::GradientBoosting(_) => unreachable!(),
    };

    let learner_name = match &scenic_params.regression_learner {
        RegressionLearner::ExtraTrees(_) => "ExtraTrees",
        RegressionLearner::RandomForest(_) => "RandomForest",
        RegressionLearner::GradientBoosting(_) => unreachable!(),
    };

    if verbosity.normal_verbosity() {
        println!(
            "Running SCENIC ({}) on {} genes ({} TFs, {} cells, {} batches of up to {})",
            learner_name, n_genes, n_tfs, n_cells, total_batches, n_multi_output,
        );
    }

    let start_fit = Instant::now();
    let batches_done = AtomicUsize::new(0);

    let batch_results: Vec<(usize, Vec<Vec<f32>>)> = (0..total_batches)
        .into_par_iter()
        .map(|batch_idx| {
            let batch_seed = seed.wrapping_add(batch_idx.wrapping_mul(2654435761));
            let imp = fit_multi_trees_sparse(
                col_batches[batch_idx],
                tf_data,
                n_cells,
                config,
                batch_seed,
            )?;

            if verbosity.normal_verbosity() {
                let done = batches_done.fetch_add(1, Ordering::Relaxed) + 1;
                let pct = done * 100 / total_batches;
                let prev_pct = (done - 1) * 100 / total_batches;
                if pct / 10 > prev_pct / 10 || done == total_batches {
                    println!(
                        "  Progress: {}% ({}/{} batches, {:.2?} elapsed)",
                        pct,
                        done,
                        total_batches,
                        start_fit.elapsed()
                    );
                }
            }

            Ok((batch_idx, imp))
        })
        .collect::<Result<Vec<_>, BixverseErrors>>()?;

    let mut importance_scores: Vec<Vec<f32>> = vec![Vec::new(); n_genes];

    for (batch_idx, imp_vecs) in batch_results {
        let batch_gene_ids = &batches[batch_idx];
        for (local_idx, imp) in imp_vecs.into_iter().enumerate() {
            let gene_id = batch_gene_ids[local_idx];
            let original_pos = gene_id_to_pos[&gene_id];
            importance_scores[original_pos] = imp;
        }
    }

    if verbosity.normal_verbosity() {
        println!(
            "SCENIC ({}) GRN inference complete in {:.2?}",
            learner_name,
            start_total.elapsed()
        );
    }

    Ok(Mat::from_fn(n_genes, n_tfs, |i, j| {
        if j < importance_scores[i].len() {
            importance_scores[i][j]
        } else {
            0.0
        }
    }))
}

/// GBM path: single-target, parallelised across individual genes.
///
/// No gene batching strategy is applied -- each gene is an independent
/// regression. All target genes are read into memory, then fitted via a
/// single `par_iter` over the full gene set.
///
/// ### Params
///
/// * `reader` - Reader for target gene I/O.
/// * `cell_set` - Set of active cell IDs for filtering.
/// * `gene_indices` - Target gene indices.
/// * `tf_data` - Quantised TF feature store.
/// * `n_cells` - Number of active cells.
/// * `n_tfs` - Number of TFs (features).
/// * `n_genes` - Number of target genes.
/// * `config` - GBM configuration.
/// * `seed` - Base random seed.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
/// * `start_total` - Timer from the top-level call for elapsed reporting.
///
/// ### Returns
///
/// A `Mat<f32>` of shape `(n_genes, n_tfs)` containing normalised
/// per-target feature importances.
#[allow(clippy::too_many_arguments)]
fn run_scenic_gbm<S: SingleCellReading>(
    reader: &S,
    cell_set: &IndexSet<u32>,
    gene_indices: &[usize],
    tf_data: &QuantisedStore,
    n_cells: usize,
    n_tfs: usize,
    n_genes: usize,
    config: &GradientBoostingConfig,
    seed: usize,
    verbose: usize,
    start_total: Instant,
) -> Result<Mat<f32>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let start_gene_read = Instant::now();
    let mut all_sparse_cols: Vec<SparseAxis<u32, f32>> = Vec::with_capacity(n_genes);

    if verbosity.normal_verbosity() {
        println!(
            "Running GRNBoost2 on {} genes ({} TFs, {} cells, {} I/O chunks)",
            n_genes.separate_with_underscores(),
            n_tfs.separate_with_underscores(),
            n_cells.separate_with_underscores(),
            gene_indices.len().div_ceil(SCENIC_GENE_CHUNK_SIZE),
        );
    }

    for (iter, chunk) in gene_indices.chunks(SCENIC_GENE_CHUNK_SIZE).enumerate() {
        let mut gene_chunks: Vec<CscGeneChunk> = reader.read_gene_parallel(chunk)?;
        gene_chunks.par_iter_mut().for_each(|c| {
            c.filter_selected_cells(cell_set);
        });

        for gc in &gene_chunks {
            all_sparse_cols.push(gc.to_sparse_axis(n_cells));
        }

        if verbosity.detailed_verbosity() {
            println!(
                "  Read gene chunk {}/{} ({} genes)",
                iter + 1,
                gene_indices.len().div_ceil(SCENIC_GENE_CHUNK_SIZE),
                all_sparse_cols.len(),
            );
        }
    }

    if verbosity.normal_verbosity() {
        println!(
            "Read and filtered {} target genes in {:.2?}",
            n_genes,
            start_gene_read.elapsed()
        );
        println!(
            "Running GRNBoost2 on {} genes ({} TFs, {} cells)",
            n_genes, n_tfs, n_cells,
        );
    }

    let start_fit = Instant::now();
    let genes_done = AtomicUsize::new(0);

    let importance_scores: Vec<Vec<f32>> = all_sparse_cols
        .par_iter()
        .enumerate()
        .map(|(gene_idx, target)| {
            let gene_seed = seed.wrapping_add(gene_idx.wrapping_mul(2654435761));
            let imp = fit_grnboost2_sparse(target, tf_data, n_cells, config, gene_seed)?;

            if verbosity.normal_verbosity() {
                let done = genes_done.fetch_add(1, Ordering::Relaxed) + 1;
                let pct = done * 100 / n_genes;
                let prev_pct = (done - 1) * 100 / n_genes;
                if pct / 10 > prev_pct / 10 || done == n_genes {
                    println!(
                        "  Progress: {}% ({}/{} genes, {:.2?} elapsed)",
                        pct,
                        done,
                        n_genes,
                        start_fit.elapsed()
                    );
                }
            }

            Ok(imp)
        })
        .collect::<Result<Vec<_>, BixverseErrors>>()?;

    if verbosity.normal_verbosity() {
        println!(
            "GRNBoost2 GRN inference complete in {:.2?}",
            start_total.elapsed()
        );
    }

    Ok(Mat::from_fn(n_genes, n_tfs, |i, j| {
        if j < importance_scores[i].len() {
            importance_scores[i][j]
        } else {
            0.0
        }
    }))
}

/// Multi-output RF/ET streaming path: read genes in I/O chunks, slice
/// into multi-output batches within each chunk, fit in parallel, drop
/// before next chunk.
///
/// ### Params
///
/// * `reader` - Reader for target gene I/O.
/// * `cell_set` - Set of active cell IDs for filtering.
/// * `cell_indices` - Cell indices as a slice (forwarded to `batch_genes`).
/// * `gene_indices` - Target gene indices.
/// * `tf_data` - Quantised TF feature store.
/// * `n_cells` - Number of active cells.
/// * `n_tfs` - Number of TFs (features).
/// * `n_genes` - Number of target genes.
/// * `scenic_params` - SCENIC configuration.
/// * `seed` - Base random seed.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
/// * `start_total` - Timer from the top-level call for elapsed reporting.
///
/// ### Returns
///
/// A `Mat<f32>` of shape `(n_genes, n_tfs)` containing normalised
/// per-target feature importances.
#[allow(clippy::too_many_arguments)]
fn run_scenic_multi_output_streaming<S: SingleCellReading>(
    reader: &S,
    cell_set: &IndexSet<u32>,
    cell_indices: &[usize],
    gene_indices: &[usize],
    tf_data: &QuantisedStore,
    n_cells: usize,
    n_tfs: usize,
    n_genes: usize,
    scenic_params: &ScenicParams,
    seed: usize,
    verbose: usize,
    start_total: Instant,
) -> Result<Mat<f32>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let n_multi_output = scenic_params
        .gene_batch_size
        .unwrap_or(MULTI_OUTPUT_BATCH)
        .min(MULTI_OUTPUT_BATCH);

    let strategy = parse_gene_batch_strategy(
        &scenic_params.gene_batch_strategy,
        scenic_params.n_pcs,
        scenic_params.n_subsample,
    )
    .unwrap_or(GeneBatchStrategy::Random);

    let batches = batch_genes(
        reader,
        gene_indices,
        cell_indices,
        n_multi_output,
        &strategy,
        seed,
        verbose,
    )?;
    let io_groups = group_batches_for_io(&batches, SCENIC_GENE_CHUNK_SIZE);

    let gene_id_to_pos: FxHashMap<usize, usize> = gene_indices
        .iter()
        .enumerate()
        .map(|(pos, &gid)| (gid, pos))
        .collect();

    let config: &dyn TreeRegressorConfig = match &scenic_params.regression_learner {
        RegressionLearner::ExtraTrees(cfg) => cfg,
        RegressionLearner::RandomForest(cfg) => cfg,
        RegressionLearner::GradientBoosting(_) => unreachable!(),
    };

    let learner_name = match &scenic_params.regression_learner {
        RegressionLearner::ExtraTrees(_) => "ExtraTrees",
        RegressionLearner::RandomForest(_) => "RandomForest",
        RegressionLearner::GradientBoosting(_) => unreachable!(),
    };

    let total_io_chunks = io_groups.len();
    let mut importance_scores: Vec<Vec<f32>> = vec![Vec::new(); n_genes];
    let mut global_batch_offset: usize = 0;
    let mut genes_processed: usize = 0;

    if verbosity.normal_verbosity() {
        println!(
            "Running SCENIC ({}, streaming) on {} genes ({} TFs, {} cells, {} I/O chunks, batches of {})",
            learner_name,
            n_genes.separate_with_underscores(),
            n_tfs.separate_with_underscores(),
            n_cells.separate_with_underscores(),
            total_io_chunks,
            n_multi_output,
        );
    }

    for (chunk_idx, &(g_start, g_end)) in io_groups.iter().enumerate() {
        let start_chunk = Instant::now();
        let group = &batches[g_start..g_end];
        let io_chunk: Vec<usize> = group.iter().flatten().copied().collect();

        // read and filter chunk
        let start_io = Instant::now();
        let mut gene_chunks: Vec<CscGeneChunk> = reader.read_gene_parallel(&io_chunk)?;
        gene_chunks.par_iter_mut().for_each(|c| {
            c.filter_selected_cells(cell_set);
        });

        let sparse_columns: Vec<SparseAxis<u32, f32>> = gene_chunks
            .iter()
            .map(|c| c.to_sparse_axis(n_cells))
            .collect();
        drop(gene_chunks);

        if verbosity.normal_verbosity() {
            println!(
                "  Chunk {}/{}: loaded and filtered {} genes in {:.2?}",
                chunk_idx + 1,
                total_io_chunks,
                io_chunk.len(),
                start_io.elapsed()
            );
        }

        // offsets into sparse_columns for each batch in this group
        let mut col_offsets = Vec::with_capacity(group.len() + 1);
        let mut off = 0usize;
        col_offsets.push(0);
        for b in group {
            off += b.len();
            col_offsets.push(off);
        }
        let n_batches_this_chunk = group.len();

        let start_fit = Instant::now();
        let batches_done = AtomicUsize::new(0);

        let batch_results: Vec<(usize, Vec<Vec<f32>>)> = (0..n_batches_this_chunk)
            .into_par_iter()
            .map(|local_batch_idx| {
                let batch_seed = seed
                    .wrapping_add((global_batch_offset + local_batch_idx).wrapping_mul(2654435761));
                let cols =
                    &sparse_columns[col_offsets[local_batch_idx]..col_offsets[local_batch_idx + 1]];
                let imp = fit_multi_trees_sparse(cols, tf_data, n_cells, config, batch_seed)?;

                if verbosity.detailed_verbosity() && n_batches_this_chunk >= 4 {
                    let done = batches_done.fetch_add(1, Ordering::Relaxed) + 1;
                    let pct = done * 100 / n_batches_this_chunk;
                    let prev_pct = (done - 1) * 100 / n_batches_this_chunk;
                    if [25, 50, 75, 100].iter().any(|&q| prev_pct < q && pct >= q) {
                        println!(
                            "    Chunk {}: ~{}% of batches done ({}/{}, {:.2?} elapsed)",
                            chunk_idx + 1,
                            pct,
                            done,
                            n_batches_this_chunk,
                            start_fit.elapsed()
                        );
                    }
                }

                Ok((local_batch_idx, imp))
            })
            .collect::<Result<Vec<_>, BixverseErrors>>()?;

        for (local_batch_idx, imp_vecs) in batch_results {
            let batch_gene_ids = &group[local_batch_idx];
            for (local_idx, imp) in imp_vecs.into_iter().enumerate() {
                let gene_id = batch_gene_ids[local_idx];
                let original_pos = gene_id_to_pos[&gene_id];
                importance_scores[original_pos] = imp;
            }
        }

        global_batch_offset += n_batches_this_chunk;
        genes_processed += io_chunk.len();

        if verbosity.normal_verbosity() {
            println!(
                "  Chunk {}/{}: {}/{} genes done in {:.2?} (fit: {:.2?})",
                chunk_idx + 1,
                total_io_chunks,
                genes_processed,
                n_genes,
                start_chunk.elapsed(),
                start_fit.elapsed()
            );
        }
        // sparse_columns dropped here
    }

    if verbosity.normal_verbosity() {
        println!(
            "SCENIC ({}) GRN inference (streaming) complete in {:.2?}",
            learner_name,
            start_total.elapsed()
        );
    }

    Ok(Mat::from_fn(n_genes, n_tfs, |i, j| {
        if j < importance_scores[i].len() {
            importance_scores[i][j]
        } else {
            0.0
        }
    }))
}

/// GBM streaming path: read genes in I/O chunks, fit each gene
/// individually within the chunk, drop chunk data before reading next.
///
/// Memory is bounded to one I/O chunk of sparse columns at a time. Each
/// gene within a chunk is an independent work unit for rayon.
///
/// ### Params
///
/// * `reader` - Reader for target gene I/O.
/// * `cell_set` - Set of active cell IDs for filtering.
/// * `gene_indices` - Target gene indices.
/// * `tf_data` - Quantised TF feature store.
/// * `n_cells` - Number of active cells.
/// * `n_tfs` - Number of TFs (features).
/// * `n_genes` - Number of target genes.
/// * `config` - GBM configuration.
/// * `seed` - Base random seed.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
/// * `start_total` - Timer from the top-level call for elapsed reporting.
///
/// ### Returns
///
/// A `Mat<f32>` of shape `(n_genes, n_tfs)` containing normalised
/// per-target feature importances.
#[allow(clippy::too_many_arguments)]
fn run_scenic_gbm_streaming<S: SingleCellReading>(
    reader: &S,
    cell_set: &IndexSet<u32>,
    gene_indices: &[usize],
    tf_data: &QuantisedStore,
    n_cells: usize,
    n_tfs: usize,
    n_genes: usize,
    config: &GradientBoostingConfig,
    seed: usize,
    verbose: usize,
    start_total: Instant,
) -> Result<Mat<f32>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let total_io_chunks = gene_indices.len().div_ceil(SCENIC_GENE_CHUNK_SIZE);
    let mut importance_scores: Vec<Vec<f32>> = vec![Vec::new(); n_genes];
    let mut global_gene_offset: usize = 0;

    if verbosity.normal_verbosity() {
        println!(
            "Running GRNBoost2 (streaming) on {} genes ({} TFs, {} cells, {} I/O chunks)",
            n_genes.separate_with_underscores(),
            n_tfs.separate_with_underscores(),
            n_cells.separate_with_underscores(),
            total_io_chunks,
        );
    }

    for (chunk_idx, io_chunk) in gene_indices.chunks(SCENIC_GENE_CHUNK_SIZE).enumerate() {
        let start_chunk = Instant::now();

        let start_io = Instant::now();
        let mut gene_chunks: Vec<CscGeneChunk> = reader.read_gene_parallel(io_chunk)?;
        gene_chunks.par_iter_mut().for_each(|c| {
            c.filter_selected_cells(cell_set);
        });

        let sparse_columns: Vec<SparseAxis<u32, f32>> = gene_chunks
            .iter()
            .map(|c| c.to_sparse_axis(n_cells))
            .collect();
        drop(gene_chunks);

        if verbosity.normal_verbosity() {
            println!(
                "  Chunk {}/{}: loaded and filtered {} genes in {:.2?}",
                chunk_idx + 1,
                total_io_chunks,
                io_chunk.len(),
                start_io.elapsed()
            );
        }

        let start_fit = Instant::now();
        let genes_done = AtomicUsize::new(0);
        let n_genes_this_chunk = sparse_columns.len();

        let chunk_results: Vec<Vec<f32>> = sparse_columns
            .par_iter()
            .enumerate()
            .map(|(local_idx, target)| {
                let gene_seed =
                    seed.wrapping_add((global_gene_offset + local_idx).wrapping_mul(2654435761));
                let imp = fit_grnboost2_sparse(target, tf_data, n_cells, config, gene_seed)?;

                if verbosity.detailed_verbosity() && n_genes_this_chunk >= 4 {
                    let done = genes_done.fetch_add(1, Ordering::Relaxed) + 1;
                    let pct = done * 100 / n_genes_this_chunk;
                    let prev_pct = (done - 1) * 100 / n_genes_this_chunk;
                    if [25, 50, 75, 100].iter().any(|&q| prev_pct < q && pct >= q) {
                        println!(
                            "    Chunk {}: ~{}% of genes done ({}/{}, {:.2?} elapsed)",
                            chunk_idx + 1,
                            pct,
                            done,
                            n_genes_this_chunk,
                            start_fit.elapsed()
                        );
                    }
                }

                Ok(imp)
            })
            .collect::<Result<Vec<_>, BixverseErrors>>()?;

        for (local_idx, imp) in chunk_results.into_iter().enumerate() {
            importance_scores[global_gene_offset + local_idx] = imp;
        }

        global_gene_offset += n_genes_this_chunk;

        if verbosity.normal_verbosity() {
            println!(
                "  Chunk {}/{}: {}/{} genes done in {:.2?} (fit: {:.2?})",
                chunk_idx + 1,
                total_io_chunks,
                global_gene_offset,
                n_genes,
                start_chunk.elapsed(),
                start_fit.elapsed()
            );
        }
        // sparse_columns dropped here
    }

    if verbosity.normal_verbosity() {
        println!(
            "GRNBoost2 GRN inference (streaming) complete in {:.2?}",
            start_total.elapsed()
        );
    }

    Ok(Mat::from_fn(n_genes, n_tfs, |i, j| {
        if j < importance_scores[i].len() {
            importance_scores[i][j]
        } else {
            0.0
        }
    }))
}

////////////
// Params //
////////////

/// Defines the parameters to run SCENIC within the bixverse-rs implementation
pub struct ScenicParams {
    /// Min total counts that a gene needs to reach to be included
    pub min_counts: usize,
    /// Min proportion of cells (between 0 and 1) that need to have a gene
    /// expressed to be considered for the analysis
    pub min_cells: f32,
    /// Regression learner - enum defining which regression learner to use
    pub regression_learner: RegressionLearner,
    /// Strategy for gene batching to use
    pub gene_batch_strategy: String,
    /// Optional gene batch size
    pub gene_batch_size: Option<usize>,
    /// Number of PCs to use for correlated gene batch
    pub n_pcs: usize,
    /// Cell subsampling threshold. If n ≥ n_subsample, n_subsample cells will
    /// be randomly selected prior to running randomised SVD for the correlated
    /// gene batch strategy
    pub n_subsample: usize,
}

/// Default implementations of the Scenic parameters
impl Default for ScenicParams {
    fn default() -> Self {
        Self {
            min_counts: 50,
            min_cells: 0.03,
            regression_learner: RegressionLearner::RandomForest(RandomForestConfig::default()),
            gene_batch_strategy: "correlated".to_string(),
            gene_batch_size: None,
            n_pcs: 50,
            n_subsample: 100_000,
        }
    }
}

//////////
// Main //
//////////

/// Shared setup produced by [`scenic_common_setup`]: the artefacts every
/// top-level SCENIC entry point needs before dispatching on the regression
/// learner.
///
/// ### Fields
///
/// * `reader` - Borrowed reader over the sparse expression store.
/// * `cell_set` - Active cell IDs, deduplicated and preserved in input order.
/// * `tf_data` - Quantised TF feature store (`u8` columns, column-major).
/// * `n_cells` - Number of active cells.
/// * `n_tfs` - Number of TFs (features).
/// * `start_total` - Wall-clock timer started at the top of the call.
pub(crate) struct ScenicSetup<'a, S: SingleCellReading> {
    /// Open reader over the sparse expression file.
    pub reader: &'a S,
    /// Active cell IDs, deduplicated and preserved in input order.
    pub cell_set: IndexSet<u32>,
    /// Quantised TF feature store (`u8` columns, column-major).
    pub tf_data: QuantisedStore,
    /// Number of active cells.
    pub n_cells: usize,
    /// Number of TFs (features).
    pub n_tfs: usize,
    /// Wall-clock timer started at the top of the call.
    pub start_total: Instant,
}

/// Perform the prelude shared by every top-level SCENIC entry point: open the
/// reader, build the cell set, read + filter + quantise the TF columns, and
/// print a timing line.
///
/// ### Params
///
/// * `reader` - Reader over the sparse gene expression store.
/// * `cell_indices` - Indices of the active cells.
/// * `tf_indices` - Transcription-factor gene indices.
/// * `verbose` - `0` -> silent, `1` -> normal, `2` -> detailed.
///
/// ### Returns
///
/// A populated [`ScenicSetup`] on success.
pub(crate) fn scenic_common_setup<'a, S: SingleCellReading>(
    reader: &'a S,
    cell_indices: &[usize],
    tf_indices: &[usize],
    verbose: usize,
) -> Result<ScenicSetup<'a, S>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let start_total = Instant::now();
    let cell_set: IndexSet<u32> = cell_indices.iter().map(|&x| x as u32).collect();
    let n_cells = cell_set.len();

    let start_reading = Instant::now();

    let mut tf_chunks: Vec<CscGeneChunk> = reader.read_gene_parallel(tf_indices)?;
    tf_chunks.par_iter_mut().for_each(|chunk| {
        chunk.filter_selected_cells(&cell_set);
    });

    let tf_csc: CompressedSparseData2<u16, f32> =
        from_gene_chunks::<u16>(tf_chunks, &DataLayerReturn::Norm, n_cells)?;
    let tf_data = QuantisedStore::from_csc(&tf_csc, n_cells)?;
    drop(tf_csc);

    let n_tfs = tf_data.n_features;

    if verbosity.normal_verbosity() {
        println!(
            "Loaded, filtered and quantised TF data (n: {}) in: {:.2?}",
            n_tfs.separate_with_underscores(),
            start_reading.elapsed()
        );
    }

    Ok(ScenicSetup {
        reader,
        cell_set,
        tf_data,
        n_cells,
        n_tfs,
        start_total,
    })
}

/// Filter genes by minimum total counts and minimum expressed-cell fraction
///
/// ### Params
///
/// * `reader` - Reader over the sparse gene expression store.
/// * `cell_indices` - Indices of the cells to restrict to.
/// * `scenic_params` - Reference to the SCENIC parameters indicating minimum
///   couts per gene and minimum proportions of cells expressing a gene to
///   be included.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// Indices of genes passing both filters, in their original order.
pub fn scenic_gene_filter<S: SingleCellReading>(
    reader: &S,
    cell_indices: &[usize],
    scenic_params: &ScenicParams,
    verbose: usize,
) -> Result<Vec<usize>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let total_genes = reader.get_header().total_genes;
    let all_gene_indices: Vec<usize> = (0..total_genes).collect();
    let cell_set: IndexSet<u32> = cell_indices.iter().map(|&x| x as u32).collect();
    let n_cells = cell_indices.len();

    let mut passing = Vec::new();

    for (iter, chunk) in all_gene_indices.chunks(SCENIC_GENE_CHUNK_SIZE).enumerate() {
        let mut gene_chunks = reader.read_gene_parallel(chunk)?;
        gene_chunks.par_iter_mut().for_each(|c| {
            c.filter_selected_cells(&cell_set);
        });

        for gene in &gene_chunks {
            let total_counts: u32 = gene.data_raw.iter().sum();
            let expressed_fraction = gene.nnz as f32 / n_cells as f32;

            if total_counts >= scenic_params.min_counts as u32
                && expressed_fraction >= scenic_params.min_cells
            {
                passing.push(gene.original_index);
            }
        }
        if verbosity.normal_verbosity() {
            println!(
                "Processed chunk {} out of {} for SCENIC inclusion criteria.",
                iter + 1,
                total_genes.div_ceil(SCENIC_GENE_CHUNK_SIZE)
            );
        }
    }

    Ok(passing)
}

/// Run SCENIC GRN inference and return a TF-by-gene importance matrix.
///
/// ### Params
///
/// * `reader` - Reader over the sparse gene expression store.
/// * `cell_indices` - Indices of cells to use.
/// * `gene_indices` - Target gene indices.
/// * `tf_indices` - Transcription factor gene indices (predictors).
/// * `scenic_params` - Reference to the SCENIC parameters.
/// * `seed` - Base random seed for reproducibility.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// A `Mat<f32>` of shape `(n_genes, n_tfs)` where entry `[i, j]` is the
/// normalised importance of TF `j` for target gene `i`.
///
/// ### Implementation details
///
/// TF expression data is loaded once, filtered and quantised into a
/// `QuantisedStore`. Target genes are read in I/O chunks of
/// `SCENIC_GENE_CHUNK_SIZE`, filtered and converted to sparse columns,
/// then collected into a single flat buffer. This buffer is sliced into
/// multi-output batches of `n_multi_output` targets each. All batches are
/// fitted in a single `par_iter` pass, giving rayon full work-stealing
/// flexibility across all batches rather than being constrained to the ~16
/// batches within a single I/O chunk.
///
/// Memory note: all target gene sparse columns are held simultaneously. At
/// 100k cells and 10% sparsity this is roughly 60KB per gene (~1.2GB for
/// 20k genes). For datasets where this is prohibitive, a wave-based
/// approach processing a few thousand genes at a time can be substituted.
pub fn run_scenic_grn<S: SingleCellReading>(
    reader: &S,
    cell_indices: &[usize],
    gene_indices: &[usize],
    tf_indices: &[usize],
    scenic_params: &ScenicParams,
    seed: usize,
    verbose: usize,
) -> Result<Mat<f32>, BixverseErrors> {
    let setup = scenic_common_setup(reader, cell_indices, tf_indices, verbose)?;
    let n_genes = gene_indices.len();

    match &scenic_params.regression_learner {
        RegressionLearner::GradientBoosting(gbm_config) => run_scenic_gbm(
            setup.reader,
            &setup.cell_set,
            gene_indices,
            &setup.tf_data,
            setup.n_cells,
            setup.n_tfs,
            n_genes,
            gbm_config,
            seed,
            verbose,
            setup.start_total,
        ),
        _ => run_scenic_multi_output(
            setup.reader,
            &setup.cell_set,
            cell_indices,
            gene_indices,
            &setup.tf_data,
            setup.n_cells,
            setup.n_tfs,
            n_genes,
            scenic_params,
            seed,
            verbose,
            setup.start_total,
        ),
    }
}

/// Run SCENIC GRN inference in streaming mode, processing target genes in
/// waves to bound memory usage.
///
/// ### Params
///
/// * `reader` - Reader over the sparse gene expression store.
/// * `cell_indices` - Indices of cells to use.
/// * `gene_indices` - Target gene indices.
/// * `tf_indices` - Transcription factor gene indices (predictors).
/// * `scenic_params` - Reference to the SCENIC parameters.
/// * `seed` - Base random seed for reproducibility.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// A `Mat<f32>` of shape `(n_genes, n_tfs)` where entry `[i, j]` is the
/// normalised importance of TF `j` for target gene `i`.
///
/// ### Implementation details
///
/// Unlike `run_scenic_grn` which loads all target gene data upfront, this
/// version reads target genes one I/O chunk at a time
/// (`SCENIC_GENE_CHUNK_SIZE`), slices each chunk into multi-output batches,
/// fits all batches within the chunk in parallel, stores the results, then
/// drops the chunk data before reading the next. Peak memory for target gene
/// data is bounded to one chunk (~1024 genes) rather than all targets.
///
/// There is a minor parallelism penalty at chunk boundaries: the last wave
/// of batches in a chunk may not fully saturate all cores. For most
/// configurations (1024 genes / 64 per batch = 16 batches) this is
/// acceptable. Use `run_scenic_grn` when memory permits for maximum
/// throughput.
pub fn run_scenic_grn_streaming<S: SingleCellReading>(
    reader: &S,
    cell_indices: &[usize],
    gene_indices: &[usize],
    tf_indices: &[usize],
    scenic_params: &ScenicParams,
    seed: usize,
    verbose: usize,
) -> Result<Mat<f32>, BixverseErrors> {
    let setup = scenic_common_setup(reader, cell_indices, tf_indices, verbose)?;
    let n_genes = gene_indices.len();

    match &scenic_params.regression_learner {
        RegressionLearner::GradientBoosting(gbm_config) => run_scenic_gbm_streaming(
            setup.reader,
            &setup.cell_set,
            gene_indices,
            &setup.tf_data,
            setup.n_cells,
            setup.n_tfs,
            n_genes,
            gbm_config,
            seed,
            verbose,
            setup.start_total,
        ),
        _ => run_scenic_multi_output_streaming(
            setup.reader,
            &setup.cell_set,
            cell_indices,
            gene_indices,
            &setup.tf_data,
            setup.n_cells,
            setup.n_tfs,
            n_genes,
            scenic_params,
            seed,
            verbose,
            setup.start_total,
        ),
    }
}

///////////
// Tests //
///////////

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

    /// Variance is recovered from the running sum and sum of squares, with no second pass.
    #[test]
    fn node_variance_basic() {
        // values: 1, 2, 3 -> mean=2, var=2/3
        let (sum, sum_sq) = (6.0f32, 14.0f32);
        let v = node_variance_f32(sum, sum_sq, 3);
        assert!((v - 2.0 / 3.0).abs() < 1e-6, "got {v}");
    }

    /// A constant node gives exactly zero, not a small floating point residue.
    #[test]
    fn node_variance_uniform() {
        // values: 3, 3, 3 -> var=0
        let (sum, sum_sq) = (9.0f32, 27.0f32);
        let v = node_variance_f32(sum, sum_sq, 3);
        assert_eq!(v, 0.0);
    }

    /// Bins accumulate every target side by side and the cumulative sums roll up over bins.
    #[test]
    fn histogram_build_multi_target() {
        let n_targets = 2;
        let n_features = 1;
        let n_samples = 4;
        let mut bufs = TreeBuffers::new(n_features, n_samples, n_targets);

        let tf_col: Vec<u8> = vec![0, 0, 0, 10, 10, 0];
        let sample_slice: Vec<u32> = vec![0, 1, 3, 4];
        // Interleaved y: sample0=[1,2], sample1=[3,4], sample2=[5,6], sample3=[7,8]
        let y_slice: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];

        bufs.build_histograms(&tf_col, &sample_slice, &y_slice, n_targets);

        assert_eq!(bufs.counts[0], 2);
        assert!((bufs.y_sums[0] - 4.0).abs() < 1e-10);
        assert!((bufs.y_sums[1] - 6.0).abs() < 1e-10);
        assert!((bufs.y_sum_sqs[0] - 10.0).abs() < 1e-10);
        assert!((bufs.y_sum_sqs[1] - 20.0).abs() < 1e-10);

        assert_eq!(bufs.counts[10], 2);
        assert!((bufs.y_sums[10 * n_targets] - 12.0).abs() < 1e-10);
        assert!((bufs.y_sums[10 * n_targets + 1] - 14.0).abs() < 1e-10);
        assert!((bufs.y_sum_sqs[10 * n_targets] - 74.0).abs() < 1e-10);
        assert!((bufs.y_sum_sqs[10 * n_targets + 1] - 100.0).abs() < 1e-10);

        assert_eq!(bufs.cum_counts[10], 4);
        assert!((bufs.cum_y_sums[10 * n_targets] - 16.0).abs() < 1e-10);
        assert!((bufs.cum_y_sums[10 * n_targets + 1] - 20.0).abs() < 1e-10);
    }

    /// The CSR-style offsets slice out each cell's target entries, empty cells included.
    #[test]
    fn sparse_y_batch_construction() {
        // Simulate SparseAxis with just indices and values
        // We can't construct SparseAxis directly here, so test via
        // SparseYBatch manually
        let offsets = vec![0u32, 1, 2, 4, 4];
        let target_indices = vec![0u8, 1, 0, 1];
        let values = vec![1.0f32, 2.0, 3.0, 4.0];

        let sy = SparseYBatch {
            offsets,
            target_indices,
            values,
        };

        // cell 0: target 0 = 1.0
        let (ti, tv) = sy.cell_entries(0);
        assert_eq!(ti, &[0]);
        assert_eq!(tv, &[1.0]);

        // cell 1: target 1 = 2.0
        let (ti, tv) = sy.cell_entries(1);
        assert_eq!(ti, &[1]);
        assert_eq!(tv, &[2.0]);

        // cell 2: both targets
        let (ti, tv) = sy.cell_entries(2);
        assert_eq!(ti, &[0, 1]);
        assert_eq!(tv, &[3.0, 4.0]);

        // cell 3: empty
        let (ti, tv) = sy.cell_entries(3);
        assert!(ti.is_empty());
        assert!(tv.is_empty());
    }

    /// The sparse histogram builder agrees bin for bin with the dense one on equivalent data.
    #[test]
    fn sparse_histogram_matches_dense() {
        // Build a scenario where we can compare dense and sparse histogram
        // outputs directly.
        let n_targets = 2;
        let n_features = 1;
        let n_samples = 6;
        let tf_col: Vec<u8> = vec![0, 0, 10, 10, 0, 10];
        let sample_slice: Vec<u32> = vec![0, 1, 2, 3, 4, 5];

        // Dense y: interleaved [s0t0, s0t1, s1t0, s1t1, ...]
        let y_dense: Vec<f32> = vec![1.0, 2.0, 0.0, 0.0, 3.0, 4.0, 0.0, 0.0, 5.0, 6.0, 7.0, 8.0];

        // Equivalent sparse representation
        let sparse_y = SparseYBatch {
            offsets: vec![0, 2, 2, 4, 4, 6, 8],
            target_indices: vec![0, 1, 0, 1, 0, 1, 0, 1],
            values: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
        };

        let mut bufs_dense = TreeBuffers::new(n_features, n_samples, n_targets);
        let mut bufs_sparse = TreeBuffers::new_sparse(n_features, n_samples, n_targets);

        let (min_d, max_d) =
            bufs_dense.build_histograms(&tf_col, &sample_slice, &y_dense, n_targets);
        let (min_s, max_s) =
            bufs_sparse.build_histograms_sparse(&tf_col, &sample_slice, &sparse_y, n_targets);

        assert_eq!(min_d, min_s);
        assert_eq!(max_d, max_s);

        // Counts must match exactly within the active bounds
        assert_eq!(
            &bufs_dense.counts[min_d..=max_d],
            &bufs_sparse.counts[min_s..=max_s]
        );
        assert_eq!(
            &bufs_dense.cum_counts[min_d..=max_d],
            &bufs_sparse.cum_counts[min_s..=max_s]
        );

        // Y sums and sum_sqs must match within the active bounds
        let start_idx = min_d * n_targets;
        let end_idx = (max_d + 1) * n_targets;

        for i in start_idx..end_idx {
            assert!(
                (bufs_dense.y_sums[i] - bufs_sparse.y_sums[i]).abs() < 1e-10,
                "y_sums mismatch at {i}: dense={} sparse={}",
                bufs_dense.y_sums[i],
                bufs_sparse.y_sums[i]
            );
            assert!(
                (bufs_dense.y_sum_sqs[i] - bufs_sparse.y_sum_sqs[i]).abs() < 1e-10,
                "y_sum_sqs mismatch at {i}"
            );
            assert!(
                (bufs_dense.cum_y_sums[i] - bufs_sparse.cum_y_sums[i]).abs() < 1e-10,
                "cum_y_sums mismatch at {i}"
            );
            assert!(
                (bufs_dense.cum_y_sum_sqs[i] - bufs_sparse.cum_y_sum_sqs[i]).abs() < 1e-10,
                "cum_y_sum_sqs mismatch at {i}"
            );
        }
    }

    /// Pins the variance-reduction score of a split against a hand-computed value.
    #[test]
    fn evaluate_split_basic() {
        let n_targets = 1;
        let mut cum_counts = [0u32; 256];
        let mut cum_y_sums = vec![0.0f32; 256];
        let mut cum_y_sum_sqs = vec![0.0f32; 256];

        // bin 0: count=2, sum=4, sumsq=10
        // bin 10: count=2, sum=22, sumsq=244
        cum_counts[0] = 2;
        cum_y_sums[0] = 4.0;
        cum_y_sum_sqs[0] = 10.0;
        for b in 1..256 {
            cum_counts[b] = cum_counts[b - 1];
            cum_y_sums[b] = cum_y_sums[b - 1];
            cum_y_sum_sqs[b] = cum_y_sum_sqs[b - 1];
            if b == 10 {
                cum_counts[b] += 2;
                cum_y_sums[b] += 22.0;
                cum_y_sum_sqs[b] += 244.0;
            }
        }

        let parent_vars = vec![node_variance_f32(26.0, 254.0, 4)];
        let y_sums_total = vec![26.0f32];
        let y_sum_sqs_total = vec![254.0f32];

        let mut best_score = 0.0f32;
        let mut best_feature = usize::MAX;
        let mut best_threshold = 0u8;
        let mut best_n_left = 0u32;
        let mut best_ys_l = vec![0.0f32; 1];
        let mut best_yss_l = vec![0.0f32; 1];

        evaluate_split_multi(
            0, // threshold: bin <= 0 goes left
            0, // feature
            &parent_vars,
            4,
            &y_sums_total,
            &y_sum_sqs_total,
            &cum_counts,
            &cum_y_sums,
            &cum_y_sum_sqs,
            n_targets,
            1, // min_samples_leaf
            &mut best_score,
            &mut best_feature,
            &mut best_threshold,
            &mut best_n_left,
            &mut best_ys_l,
            &mut best_yss_l,
        );

        assert_eq!(best_feature, 0);
        assert_eq!(best_threshold, 0);
        assert_eq!(best_n_left, 2);
        // parent_var = 254/4 - (26/4)^2 = 63.5 - 42.25 = 21.25
        // var_l = 10/2 - (4/2)^2 = 5 - 4 = 1
        // var_r = 244/2 - (22/2)^2 = 122 - 121 = 1
        // reduction = 21.25 - 0.5*1 - 0.5*1 = 20.25
        assert!((best_score - 20.25).abs() < 1e-10, "got {}", best_score);
    }

    /// Batching returns every gene exactly once and does not hand back the input order.
    #[test]
    fn batch_genes_random_is_permutation() {
        let genes: Vec<usize> = (0..100).collect();
        let shuffled = batch_genes_random(&genes, 42);
        assert_eq!(shuffled.len(), genes.len());

        let mut sorted = shuffled.clone();
        sorted.sort();
        assert_eq!(sorted, genes);

        // Should actually shuffle (not identity)
        assert_ne!(shuffled, genes);
    }

    /// The same seed gives the same gene order, so a run can be reproduced.
    #[test]
    fn batch_genes_random_deterministic() {
        let genes: Vec<usize> = (0..100).collect();
        let a = batch_genes_random(&genes, 42);
        let b = batch_genes_random(&genes, 42);
        assert_eq!(a, b);
    }

    /// Asking for more cells than exist returns all of them rather than padding or resampling.
    #[test]
    fn subsample_cells_all_when_small() {
        let cells: Vec<usize> = (0..50).collect();
        let sub = subsample_cells(&cells, 100, 0);
        assert_eq!(sub.len(), 50);
        // Should contain all original indices
        let mut sorted = sub.clone();
        sorted.sort();
        assert_eq!(sorted, cells);
    }

    /// Subsampling yields the requested count with no duplicates and no out-of-range indices.
    #[test]
    fn subsample_cells_correct_count() {
        let cells: Vec<usize> = (0..1000).collect();
        let sub = subsample_cells(&cells, 100, 42);
        assert_eq!(sub.len(), 100);
        // All elements should be valid cell indices
        for &c in &sub {
            assert!(c < 1000);
        }
        // No duplicates (Fisher-Yates guarantees this)
        let mut sorted = sub.clone();
        sorted.sort();
        sorted.dedup();
        assert_eq!(sorted.len(), 100);
    }

    /// Each feature gets its own 256-bin block, indexed by the quantised value of the cell.
    #[test]
    fn node_hist_build_from_samples() {
        let data = vec![
            0, 10, 10, 20, // feature 0, cells 0..3
            5, 5, 15, 15, // feature 1, cells 0..3
        ];
        let store = QuantisedStore::from_raw(data, 4, 2);
        let residuals = vec![1.0, 2.0, 3.0, 4.0];
        let samples: Vec<u32> = vec![0, 1, 2, 3];

        let mut hist = NodeHistograms::new(2);
        hist.build_from_samples(&store, &samples, &residuals);

        // Feature 0: bin 0 -> cell 0 (r=1), bin 10 -> cells 1,2 (r=2,3),
        //            bin 20 -> cell 3 (r=4)
        assert_eq!(hist.counts[0], 1);
        assert_eq!(hist.counts[10], 2);
        assert_eq!(hist.counts[20], 1);
        assert!((hist.y_sums[0] - 1.0).abs() < 1e-10);
        assert!((hist.y_sums[10] - 5.0).abs() < 1e-10);
        assert!((hist.y_sums[20] - 4.0).abs() < 1e-10);
        assert!((hist.y_sum_sqs[10] - 13.0).abs() < 1e-10); // 4+9

        // Feature 1: bin 5 -> cells 0,1 (r=1,2), bin 15 -> cells 2,3 (r=3,4)
        assert_eq!(hist.counts[256 + 5], 2);
        assert_eq!(hist.counts[256 + 15], 2);
        assert!((hist.y_sums[256 + 5] - 3.0).abs() < 1e-10);
        assert!((hist.y_sums[256 + 15] - 7.0).abs() < 1e-10);
    }

    /// Restricting the sample list accumulates only those cells and leaves the other bins empty.
    #[test]
    fn node_hist_build_subset() {
        // Same store, only use cells 0 and 2
        let data = vec![0, 10, 10, 20, 5, 5, 15, 15];
        let store = QuantisedStore::from_raw(data, 4, 2);
        let residuals = vec![1.0, 2.0, 3.0, 4.0];
        let samples: Vec<u32> = vec![0, 2];

        let mut hist = NodeHistograms::new(2);
        hist.build_from_samples(&store, &samples, &residuals);

        // Feature 0: bin 0 -> cell 0 (r=1), bin 10 -> cell 2 (r=3)
        assert_eq!(hist.counts[0], 1);
        assert_eq!(hist.counts[10], 1);
        assert_eq!(hist.counts[20], 0);
        assert!((hist.y_sums[0] - 1.0).abs() < 1e-10);
        assert!((hist.y_sums[10] - 3.0).abs() < 1e-10);
    }

    /// Parent minus child must equal building the sibling histogram from scratch.
    #[test]
    fn histogram_subtraction_matches_complement() {
        // Build parent from all 4 samples, child from {0,2},
        // subtraction should match building from {1,3}
        let data = vec![0, 10, 10, 20, 5, 5, 15, 15];
        let store = QuantisedStore::from_raw(data, 4, 2);
        let residuals = vec![1.0, 2.0, 3.0, 4.0];

        let mut parent = NodeHistograms::new(2);
        parent.build_from_samples(&store, &[0, 1, 2, 3], &residuals);

        let mut child = NodeHistograms::new(2);
        child.build_from_samples(&store, &[0, 2], &residuals);

        let mut complement_direct = NodeHistograms::new(2);
        complement_direct.build_from_samples(&store, &[1, 3], &residuals);

        // Compute subtraction via pool
        let mut pool = HistogramPool::new(3, 2);
        // Copy parent into slot 0
        pool.histograms[0].counts.copy_from_slice(&parent.counts);
        pool.histograms[0].y_sums.copy_from_slice(&parent.y_sums);
        pool.histograms[0]
            .y_sum_sqs
            .copy_from_slice(&parent.y_sum_sqs);
        // Copy child into slot 1
        pool.histograms[1].counts.copy_from_slice(&child.counts);
        pool.histograms[1].y_sums.copy_from_slice(&child.y_sums);
        pool.histograms[1]
            .y_sum_sqs
            .copy_from_slice(&child.y_sum_sqs);

        pool.subtract(0, 1, 2);

        let result = &pool.histograms[2];
        let n = 2 * 256;
        for i in 0..n {
            assert_eq!(
                result.counts[i], complement_direct.counts[i],
                "counts mismatch at {i}"
            );
            assert!(
                (result.y_sums[i] - complement_direct.y_sums[i]).abs() < 1e-6,
                "y_sums mismatch at {i}: {} vs {}",
                result.y_sums[i],
                complement_direct.y_sums[i]
            );
            assert!(
                (result.y_sum_sqs[i] - complement_direct.y_sum_sqs[i]).abs() < 1e-6,
                "y_sum_sqs mismatch at {i}"
            );
        }
    }

    /// A cleanly separable feature yields the expected partition and left-hand sum.
    #[test]
    fn find_split_obvious_partition() {
        // Obvious split at threshold ~50: left=[1,1,1,1], right=[10,10,10,10]
        let data: Vec<u8> = vec![0, 0, 0, 0, 100, 100, 100, 100];
        let store = QuantisedStore::from_raw(data, 8, 1);
        let residuals = vec![1.0, 1.0, 1.0, 1.0, 10.0, 10.0, 10.0, 10.0];
        let samples: Vec<u32> = (0..8).collect();

        let mut hist = NodeHistograms::new(1);
        hist.build_from_samples(&store, &samples, &residuals);

        let total_sum: f32 = residuals.iter().sum();
        let total_sum_sq: f32 = residuals.iter().map(|&r| r * r).sum();
        let mut feat_buf: Vec<usize> = vec![0];
        let mut rng = SmallRng::seed_from_u64(42);

        let split = hist
            .find_best_split(total_sum, total_sum_sq, 8, 1, 0, &mut feat_buf, &mut rng)
            .expect("should find a split");

        assert_eq!(split.feature, 0);
        // Any threshold in [0, 99] gives the same partition
        assert!(split.threshold < 100, "threshold={}", split.threshold);
        assert_eq!(split.n_left, 4);
        assert!((split.y_sum_left - 4.0).abs() < 1e-6);
    }

    /// When no candidate leaves both sides large enough the search returns None, not a bad split.
    #[test]
    fn find_split_respects_min_samples_leaf() {
        // With min_samples_leaf=3, only threshold that puts 3 left / 1 right
        // violates the right side. Should return None.
        let data: Vec<u8> = vec![0, 0, 0, 100];
        let store = QuantisedStore::from_raw(data, 4, 1);
        let residuals = vec![1.0, 2.0, 3.0, 100.0];
        let samples: Vec<u32> = vec![0, 1, 2, 3];

        let mut hist = NodeHistograms::new(1);
        hist.build_from_samples(&store, &samples, &residuals);

        let total_sum: f32 = residuals.iter().sum();
        let total_sum_sq: f32 = residuals.iter().map(|&r| r * r).sum();
        let mut feat_buf: Vec<usize> = vec![0];
        let mut rng = SmallRng::seed_from_u64(0);

        let split = hist.find_best_split(total_sum, total_sum_sq, 4, 3, 0, &mut feat_buf, &mut rng);
        assert!(split.is_none(), "should not find a valid split");
    }

    /// Constant residuals leave nothing to reduce, so no split is proposed.
    #[test]
    fn find_split_no_variance() {
        // All residuals identical => no useful split
        let data: Vec<u8> = vec![0, 50, 100, 200];
        let store = QuantisedStore::from_raw(data, 4, 1);
        let residuals = vec![5.0, 5.0, 5.0, 5.0];
        let samples: Vec<u32> = vec![0, 1, 2, 3];

        let mut hist = NodeHistograms::new(1);
        hist.build_from_samples(&store, &samples, &residuals);

        let mut feat_buf: Vec<usize> = vec![0];
        let mut rng = SmallRng::seed_from_u64(0);

        let split = hist.find_best_split(20.0, 100.0, 4, 1, 0, &mut feat_buf, &mut rng);
        assert!(split.is_none());
    }

    /// The leaf prediction is shrunk by the learning rate, applied to every cell, and scored out of bag.
    #[test]
    fn apply_leaf_updates_residuals() {
        let mut residuals = vec![10.0, 20.0, 30.0, 40.0];
        let train: Vec<u32> = vec![0, 1];
        let oob: Vec<u32> = vec![2, 3];
        // y_sum = 10 + 20 = 30, n_train = 2, pred = 15
        // lr_pred = 0.1 * 15 = 1.5
        let mut improvement = 0.0f32;

        apply_regression_leaf(&mut residuals, &train, &oob, 30.0, 2, 0.1, &mut improvement);

        // All residuals reduced by lr * pred = 1.5
        assert!((residuals[0] - 8.5).abs() < 1e-6);
        assert!((residuals[1] - 18.5).abs() < 1e-6);
        assert!((residuals[2] - 28.5).abs() < 1e-6);
        assert!((residuals[3] - 38.5).abs() < 1e-6);

        // OOB improvement for cell 2: 2*1.5*30 - 1.5^2 = 90 - 2.25 = 87.75
        // OOB improvement for cell 3: 2*1.5*40 - 1.5^2 = 120 - 2.25 = 117.75
        // Total = 205.5
        assert!((improvement - 205.5).abs() < 1e-4, "got {}", improvement);
    }

    // ---- fit_grnboost2_sparse (integration) ----

    /// The TF that actually drives the target takes the top importance, and importances sum to one.
    #[test]
    fn fit_grnboost2_strong_signal() {
        // 200 cells, 3 features (TFs). Target = 2.0 * TF0 + noise.
        // TF0 should dominate importance.
        let n_cells = 200;
        let n_features = 3;
        let mut rng = SmallRng::seed_from_u64(123);

        // Build quantised feature store directly
        let mut data = vec![0u8; n_features * n_cells];
        for f in 0..n_features {
            for c in 0..n_cells {
                data[f * n_cells + c] = rng.random_range(0..=255u8);
            }
        }
        let store = QuantisedStore::from_raw(data, n_cells, n_features);

        // Target: 2 * (TF0 quantised value) + small noise
        // Build as a sparse axis
        let tf0_col = store.get_col(0);
        let mut target_indices: Vec<usize> = Vec::new();
        let mut target_values: Vec<f32> = Vec::new();
        for c in 0..n_cells {
            let val = 2.0 * tf0_col[c] as f32 + rng.random_range(-5.0..5.0f32);
            if val.abs() > 0.01 {
                target_indices.push(c);
                target_values.push(val);
            }
        }

        let target = SparseAxis {
            indices: target_indices.to_vec(),
            data: Vec::new(),
            data_2: Some(target_values),
            cs_type: CompressedSparseFormat::Csc,
            len: n_cells,
        };

        let config = GradientBoostingConfig {
            n_trees_max: 200,
            learning_rate: 0.05,
            max_depth: 3,
            min_samples_leaf: 5,
            early_stop_window: 20,
            subsample_rate: 0.9,
            n_features_split: 0,
        };

        let imp = fit_grnboost2_sparse(&target, &store, n_cells, &config, 42).unwrap();

        assert_eq!(imp.len(), n_features);
        // TF0 should have the highest importance
        assert!(
            imp[0] > imp[1] && imp[0] > imp[2],
            "TF0 importance ({}) should dominate: {:?}",
            imp[0],
            imp
        );
        // Normalised to sum to 1
        let sum: f32 = imp.iter().sum();
        assert!((sum - 1.0).abs() < 1e-6, "sum={}", sum);
    }

    /// An all-zero target has no variance to explain, so no TF picks up spurious importance.
    #[test]
    fn fit_grnboost2_early_stopping() {
        // Pure noise target: all zero. Early stopping should kick in fast.
        let n_cells = 100;
        let n_features = 2;
        let mut rng = SmallRng::seed_from_u64(99);

        let mut data = vec![0u8; n_features * n_cells];
        for i in 0..data.len() {
            data[i] = rng.random_range(0..=255u8);
        }
        let store = QuantisedStore::from_raw(data, n_cells, n_features);

        // Empty target (all zeros)
        let target = SparseAxis {
            indices: Vec::new(),
            data: Vec::new(),
            data_2: Some(Vec::new()),
            cs_type: CompressedSparseFormat::Csc,
            len: n_cells,
        };

        let config = GradientBoostingConfig {
            n_trees_max: 500,
            learning_rate: 0.01,
            max_depth: 3,
            min_samples_leaf: 5,
            early_stop_window: 10,
            subsample_rate: 0.9,
            n_features_split: 0,
        };

        let imp = fit_grnboost2_sparse(&target, &store, n_cells, &config, 0).unwrap();

        // With zero target, importances should all be zero (no variance to
        // reduce) or the function should terminate quickly.
        let total: f32 = imp.iter().sum();
        assert!(
            total < 1e-6,
            "expected near-zero importance for zero target, got {}",
            total
        );
    }

    /// The whole boosting loop, subsampling included, is reproducible from the seed alone.
    #[test]
    fn fit_grnboost2_deterministic() {
        let n_cells = 100;
        let n_features = 2;
        let mut rng = SmallRng::seed_from_u64(7);
        let mut data = vec![0u8; n_features * n_cells];
        for i in 0..data.len() {
            data[i] = rng.random_range(0..=255u8);
        }
        let store = QuantisedStore::from_raw(data, n_cells, n_features);
        let mut t_idx: Vec<u32> = Vec::new();
        let mut t_val: Vec<f32> = Vec::new();
        let col = store.get_col(0);
        for c in 0..n_cells {
            let v = col[c] as f32;
            if v > 0.0 {
                t_idx.push(c as u32);
                t_val.push(v);
            }
        }
        let target = SparseAxis {
            indices: t_idx.iter().map(|&i| i as usize).collect(),
            data: Vec::new(),
            data_2: Some(t_val),
            cs_type: CompressedSparseFormat::Csc,
            len: n_cells,
        };
        let config = GradientBoostingConfig::default();
        let a = fit_grnboost2_sparse(&target, &store, n_cells, &config, 42).unwrap();
        let b = fit_grnboost2_sparse(&target, &store, n_cells, &config, 42).unwrap();
        assert_eq!(a, b, "same seed should produce identical results");
    }

    /// Released slots go back on the free list and get handed out again instead of leaking.
    #[test]
    fn pool_acquire_release_cycle() {
        let mut pool = HistogramPool::new(3, 1);
        let a = pool.acquire();
        let b = pool.acquire();
        let c = pool.acquire();
        assert_ne!(a, b);
        assert_ne!(b, c);
        pool.release(b);
        let d = pool.acquire();
        assert_eq!(d, b, "released slot should be reused");
    }

    /// Over-acquiring is a caller bug, so the pool panics rather than growing silently.
    #[test]
    #[should_panic(expected = "histogram pool exhausted")]
    fn pool_exhaustion_panics() {
        let mut pool = HistogramPool::new(2, 1);
        let _ = pool.acquire();
        let _ = pool.acquire();
        let _ = pool.acquire(); // should panic
    }
}