legume-numeric 0.8.11

use crate::matrix::rand_util::mix_seed;
use rand::prelude::SliceRandom;
use rand::rngs::SmallRng;
use rand::SeedableRng;
use rayon::prelude::*;
use rustc_hash::FxHashMap as HashMap;
use std::hash::Hash;

/// Fixed base seed for the pseudobulk down-sample shuffle. Mixed with each
/// group's stable identity so the kept subset is reproducible and independent
/// of thread scheduling (the shuffle previously drew from OS entropy).
const PARTITION_SHUFFLE_SEED: u64 = 0x5041_5254_5348_5546; // "PARTSHUF"

/// Median of a slice (sorts a copy). Returns `0.0` for an empty slice and
/// is NaN-tolerant (NaNs compare as equal rather than panicking).
pub fn median(values: &[f32]) -> f32 {
    if values.is_empty() {
        return 0.0;
    }
    let mut sorted = values.to_vec();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let n = sorted.len();
    if n % 2 == 1 {
        sorted[n / 2]
    } else {
        0.5 * (sorted[n / 2 - 1] + sorted[n / 2])
    }
}

/// Quantiles of a slice at the given probabilities (sorts a copy; each
/// quantile is the element at `⌊q·(n − 1)⌋`). Empty input gives zeros;
/// NaN-tolerant like [`median`].
pub fn quantiles(values: &[f32], qs: &[f64]) -> Vec<f32> {
    if values.is_empty() {
        return vec![0.0; qs.len()];
    }
    let mut sorted = values.to_vec();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let last = sorted.len() - 1;
    qs.iter()
        .map(|&q| sorted[((q.clamp(0.0, 1.0) * last as f64).floor() as usize).min(last)])
        .collect()
}

/// Cosine similarity of two vectors; `0.0` when either has no length.
pub fn cosine(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if na > 0.0 && nb > 0.0 {
        dot / (na * nb)
    } else {
        0.0
    }
}

/// partition membership vector into groups of indexes
/// # Arguments
/// * `membership` - a vector of membership (E.g., cluster assignment)
/// * `nelem_per_group` - number of elements per group (if None, no downsampling)
/// # Returns
/// A hashmap: cluster/group name -> indexes of the elements
pub fn partition_by_membership<T>(
    membership: &[T],
    nelem_per_group: Option<usize>,
) -> HashMap<T, Vec<usize>>
where
    T: Eq + Hash + Clone + Send + Sync,
{
    let mut pb_elems: HashMap<T, Vec<usize>> = HashMap::default();
    for (cell, k) in membership.iter().enumerate() {
        pb_elems.entry(k.clone()).or_default().push(cell);
    }

    // Down sample elements if needed
    pb_elems.par_iter_mut().for_each(|(_k, cells)| {
        let ncells = cells.len();
        if let Some(ntarget) = nelem_per_group {
            if ncells > ntarget {
                // Seed per group from its smallest member — stable, unique per
                // group, and independent of thread order — so the down-sampled
                // subset is reproducible run-to-run.
                let key = cells.first().copied().unwrap_or(0) as u64;
                let mut rng = SmallRng::seed_from_u64(mix_seed(PARTITION_SHUFFLE_SEED, key));
                cells.shuffle(&mut rng);
                cells.truncate(ntarget);
            }
        }
        // dbg!(cells.len());
    });
    pb_elems
}

/// Pick a per-job cell block size given the feature count.
/// Target ≈ `MIN_BLOCK_SIZE × 10 000` cells×features of work per job;
/// clamped to `[MIN_BLOCK_SIZE, MAX_BLOCK_SIZE]`.
/// `None` passed to `generate_minibatch_intervals` / `create_jobs` resolves through this.
///
/// This bounds READ work per job. It is **not** a bound on what a job
/// allocates, and it cannot be: it never sees the caller's accumulator.
///
/// So if a job's accumulator is sized by an OUTPUT axis rather than by the
/// block it reads — `[n_features × k]`, one dense column per group — then
/// collecting one per job with `map().collect()` makes peak memory scale with
/// the JOB COUNT, which this function sets. Note the shape of that hazard: a
/// wide, sparse input drives the block size DOWN to the floor, so the job
/// count goes up exactly when each accumulator is at its largest.
///
/// Accumulate with `try_fold`/`try_reduce` so live accumulators track the
/// worker count instead; and where the accumulator's fixed axis is itself
/// unbounded, key the jobs by the output the job owns and write into it
/// directly, rather than giving every job a copy of the whole thing.
pub fn default_block_size(num_features: usize) -> usize {
    const MIN_BLOCK_SIZE: usize = 100;
    const MAX_BLOCK_SIZE: usize = 10_000;
    const TARGET_WORK: usize = MIN_BLOCK_SIZE * 10_000;
    if num_features == 0 {
        return MIN_BLOCK_SIZE;
    }
    (TARGET_WORK / num_features).clamp(MIN_BLOCK_SIZE, MAX_BLOCK_SIZE)
}

/// Generate minibatch intervals.
/// * `ntot` - number of total samples (columns/cells/edges/...)
/// * `num_features` - number of features per sample; used only when `batch_size` is `None`
///   to derive an adaptive default via [`default_block_size`]
/// * `batch_size` - `Some(n)` to use `n` explicitly, `None` to auto-scale by feature count
pub fn generate_minibatch_intervals(
    ntot: usize,
    num_features: usize,
    batch_size: Option<usize>,
) -> Vec<(usize, usize)> {
    let batch_size = batch_size.unwrap_or_else(|| default_block_size(num_features));
    let num_batches = ntot.div_ceil(batch_size);
    (0..num_batches)
        .map(|b| {
            let lb: usize = b * batch_size;
            let ub: usize = ((b + 1) * batch_size).min(ntot);
            (lb, ub)
        })
        .collect::<Vec<_>>()
}

/// Column blocks bounded by BYTES of triplets, not by column count.
///
/// [`default_block_size`] bounds *work* — columns × features — which says
/// nothing about residency: a block of 100 dense columns and a block of 100
/// empty ones cost the same under it, and the sparse cost that actually OOMs is
/// the nnz a block materialises. This takes the measured per-column nnz (from
/// the backend's resident indptr) and cuts blocks so each stays under
/// `budget_bytes` at `bytes_per_nnz` — 24 for a `(u64, u64, f32)` triplet.
///
/// A single column heavier than the whole budget still becomes its own block:
/// it cannot be split at this layer, so the bound is `max(budget, heaviest
/// single column)` and the caller sees that in the block rather than an error.
/// Empty trailing columns still land in a block, so every column is covered.
pub fn byte_budget_intervals(
    nnz_per_col: &[u64],
    budget_bytes: usize,
    bytes_per_nnz: usize,
) -> Vec<(usize, usize)> {
    let budget_nnz = (budget_bytes / bytes_per_nnz.max(1)).max(1) as u64;
    let mut intervals = Vec::new();
    let mut lb = 0usize;
    let mut acc = 0u64;
    for (col, &nnz) in nnz_per_col.iter().enumerate() {
        if acc > 0 && acc + nnz > budget_nnz {
            intervals.push((lb, col));
            lb = col;
            acc = 0;
        }
        acc += nnz;
    }
    if lb < nnz_per_col.len() {
        intervals.push((lb, nnz_per_col.len()));
    }
    intervals
}

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