//! Per-bin expression normalization and panel aggregation.
//!
//! This is the scientific pre-processing the operators depend on: raw 10x counts
//! are zero-inflated and library-size-confounded, so we apply CPM-style
//! library-size scaling + log1p before any differential operator, and aggregate
//! multi-gene panels with a Seurat-style module score (z-score per gene, then
//! mean) rather than summing raw counts of a biologically heterogeneous panel.

use crate::cli::{NormalizeMode, PanelMode};
use crate::error::OrchestratorError;

/// Numerically safe `ln(1 + max(v, 0))` — counts are non-negative; guards NaN/Inf.
#[inline]
fn log1p_nonneg(v: f32) -> f32 {
    let x = if v.is_finite() { v.max(0.0) } else { 0.0 };
    (1.0_f32 + x).ln()
}

/// Per-bin library sizes (total counts per bin) from a bin-major CSR.
///
/// `indptr[bin]..indptr[bin+1]` are the non-zero entries of `data` for that bin.
pub(crate) fn bin_library_sizes(indptr: &[u64], data: &[f32]) -> Vec<f32> {
    if indptr.len() < 2 {
        return Vec::new();
    }
    let n_bins = indptr.len() - 1;
    let mut out = vec![0.0_f32; n_bins];
    for (bin, slot) in out.iter_mut().enumerate() {
        let lo = indptr[bin] as usize;
        let hi = indptr[bin + 1] as usize;
        let mut acc = 0.0_f64;
        for &v in &data[lo..hi.min(data.len())] {
            if v.is_finite() && v > 0.0 {
                acc += v as f64;
            }
        }
        *slot = acc as f32;
    }
    out
}

/// Apply per-bin normalization to one gene's dense values in place.
/// `libsize` (per bin total counts) is required for CPM; when absent, the
/// `CpmLog1p` mode degrades to `Log1p`.
pub(crate) fn normalize_gene(
    values: &mut [f32],
    libsize: Option<&[f32]>,
    mode: NormalizeMode,
    target_sum: f32,
) {
    match mode {
        NormalizeMode::None => {}
        NormalizeMode::Log1p => {
            for v in values.iter_mut() {
                *v = log1p_nonneg(*v);
            }
        }
        NormalizeMode::CpmLog1p => match libsize {
            Some(ls) if ls.len() == values.len() && target_sum.is_finite() && target_sum > 0.0 => {
                for (i, v) in values.iter_mut().enumerate() {
                    let l = ls[i];
                    let scaled = if l > 0.0 { *v / l * target_sum } else { 0.0 };
                    *v = log1p_nonneg(scaled);
                }
            }
            // No library size available (e.g. feature-slice path): log1p only.
            _ => {
                for v in values.iter_mut() {
                    *v = log1p_nonneg(*v);
                }
            }
        },
    }
}

/// Aggregate already-normalized per-gene fields into a single field.
///
/// A single gene passes through unchanged. Multi-gene panels use either a plain
/// sum (legacy) or a Seurat-style module score (default). `valid` marks bins
/// included in the module-score statistics; its length must match each field.
pub(crate) fn aggregate_panel(
    per_gene: &[Vec<f32>],
    valid: &[bool],
    mode: PanelMode,
) -> Result<Vec<f32>, OrchestratorError> {
    if per_gene.is_empty() {
        return Err(OrchestratorError::InvalidInput(
            "panel aggregation requires at least one gene",
        ));
    }
    if per_gene.len() == 1 {
        return Ok(per_gene[0].clone());
    }
    match mode {
        PanelMode::Sum => {
            let n = per_gene[0].len();
            let mut out = vec![0.0_f32; n];
            for gene in per_gene {
                if gene.len() != n {
                    return Err(OrchestratorError::InvalidInput(
                        "per-gene field length mismatch in panel sum",
                    ));
                }
                for (o, &v) in out.iter_mut().zip(gene.iter()) {
                    *o += v;
                }
            }
            Ok(out)
        }
        PanelMode::ModuleScore => {
            let refs: Vec<&[f32]> = per_gene.iter().map(|v| v.as_slice()).collect();
            Ok(kira_spatial_field::module_score(&refs, valid)?)
        }
    }
}
