datarust 0.6.6

//! Column-wise statistics, covariance and correlation helpers.

use crate::error::{DatarustError, Result};
#[cfg(feature = "rayon")]
use rayon::prelude::*;

/// Returns the mean of each column.
///
/// ```rust
/// use datarust::stats::column_mean;
///
/// let data = vec![vec![1.0, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]];
/// let means = column_mean(&data);
/// assert!((means[0] - 3.0).abs() < 1e-12);
/// assert!((means[1] - 4.0).abs() < 1e-12);
/// ```
pub fn column_mean(data: &[Vec<f64>]) -> Vec<f64> {
    if data.is_empty() {
        return vec![];
    }
    let cols = data[0].len();
    let n = data.len() as f64;
    #[cfg(feature = "rayon")]
    {
        (0..cols)
            .into_par_iter()
            .map(|j| {
                let s: f64 = data.iter().map(|r| r[j]).sum();
                s / n
            })
            .collect()
    }
    #[cfg(not(feature = "rayon"))]
    {
        (0..cols)
            .map(|j| {
                let s: f64 = data.iter().map(|r| r[j]).sum();
                s / n
            })
            .collect()
    }
}

/// Per-column mean over flat row-major data (single fused pass).
///
/// This is the flat-storage counterpart of [`column_mean`], walking the
/// contiguous buffer once with stride-1 inner access instead of gathering
/// column-by-column across scattered row allocations.
pub fn column_mean_flat(data: &[f64], rows: usize, cols: usize) -> Vec<f64> {
    if rows == 0 || cols == 0 {
        return vec![];
    }
    let n = rows as f64;
    let mut sums = vec![0.0; cols];
    for i in 0..rows {
        let base = i * cols;
        for j in 0..cols {
            sums[j] += data[base + j];
        }
    }
    sums.iter().map(|&s| s / n).collect()
}

/// Returns the variance of each column using the given delta degrees of freedom.
///
/// ```rust
/// use datarust::stats::column_variance;
///
/// let data = vec![vec![1.0, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]];
/// let var = column_variance(&data, 0);
/// assert!((var[0] - 8.0 / 3.0).abs() < 1e-12);
/// ```
pub fn column_variance(data: &[Vec<f64>], ddof: usize) -> Vec<f64> {
    if data.is_empty() {
        return vec![];
    }
    let cols = data[0].len();
    let n = data.len();
    // Guard against a non-positive denominator (ddof >= n); fall back to NaN
    // rather than producing +/-inf and propagating it through downstream transforms.
    let denom = n.saturating_sub(ddof) as f64;
    let means = column_mean(data);
    #[cfg(feature = "rayon")]
    {
        (0..cols)
            .into_par_iter()
            .map(|j| {
                let m = means[j];
                let s: f64 = data.iter().map(|r| (r[j] - m).powi(2)).sum();
                if denom > 0.0 {
                    s / denom
                } else {
                    f64::NAN
                }
            })
            .collect()
    }
    #[cfg(not(feature = "rayon"))]
    {
        (0..cols)
            .map(|j| {
                let m = means[j];
                let s: f64 = data.iter().map(|r| (r[j] - m).powi(2)).sum();
                if denom > 0.0 {
                    s / denom
                } else {
                    f64::NAN
                }
            })
            .collect()
    }
}

/// Returns the standard deviation of each column using the given delta degrees of freedom.
pub fn column_std(data: &[Vec<f64>], ddof: usize) -> Vec<f64> {
    column_variance(data, ddof)
        .iter()
        .map(|v| v.sqrt())
        .collect()
}

/// Returns the minimum value of each column.
///
/// ```rust
/// use datarust::stats::column_min;
///
/// let data = vec![vec![3.0, 1.0], vec![1.0, 5.0], vec![2.0, 2.0]];
/// let mins = column_min(&data);
/// assert_eq!(mins, vec![1.0, 1.0]);
/// ```
pub fn column_min(data: &[Vec<f64>]) -> Vec<f64> {
    if data.is_empty() {
        return vec![];
    }
    let cols = data[0].len();
    #[cfg(feature = "rayon")]
    {
        (0..cols)
            .into_par_iter()
            .map(|j| data.iter().map(|r| r[j]).fold(f64::INFINITY, f64::min))
            .collect()
    }
    #[cfg(not(feature = "rayon"))]
    {
        (0..cols)
            .map(|j| data.iter().map(|r| r[j]).fold(f64::INFINITY, f64::min))
            .collect()
    }
}

/// Returns the maximum value of each column.
///
/// ```rust
/// use datarust::stats::column_max;
///
/// let data = vec![vec![3.0, 1.0], vec![1.0, 5.0], vec![2.0, 2.0]];
/// let maxs = column_max(&data);
/// assert_eq!(maxs, vec![3.0, 5.0]);
/// ```
pub fn column_max(data: &[Vec<f64>]) -> Vec<f64> {
    if data.is_empty() {
        return vec![];
    }
    let cols = data[0].len();
    #[cfg(feature = "rayon")]
    {
        (0..cols)
            .into_par_iter()
            .map(|j| data.iter().map(|r| r[j]).fold(f64::NEG_INFINITY, f64::max))
            .collect()
    }
    #[cfg(not(feature = "rayon"))]
    {
        (0..cols)
            .map(|j| data.iter().map(|r| r[j]).fold(f64::NEG_INFINITY, f64::max))
            .collect()
    }
}

/// Sum of the values in a slice.
///
/// The 1-D counterpart of [`column_sum`]. Returns `0.0` for an empty slice.
pub fn sum(data: &[f64]) -> f64 {
    data.iter().sum()
}

/// Arithmetic mean (average) of a slice.
///
/// The 1-D counterpart of [`column_mean`]. Returns [`f64::NAN`] for an empty
/// slice, matching numpy.
///
/// ```rust
/// use datarust::stats::mean;
/// assert!((mean(&[1.0, 2.0, 3.0, 4.0]) - 2.5).abs() < 1e-12);
/// ```
pub fn mean(data: &[f64]) -> f64 {
    if data.is_empty() {
        return f64::NAN;
    }
    sum(data) / data.len() as f64
}

/// Minimum value of a slice.
///
/// The 1-D counterpart of [`column_min`]. Returns [`f64::INFINITY`] for an
/// empty slice (the identity element of `min`), so reducing a non-empty slice
/// folded over it never overflows.
pub fn min(data: &[f64]) -> f64 {
    data.iter().copied().fold(f64::INFINITY, f64::min)
}

/// Maximum value of a slice.
///
/// The 1-D counterpart of [`column_max`]. Returns [`f64::NEG_INFINITY`] for an
/// empty slice (the identity element of `max`).
pub fn max(data: &[f64]) -> f64 {
    data.iter().copied().fold(f64::NEG_INFINITY, f64::max)
}

/// Variance of a slice using the given delta degrees of freedom.
///
/// The 1-D counterpart of [`column_variance`]. `ddof = 0` yields the
/// population variance, `ddof = 1` the sample variance (the numpy default).
/// Returns [`f64::NAN`] for an empty slice or when `ddof >= n` (a
/// non-positive denominator), mirroring [`column_variance`].
///
/// ```rust
/// use datarust::stats::variance;
///
/// let v = variance(&[1.0, 2.0, 3.0, 4.0, 5.0], 1);
/// assert!((v - 2.5).abs() < 1e-12);
/// ```
pub fn variance(data: &[f64], ddof: usize) -> f64 {
    let n = data.len();
    if n == 0 || ddof >= n {
        return f64::NAN;
    }
    let denom = n.saturating_sub(ddof) as f64;
    let m = mean(data);
    let s: f64 = data.iter().map(|&x| (x - m) * (x - m)).sum();
    s / denom
}

/// Standard deviation of a slice using the given delta degrees of freedom.
///
/// The 1-D counterpart of [`column_std`]. Simply `variance(data, ddof).sqrt()`;
/// see [`variance`] for the `ddof` and empty-slice semantics.
///
/// ```rust
/// use datarust::stats::std;
///
/// let s = std(&[1.0, 2.0, 3.0, 4.0, 5.0], 1);
/// assert!((s - 1.5811).abs() < 1e-3);
/// ```
pub fn std(data: &[f64], ddof: usize) -> f64 {
    variance(data, ddof).sqrt()
}

/// Median of a slice that is assumed to be sorted in non-decreasing order.
///
/// Returns `None` for an empty slice instead of panicking. Callers that can
/// guarantee a non-empty slice may safely [`Option::unwrap`] the result.
///
/// ```rust
/// use datarust::stats::median_sorted;
///
/// let sorted = vec![1.0, 2.0, 3.0, 4.0, 5.0];
/// assert!((median_sorted(&sorted).unwrap() - 3.0).abs() < 1e-12);
/// ```
pub fn median_sorted(sorted: &[f64]) -> Option<f64> {
    let n = sorted.len();
    if n == 0 {
        return None;
    }
    if n % 2 == 1 {
        Some(sorted[n / 2])
    } else {
        Some((sorted[n / 2 - 1] + sorted[n / 2]) / 2.0)
    }
}

/// Quantile with linear interpolation, matching numpy's default ("linear") method.
///
/// Returns `None` if the slice is empty or if `q` is outside `[0, 1]`.
pub fn quantile(sorted: &[f64], q: f64) -> Option<f64> {
    let n = sorted.len();
    if n == 0 || !(0.0..=1.0).contains(&q) {
        return None;
    }
    if n == 1 {
        return Some(sorted[0]);
    }
    let pos = q * (n - 1) as f64;
    let lo = pos.floor() as usize;
    let hi = pos.ceil() as usize;
    if lo == hi {
        return Some(sorted[lo]);
    }
    let frac = pos - lo as f64;
    Some(sorted[lo] * (1.0 - frac) + sorted[hi] * frac)
}

/// Median (50th percentile) of a slice.
///
/// Unlike [`median_sorted`], this sorts a copy of the input internally, so the
/// caller does not need to pre-sort. Returns `None` for an empty slice.
///
/// ```rust
/// use datarust::stats::median;
/// assert!((median(&[3.0, 1.0, 2.0]).unwrap() - 2.0).abs() < 1e-12);
/// assert!((median(&[1.0, 2.0, 3.0, 4.0]).unwrap() - 2.5).abs() < 1e-12);
/// ```
pub fn median(data: &[f64]) -> Option<f64> {
    if data.is_empty() {
        return None;
    }
    let mut sorted = data.to_vec();
    sorted.sort_by(|a, b| a.total_cmp(b));
    median_sorted(&sorted)
}

/// Returns the requested quantile of each column using linear interpolation.
pub fn quantile_column(data: &[Vec<f64>], q: f64) -> Result<Vec<f64>> {
    if !(0.0..=1.0).contains(&q) {
        return Err(DatarustError::InvalidInput(format!(
            "quantile q must be in [0, 1], got {}",
            q
        )));
    }
    if data.is_empty() {
        return Ok(vec![]);
    }
    let cols = data[0].len();
    #[cfg(feature = "rayon")]
    {
        Ok((0..cols)
            .into_par_iter()
            .map(|j| {
                let mut col: Vec<f64> = data.iter().map(|r| r[j]).collect();
                col.sort_by(|a, b| a.total_cmp(b));
                quantile(&col, q).expect("non-empty column with q in [0,1]")
            })
            .collect())
    }
    #[cfg(not(feature = "rayon"))]
    {
        Ok((0..cols)
            .map(|j| {
                let mut col: Vec<f64> = data.iter().map(|r| r[j]).collect();
                col.sort_by(|a, b| a.total_cmp(b));
                quantile(&col, q).expect("non-empty column with q in [0,1]")
            })
            .collect())
    }
}

/// Returns the median of each column.
pub fn median_column(data: &[Vec<f64>]) -> Vec<f64> {
    if data.is_empty() {
        return vec![];
    }
    let cols = data[0].len();
    #[cfg(feature = "rayon")]
    {
        (0..cols)
            .into_par_iter()
            .map(|j| {
                let mut col: Vec<f64> = data.iter().map(|r| r[j]).collect();
                col.sort_by(|a, b| a.total_cmp(b));
                // INVARIANT: `data` is non-empty (checked above), so each column is non-empty.
                median_sorted(&col).expect("non-empty column")
            })
            .collect()
    }
    #[cfg(not(feature = "rayon"))]
    {
        (0..cols)
            .map(|j| {
                let mut col: Vec<f64> = data.iter().map(|r| r[j]).collect();
                col.sort_by(|a, b| a.total_cmp(b));
                median_sorted(&col).expect("non-empty column")
            })
            .collect()
    }
}

/// Most frequent value. Ties broken by smallest value (deterministic).
///
/// A column of all-equal (or empty) entries yields [`f64::NAN`] for that column.
///
/// Uses the same sort-then-scan strategy as [`mode`]: each column is gathered
/// once, sorted, and the longest run of equal bit patterns is scanned. This is
/// dramatically faster than hash-map counting on large mostly-distinct columns
/// (the previous `HashMap<u64, …>` grew to one entry per distinct value) and
/// breaks ties identically (smallest value wins).
pub fn mode_column(data: &[Vec<f64>]) -> Vec<f64> {
    if data.is_empty() {
        return vec![];
    }
    let cols = data[0].len();
    #[cfg(feature = "rayon")]
    {
        (0..cols)
            .into_par_iter()
            .map(|j| mode_column_one(data, j))
            .collect()
    }
    #[cfg(not(feature = "rayon"))]
    {
        (0..cols).map(|j| mode_column_one(data, j)).collect()
    }
}

/// Sorted-run mode of a single gathered column, returning [`f64::NAN`] for an
/// empty column (matching [`mode_column`]'s contract).
fn mode_column_one(data: &[Vec<f64>], j: usize) -> f64 {
    let mut col: Vec<f64> = data.iter().map(|r| r[j]).collect();
    if col.is_empty() {
        return f64::NAN;
    }
    col.sort_by(|a, b| a.total_cmp(b));
    mode_sorted(&col).unwrap_or(f64::NAN)
}

/// Most frequent value of a slice. Ties are broken by the smallest value
/// (deterministic), matching [`mode_column`].
///
/// The 1-D counterpart of [`mode_column`]. Returns `None` for an empty slice.
/// A slice where every value is distinct returns any single element (all are
/// tied at count 1; the smallest wins).
pub fn mode(data: &[f64]) -> Option<f64> {
    if data.is_empty() {
        return None;
    }
    // Sort-then-scan: one copy + one sort, then find the longest run of equal
    // bit patterns via [`mode_sorted`]. This is dramatically faster than
    // hash-map counting on large mostly-distinct columns (e.g. imputing
    // continuous features), where the old `HashMap<u64, …>` grew to one entry
    // per distinct value. Ties are broken by the smallest value because the
    // scan keeps the first (sorted) value that achieves the maximum run length.
    let mut sorted = data.to_vec();
    sorted.sort_by(|a, b| a.total_cmp(b));
    mode_sorted(&sorted)
}

/// Scans an ascending-sorted slice for the longest run of equal bit patterns,
/// returning the first value that achieves the maximum run length.
///
/// Because the input is sorted ascending, ties are broken toward the smallest
/// value. Returns `None` for an empty slice.
fn mode_sorted(sorted: &[f64]) -> Option<f64> {
    let n = sorted.len();
    if n == 0 {
        return None;
    }
    let mut best_val = sorted[0];
    let mut best_len = 1usize;
    let mut run_val = sorted[0];
    let mut run_len = 1usize;
    for &x in &sorted[1..] {
        if x.to_bits() == run_val.to_bits() {
            run_len += 1;
        } else if run_len > best_len {
            best_len = run_len;
            best_val = run_val;
            run_val = x;
            run_len = 1;
        } else {
            run_val = x;
            run_len = 1;
        }
    }
    if run_len > best_len {
        best_val = run_val;
    }
    Some(best_val)
}

/// Sum of each column.
pub fn column_sum(data: &[Vec<f64>]) -> Vec<f64> {
    if data.is_empty() {
        return vec![];
    }
    let cols = data[0].len();
    let mut sums = vec![0.0; cols];
    for row in data {
        for (j, &v) in row.iter().enumerate() {
            sums[j] += v;
        }
    }
    sums
}

/// Per-column mean and variance in a single row-major sweep using Welford's
/// online algorithm.
///
/// This replaces the previous `column_mean` + `column_variance` pair (which
/// made three full passes over the data: two for the mean, one for the
/// variance) with a single fused, numerically stable pass. Returns
/// `(means, variances)` where the variance uses the supplied delta degrees of
/// freedom. If the denominator is non-positive (`ddof >= n`), variances are
/// `NaN`, mirroring [`column_variance`].
pub fn column_mean_var(data: &[Vec<f64>], ddof: usize) -> (Vec<f64>, Vec<f64>) {
    if data.is_empty() {
        return (vec![], vec![]);
    }
    let cols = data[0].len();
    let n = data.len();
    let mut mean = vec![0.0; cols];
    let mut m2 = vec![0.0; cols];
    // Welford: for each x, count++, delta = x - mean, mean += delta/count,
    // m2 += delta * (x - mean). Single row-major pass, cache-friendly.
    for (count, row) in data.iter().enumerate() {
        let c = (count + 1) as f64;
        for (j, &x) in row.iter().enumerate() {
            let delta = x - mean[j];
            mean[j] += delta / c;
            m2[j] += delta * (x - mean[j]);
        }
    }
    let denom = n.saturating_sub(ddof) as f64;
    let var = if denom > 0.0 {
        m2.iter().map(|&m| m / denom).collect()
    } else {
        vec![f64::NAN; cols]
    };
    (mean, var)
}

/// Flat-storage counterpart of [`column_mean_var`]: single fused Welford pass
/// over contiguous row-major `data` of shape `rows × cols`. Used by scalers
/// that operate directly on the flat buffer returned by
/// [`Matrix::as_slice`](crate::matrix::Matrix::as_slice).
pub fn column_mean_var_flat(
    data: &[f64],
    rows: usize,
    cols: usize,
    ddof: usize,
) -> (Vec<f64>, Vec<f64>) {
    if rows == 0 || cols == 0 {
        return (vec![], vec![]);
    }
    let mut mean = vec![0.0; cols];
    let mut m2 = vec![0.0; cols];
    for count in 0..rows {
        let c = (count + 1) as f64;
        let base = count * cols;
        for j in 0..cols {
            let x = data[base + j];
            let delta = x - mean[j];
            mean[j] += delta / c;
            m2[j] += delta * (x - mean[j]);
        }
    }
    let denom = rows.saturating_sub(ddof) as f64;
    let var = if denom > 0.0 {
        m2.iter().map(|&m| m / denom).collect()
    } else {
        vec![f64::NAN; cols]
    };
    (mean, var)
}

/// Per-column minimum and maximum in a single fused row-major pass.
///
/// Replaces the separate [`column_min`] + [`column_max`] pair (two passes)
/// with one pass, halving memory traffic on row-major data.
pub fn column_min_max(data: &[Vec<f64>]) -> (Vec<f64>, Vec<f64>) {
    if data.is_empty() {
        return (vec![], vec![]);
    }
    let cols = data[0].len();
    let mut min = vec![f64::INFINITY; cols];
    let mut max = vec![f64::NEG_INFINITY; cols];
    for row in data {
        for (j, &v) in row.iter().enumerate() {
            if v < min[j] {
                min[j] = v;
            }
            if v > max[j] {
                max[j] = v;
            }
        }
    }
    (min, max)
}

/// Flat-storage counterpart of [`column_min_max`].
pub fn column_min_max_flat(data: &[f64], rows: usize, cols: usize) -> (Vec<f64>, Vec<f64>) {
    if rows == 0 || cols == 0 {
        return (vec![], vec![]);
    }
    let mut min = vec![f64::INFINITY; cols];
    let mut max = vec![f64::NEG_INFINITY; cols];
    for i in 0..rows {
        let base = i * cols;
        for j in 0..cols {
            let v = data[base + j];
            if v < min[j] {
                min[j] = v;
            }
            if v > max[j] {
                max[j] = v;
            }
        }
    }
    (min, max)
}

/// Computes exact quantiles of a single (unsorted) column.
///
/// With few distinct order statistics the column is introselected with
/// `select_nth_unstable_by` instead of fully sorted; results are bit-identical
/// to the full-sort path because both use the same `total_cmp` order. When
/// many quantiles are requested the full sort remains cheaper.
fn column_quantiles_select(col: &mut [f64], qs: &[f64]) -> Vec<f64> {
    let n = col.len();
    let mut idx: Vec<usize> = Vec::with_capacity(2 * qs.len());
    for &q in qs {
        let pos = q * (n - 1) as f64;
        idx.push(pos.floor() as usize);
        idx.push(pos.ceil() as usize);
    }
    idx.sort_unstable();
    idx.dedup();
    if idx.len() > 32 {
        col.sort_by(|a, b| a.total_cmp(b));
        return qs
            .iter()
            .map(|&q| quantile(col, q).expect("non-empty column with q in [0,1]"))
            .collect();
    }
    for &k in idx.iter() {
        col.select_nth_unstable_by(k, |a, b| a.total_cmp(b));
    }
    qs.iter()
        .map(|&q| {
            let pos = q * (n - 1) as f64;
            let lo = pos.floor() as usize;
            let hi = pos.ceil() as usize;
            if lo == hi {
                col[lo]
            } else {
                let frac = pos - lo as f64;
                col[lo] * (1.0 - frac) + col[hi] * frac
            }
        })
        .collect()
}

/// Multiple quantiles of each column computed from a single selection per column.
///
/// For each column the values are gathered once, then every requested quantile
/// is read off the same working buffer via linear interpolation. This replaces
/// the previous pattern of calling [`quantile_column`] (or [`median_column`])
/// separately for each quantile, which re-sorted the same column data
/// redundantly.
///
/// `qs` must be in `[0, 1]`; an empty `qs` yields an empty `Vec` per column.
/// The result has shape `qs.len() × cols` (one row per requested quantile).
pub fn column_quantiles_many(data: &[Vec<f64>], qs: &[f64]) -> Result<Vec<Vec<f64>>> {
    if qs.iter().any(|&q| !(0.0..=1.0).contains(&q)) {
        return Err(DatarustError::InvalidInput(format!(
            "quantiles must be in [0, 1], got {:?}",
            qs
        )));
    }
    if data.is_empty() {
        return Ok(vec![vec![]; qs.len()]);
    }
    let cols = data[0].len();
    let nqs = qs.len();
    let mut out: Vec<Vec<f64>> = (0..nqs).map(|_| Vec::with_capacity(cols)).collect();
    #[cfg(feature = "rayon")]
    let iter = (0..cols).into_par_iter();
    #[cfg(not(feature = "rayon"))]
    let iter = 0..cols;
    // For each column: select/sort once, then read off every requested quantile.
    let per_col: Vec<Vec<f64>> = iter
        .map(|j| {
            let mut col: Vec<f64> = data.iter().map(|r| r[j]).collect();
            column_quantiles_select(&mut col, qs)
        })
        .collect();
    // Transpose per_col (cols × nqs) into out (nqs × cols).
    for col_vals in per_col {
        for (qi, v) in col_vals.into_iter().enumerate() {
            out[qi].push(v);
        }
    }
    Ok(out)
}

/// Flat-storage counterpart of [`column_quantiles_many`].
///
/// Each column is gathered from the contiguous `data` buffer once, its order
/// statistics are introselected once, and every requested quantile is read off
/// the same working buffer. Result shape is `qs.len() × cols`.
pub fn column_quantiles_many_flat(
    data: &[f64],
    rows: usize,
    cols: usize,
    qs: &[f64],
) -> Result<Vec<Vec<f64>>> {
    if qs.iter().any(|&q| !(0.0..=1.0).contains(&q)) {
        return Err(DatarustError::InvalidInput(format!(
            "quantiles must be in [0, 1], got {:?}",
            qs
        )));
    }
    if rows == 0 || cols == 0 {
        return Ok(vec![vec![]; qs.len()]);
    }
    let nqs = qs.len();
    let mut out: Vec<Vec<f64>> = (0..nqs).map(|_| Vec::with_capacity(cols)).collect();
    #[cfg(feature = "rayon")]
    let iter = (0..cols).into_par_iter();
    #[cfg(not(feature = "rayon"))]
    let iter = 0..cols;
    let per_col: Vec<Vec<f64>> = iter
        .map(|j| {
            let mut col: Vec<f64> = (0..rows).map(|i| data[i * cols + j]).collect();
            column_quantiles_select(&mut col, qs)
        })
        .collect();
    for col_vals in per_col {
        for (qi, v) in col_vals.into_iter().enumerate() {
            out[qi].push(v);
        }
    }
    Ok(out)
}

/// Covariance of already-centered data: `(1/(n-ddof)) * Xcᵀ Xc`.
///
/// This is the single canonical centered-covariance routine shared by
/// [`covariance_matrix`] (raw-data entry point), PCA and Truncated SVD.
/// `x_centered` is row-major `n × p`. A non-positive denominator (`ddof >= n`)
/// leaves the scale unchanged rather than producing infinities.
///
/// When the `matrixmultiply` feature is enabled, the `Xcᵀ Xc` product is
/// computed with a tuned pure-Rust GEMM (no system BLAS); otherwise a scalar
/// row-major accumulation is used.
#[allow(clippy::needless_range_loop)]
pub(crate) fn covariance_centered(x_centered: &[Vec<f64>], ddof: usize) -> Vec<Vec<f64>> {
    let n = x_centered.len();
    let p = if n > 0 { x_centered[0].len() } else { 0 };

    #[cfg(feature = "matrixmultiply")]
    {
        if n > 0 && p > 0 {
            return covariance_centered_gemm(x_centered, n, p, ddof);
        }
    }

    let mut cov = vec![vec![0.0; p]; p];
    for row in x_centered {
        for i in 0..p {
            let xi = row[i];
            if xi == 0.0 {
                continue;
            }
            for j in 0..p {
                cov[i][j] += xi * row[j];
            }
        }
    }
    let denom = n.saturating_sub(ddof) as f64;
    if denom > 0.0 {
        let inv = 1.0 / denom;
        for i in 0..p {
            for j in 0..p {
                cov[i][j] *= inv;
            }
        }
    }
    cov
}

/// Flat-storage centered covariance: `C = (1/(n-ddof)) · Xcᵀ · Xc`.
///
/// `x_centered` is a flat row-major buffer of shape `n × p` (length `n*p`).
/// Returns a flat row-major `p × p` covariance matrix. A non-positive
/// denominator (`ddof >= n`) leaves the scale unchanged.
#[allow(clippy::needless_range_loop)]
pub(crate) fn covariance_centered_flat(
    x_centered: &[f64],
    n: usize,
    p: usize,
    ddof: usize,
) -> Vec<Vec<f64>> {
    if n == 0 || p == 0 {
        return vec![];
    }
    #[cfg(feature = "matrixmultiply")]
    {
        covariance_centered_flat_gemm(x_centered, n, p, ddof)
    }
    #[cfg(not(feature = "matrixmultiply"))]
    {
        covariance_centered_flat_scalar(x_centered, n, p, ddof)
    }
}

#[cfg(not(feature = "matrixmultiply"))]
#[allow(clippy::needless_range_loop)]
fn covariance_centered_flat_scalar(
    x_centered: &[f64],
    n: usize,
    p: usize,
    ddof: usize,
) -> Vec<Vec<f64>> {
    let mut cov = vec![vec![0.0; p]; p];
    for i in 0..n {
        let base = i * p;
        for a in 0..p {
            let xi = x_centered[base + a];
            if xi == 0.0 {
                continue;
            }
            // Accumulate only the lower triangle (b in 0..=a) and mirror it
            // afterwards: roughly halves the multiply-adds versus the full
            // p×p product. Per-element accumulation order is unchanged (each
            // cov[a][b] still receives one term per row in row order), so the
            // lower triangle is bit-identical to the previous full loop.
            let row = &mut cov[a];
            for (slot, &xb) in row
                .iter_mut()
                .take(a + 1)
                .zip(&x_centered[base..base + a + 1])
            {
                *slot += xi * xb;
            }
        }
    }
    // Mirror the lower triangle into the upper triangle.
    for a in 0..p {
        for b in (a + 1)..p {
            cov[a][b] = cov[b][a];
        }
    }
    let denom = n.saturating_sub(ddof) as f64;
    if denom > 0.0 {
        let inv = 1.0 / denom;
        for row in cov.iter_mut() {
            for v in row.iter_mut() {
                *v *= inv;
            }
        }
    }
    cov
}

/// GEMM-backed flat centered covariance.
#[cfg(feature = "matrixmultiply")]
fn covariance_centered_flat_gemm(
    x_centered: &[f64],
    n: usize,
    p: usize,
    ddof: usize,
) -> Vec<Vec<f64>> {
    use matrixmultiply::dgemm;
    let mut cov_flat = vec![0.0; p * p];
    // C(p×p) = 1.0 * Xcᵀ(p×n) · Xc(n×p) + 0.0 * C (see covariance_centered_gemm
    // for the stride rationale).
    unsafe {
        dgemm(
            p,
            n,
            p,
            1.0,
            x_centered.as_ptr(),
            1,
            p as isize,
            x_centered.as_ptr(),
            p as isize,
            1,
            0.0,
            cov_flat.as_mut_ptr(),
            p as isize,
            1,
        );
    }
    let denom = n.saturating_sub(ddof) as f64;
    let mut cov: Vec<Vec<f64>> = cov_flat.chunks_exact(p).map(|row| row.to_vec()).collect();
    if denom > 0.0 {
        let inv = 1.0 / denom;
        for row in cov.iter_mut() {
            for v in row.iter_mut() {
                *v *= inv;
            }
        }
    }
    cov
}

/// `matrixmultiply`-backed centered covariance: `C = (1/(n-ddof)) · Xcᵀ · Xc`.
///
/// `Xc` is row-major `n × p`. We compute `C = Xcᵀ · Xc` as a single `dgemm`.
/// The same flat buffer is used for both operands: as `A` it is read with the
/// strides of a column-major `Xcᵀ` (rsa=1, csa=p), as `B` with the strides of
/// a row-major `Xc` (rsb=p, csb=1). The result lands in row-major `C` (p×p).
#[cfg(feature = "matrixmultiply")]
fn covariance_centered_gemm(
    x_centered: &[Vec<f64>],
    n: usize,
    p: usize,
    ddof: usize,
) -> Vec<Vec<f64>> {
    use matrixmultiply::dgemm;
    // Flatten the centered data once (rows may live in separate Vecs).
    let mut flat = Vec::with_capacity(n * p);
    for row in x_centered {
        flat.extend_from_slice(row);
    }
    let mut cov_flat = vec![0.0; p * p];
    // dgemm signature: dgemm(m, k, n, alpha, a, rsa, csa, b, rsb, csb, beta, c, rsc, csc)
    // computes C(m×n) = alpha * A(m×k) · B(k×n) + beta * C.
    // Here C(p×p) = 1.0 * Xcᵀ(p×n) · Xc(n×p) + 0.0 * C, so m=p, k=n, n=p.
    unsafe {
        dgemm(
            p,                     // m: rows of A (= Xcᵀ) and of C
            n,                     // k: inner dimension (rows of Xc)
            p,                     // n: cols of B (= Xc) and of C
            1.0,                   // alpha
            flat.as_ptr(),         // A = Xcᵀ read col-major over flat
            1,                     // rsa: stride between rows of A — adjacent in a column of Xcᵀ
            p as isize, // csa: stride between cols of A — one column of Xcᵀ = one row of Xc
            flat.as_ptr(), // B = Xc read row-major
            p as isize, // rsb: stride between rows of B (one row of Xc)
            1,          // csb: stride between cols of B (adjacent in a row of Xc)
            0.0,        // beta
            cov_flat.as_mut_ptr(), // C row-major p×p
            p as isize, // rsc
            1,          // csc
        );
    }
    let denom = n.saturating_sub(ddof) as f64;
    let mut cov: Vec<Vec<f64>> = cov_flat.chunks_exact(p).map(|row| row.to_vec()).collect();
    if denom > 0.0 {
        let inv = 1.0 / denom;
        for row in cov.iter_mut() {
            for v in row.iter_mut() {
                *v *= inv;
            }
        }
    }
    cov
}

/// Covariance matrix (p × p) for an n × p data matrix.
///
/// `ddof=0` gives population covariance, `ddof=1` gives sample covariance.
///
/// ```rust
/// use datarust::stats::covariance_matrix;
///
/// let data = vec![vec![1.0, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]];
/// let cov = covariance_matrix(&data, 0);
/// assert_eq!(cov.len(), 2);
/// assert_eq!(cov[0].len(), 2);
/// ```
#[allow(clippy::needless_range_loop)]
pub fn covariance_matrix(data: &[Vec<f64>], ddof: usize) -> Vec<Vec<f64>> {
    if data.is_empty() {
        return vec![];
    }
    let means = column_mean(data);
    // Center the data, then delegate to the shared centered-covariance routine.
    let centered: Vec<Vec<f64>> = data
        .iter()
        .map(|row| row.iter().enumerate().map(|(j, &v)| v - means[j]).collect())
        .collect();
    covariance_centered(&centered, ddof)
}

/// Pearson correlation matrix (p × p).
///
/// ```rust
/// use datarust::stats::correlation_matrix;
///
/// let data = vec![vec![1.0, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]];
/// let corr = correlation_matrix(&data);
/// assert!((corr[0][1] - 1.0).abs() < 1e-12); // perfect correlation
/// ```
pub fn correlation_matrix(data: &[Vec<f64>]) -> Vec<Vec<f64>> {
    if data.is_empty() {
        return vec![];
    }
    let p = data[0].len();
    let cov = covariance_matrix(data, 1);
    let std: Vec<f64> = (0..p).map(|j| cov[j][j].sqrt()).collect();
    let mut corr = vec![vec![0.0; p]; p];
    for i in 0..p {
        for j in i..p {
            let v = if std[i] == 0.0 || std[j] == 0.0 {
                if i == j {
                    1.0
                } else {
                    0.0
                }
            } else {
                cov[i][j] / (std[i] * std[j])
            };
            corr[i][j] = v;
            corr[j][i] = v;
        }
    }
    corr
}

/// Pearson correlation matrix (`p × p`) from flat row-major data.
///
/// This is the flat-storage counterpart of [`correlation_matrix`]: it accepts
/// a contiguous row-major buffer of shape `rows × cols` instead of a nested
/// `Vec<Vec<f64>>`. Column means, centering, and the covariance accumulation
/// all stream over contiguous memory, which is substantially faster than the
/// nested variant on wide tables (where gathering each column across scattered
/// rows dominates). The `matrixmultiply` feature accelerates the covariance
/// product automatically, just like the nested path.
///
/// ```rust
/// use datarust::stats::correlation_matrix_flat;
///
/// // Three rows, two columns; column 1 = 2 × column 0 → r = 1.
/// let flat = vec![1.0, 2.0, 2.0, 4.0, 3.0, 6.0];
/// let corr = correlation_matrix_flat(&flat, 3, 2);
/// assert!((corr[0][1] - 1.0).abs() < 1e-12);
/// ```
pub fn correlation_matrix_flat(data: &[f64], rows: usize, cols: usize) -> Vec<Vec<f64>> {
    if rows == 0 || cols == 0 {
        return vec![];
    }
    let means = column_mean_flat(data, rows, cols);
    let mut centered = vec![0.0; rows * cols];
    for i in 0..rows {
        let base = i * cols;
        for (j, m) in means.iter().enumerate() {
            centered[base + j] = data[base + j] - m;
        }
    }
    let cov = covariance_centered_flat(&centered, rows, cols, 1);
    let std: Vec<f64> = (0..cols).map(|j| cov[j][j].sqrt()).collect();
    let mut corr = vec![vec![0.0; cols]; cols];
    for i in 0..cols {
        for j in i..cols {
            let v = if std[i] == 0.0 || std[j] == 0.0 {
                if i == j {
                    1.0
                } else {
                    0.0
                }
            } else {
                cov[i][j] / (std[i] * std[j])
            };
            corr[i][j] = v;
            corr[j][i] = v;
        }
    }
    corr
}

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

    #[test]
    fn mean_basic() {
        let data = vec![vec![1.0, 10.0], vec![3.0, 20.0], vec![5.0, 30.0]];
        let m = column_mean(&data);
        assert!((m[0] - 3.0).abs() < 1e-12);
        assert!((m[1] - 20.0).abs() < 1e-12);
    }

    #[test]
    fn variance_ddof() {
        let data = vec![vec![1.0, 2.0, 3.0, 4.0]];
        let t = transpose(&data);
        let v0 = column_variance(&t, 0);
        let v1 = column_variance(&t, 1);
        // population variance = 1.25 ; sample = 1.666...
        assert!((v0[0] - 1.25).abs() < 1e-12);
        assert!((v1[0] - (5.0 / 3.0)).abs() < 1e-12);
    }

    #[test]
    fn quantile_linear() {
        // numpy quantile linear for [0,1,2,3,4]: q=0.5 -> 2, q=0.25 -> 1, q=0.75 -> 3
        let s = [0.0_f64, 1.0, 2.0, 3.0, 4.0];
        assert!((quantile(&s, 0.5).unwrap() - 2.0).abs() < 1e-12);
        assert!((quantile(&s, 0.25).unwrap() - 1.0).abs() < 1e-12);
        assert!((quantile(&s, 0.75).unwrap() - 3.0).abs() < 1e-12);
        // q=0.3 -> 0.3*4 = 1.2 -> interp between idx1 and idx2
        assert!((quantile(&s, 0.3).unwrap() - 1.2).abs() < 1e-12);
    }

    #[test]
    fn quantile_edge() {
        let s = [5.0_f64];
        assert!((quantile(&s, 0.5).unwrap() - 5.0).abs() < 1e-12);
        assert!((quantile(&s, 0.0).unwrap() - 5.0).abs() < 1e-12);
    }

    #[test]
    fn quantile_none_cases() {
        assert!(quantile(&[], 0.5).is_none());
        assert!(quantile(&[1.0, 2.0], 1.5).is_none());
        assert!(quantile(&[1.0, 2.0], -0.1).is_none());
        assert!(median_sorted(&[]).is_none());
    }

    #[test]
    fn median_even_odd() {
        assert!((median_sorted(&[1.0_f64, 2.0, 3.0]).unwrap() - 2.0).abs() < 1e-12);
        assert!((median_sorted(&[1.0_f64, 2.0, 3.0, 4.0]).unwrap() - 2.5).abs() < 1e-12);
    }

    #[test]
    fn mode_simple() {
        let data = vec![vec![1.0], vec![2.0], vec![2.0], vec![3.0]];
        let m = mode_column(&data);
        assert!((m[0] - 2.0).abs() < 1e-12);
    }

    #[test]
    fn mode_tie_smallest() {
        // tie between 1.0 and 2.0 -> smallest wins
        let data = vec![vec![1.0], vec![2.0], vec![1.0], vec![2.0]];
        let m = mode_column(&data);
        assert!((m[0] - 1.0).abs() < 1e-12);
    }

    #[test]
    fn min_max() {
        let data = vec![vec![3.0, -1.0], vec![5.0, 2.0], vec![1.0, 0.0]];
        let mn = column_min(&data);
        let mx = column_max(&data);
        assert!((mn[0] - 1.0).abs() < 1e-12);
        assert!((mx[0] - 5.0).abs() < 1e-12);
        assert!((mn[1] - -1.0).abs() < 1e-12);
    }

    #[test]
    fn basic_column_helpers_handle_values_and_empty_inputs() {
        let data = vec![vec![1.0, 6.0], vec![3.0, 2.0], vec![5.0, 4.0]];

        assert_eq!(column_mean(&data), vec![3.0, 4.0]);
        assert_eq!(
            column_mean_flat(&[1.0, 6.0, 3.0, 2.0, 5.0, 4.0], 3, 2),
            vec![3.0, 4.0]
        );
        assert_eq!(column_min(&data), vec![1.0, 2.0]);
        assert_eq!(column_max(&data), vec![5.0, 6.0]);
        assert_eq!(column_std(&data, 1), vec![2.0, 2.0]);

        assert!(column_variance(&data, 99)
            .iter()
            .all(|value| value.is_nan()));
        assert!(column_std(&data, 99).iter().all(|value| value.is_nan()));
        assert!(column_mean(&[]).is_empty());
        assert!(column_mean_flat(&[], 0, 0).is_empty());
        assert!(column_variance(&[], 0).is_empty());
        assert!(column_min(&[]).is_empty());
        assert!(column_max(&[]).is_empty());
    }

    #[test]
    fn column_quantile_and_median_validate_and_preserve_column_order() {
        let data = vec![
            vec![1.0, 4.0],
            vec![3.0, 2.0],
            vec![5.0, 6.0],
            vec![7.0, 0.0],
        ];

        assert_eq!(quantile_column(&data, 0.25).unwrap(), vec![2.5, 1.5]);
        assert_eq!(median_column(&data), vec![4.0, 3.0]);
        assert!(quantile_column(&data, -0.1).is_err());
        assert!(quantile_column(&[], 0.5).unwrap().is_empty());
        assert!(median_column(&[]).is_empty());
    }

    #[test]
    fn fused_column_helpers_match_individual_operations() {
        let data = vec![
            vec![1.0, 4.0],
            vec![3.0, 2.0],
            vec![5.0, 6.0],
            vec![7.0, 0.0],
        ];
        let flat = [1.0, 4.0, 3.0, 2.0, 5.0, 6.0, 7.0, 0.0];

        let (means, variances) = column_mean_var(&data, 1);
        assert_eq!(means, column_mean(&data));
        assert_eq!(variances, column_variance(&data, 1));

        let (flat_means, flat_variances) = column_mean_var_flat(&flat, 4, 2, 1);
        assert_eq!(flat_means, means);
        assert_eq!(flat_variances, variances);

        let (mins, maxes) = column_min_max(&data);
        assert_eq!(mins, column_min(&data));
        assert_eq!(maxes, column_max(&data));
        assert_eq!(column_min_max_flat(&flat, 4, 2), (mins, maxes));

        let (_, invalid_variances) = column_mean_var(&data, 99);
        assert!(invalid_variances.iter().all(|value| value.is_nan()));
        let (_, invalid_flat_variances) = column_mean_var_flat(&flat, 4, 2, 99);
        assert!(invalid_flat_variances.iter().all(|value| value.is_nan()));
        assert_eq!(column_mean_var(&[], 0), (vec![], vec![]));
        assert_eq!(column_mean_var_flat(&[], 0, 0, 0), (vec![], vec![]));
        assert_eq!(column_min_max(&[]), (vec![], vec![]));
        assert_eq!(column_min_max_flat(&[], 0, 0), (vec![], vec![]));
    }

    #[test]
    fn multiple_column_quantiles_match_flat_storage_and_validate_input() {
        let data = vec![
            vec![1.0, 4.0],
            vec![3.0, 2.0],
            vec![5.0, 6.0],
            vec![7.0, 0.0],
        ];
        let flat = [1.0, 4.0, 3.0, 2.0, 5.0, 6.0, 7.0, 0.0];
        let qs = [0.0, 0.5, 1.0];
        let expected = vec![vec![1.0, 0.0], vec![4.0, 3.0], vec![7.0, 6.0]];

        assert_eq!(column_quantiles_many(&data, &qs).unwrap(), expected);
        assert_eq!(
            column_quantiles_many_flat(&flat, 4, 2, &qs).unwrap(),
            expected
        );
        assert!(column_quantiles_many(&data, &[1.1]).is_err());
        assert!(column_quantiles_many_flat(&flat, 4, 2, &[-0.1]).is_err());
        assert_eq!(
            column_quantiles_many(&[], &qs).unwrap(),
            vec![Vec::<f64>::new(); 3]
        );
        assert_eq!(
            column_quantiles_many_flat(&[], 0, 0, &qs).unwrap(),
            vec![Vec::<f64>::new(); 3]
        );
    }

    #[test]
    fn covariance_with_excess_ddof_remains_finite_and_empty_inputs_stay_empty() {
        let data = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
        // With ddof >= n the documented behavior is to leave the centered
        // cross-product unscaled instead of dividing by zero.
        assert_eq!(
            covariance_matrix(&data, 99),
            vec![vec![2.0, 2.0], vec![2.0, 2.0]]
        );
        assert!(covariance_matrix(&[], 0).is_empty());
        assert!(correlation_matrix(&[]).is_empty());
    }

    // ---- 1-D (single-slice) statistics ----

    #[test]
    fn sum_basic() {
        assert!((sum(&[1.0, 2.0, 3.0, 4.0]) - 10.0).abs() < 1e-12);
        assert!((sum(&[-1.5, 0.5, 1.0]) - 0.0).abs() < 1e-12);
    }

    #[test]
    fn sum_empty_returns_zero() {
        assert_eq!(sum(&[]), 0.0);
    }

    #[test]
    fn mean_basic_1d() {
        assert!((mean(&[1.0, 2.0, 3.0, 4.0]) - 2.5).abs() < 1e-12);
        assert!((mean(&[5.0]) - 5.0).abs() < 1e-12);
    }

    #[test]
    fn mean_empty_returns_nan() {
        assert!(mean(&[]).is_nan());
    }

    #[test]
    fn min_max_basic_1d() {
        assert!((min(&[3.0, -1.0, 2.0]) - (-1.0)).abs() < 1e-12);
        assert!((max(&[3.0, -1.0, 2.0]) - 3.0).abs() < 1e-12);
    }

    #[test]
    fn min_max_empty_returns_identity() {
        assert!(min(&[]).is_infinite() && min(&[]).is_sign_positive());
        assert!(max(&[]).is_infinite() && !max(&[]).is_sign_positive());
    }

    #[test]
    fn variance_ddof_1d() {
        let data = [1.0, 2.0, 3.0, 4.0];
        // population variance = 1.25 ; sample = 5/3
        assert!((variance(&data, 0) - 1.25).abs() < 1e-12);
        assert!((variance(&data, 1) - (5.0 / 3.0)).abs() < 1e-12);
    }

    #[test]
    fn variance_empty_and_bad_ddof_returns_nan() {
        assert!(variance(&[], 0).is_nan());
        // ddof >= n -> non-positive denominator
        assert!(variance(&[1.0, 2.0], 2).is_nan());
        assert!(variance(&[1.0, 2.0], 5).is_nan());
    }

    #[test]
    fn std_matches_variance_sqrt() {
        let data = [1.0, 2.0, 3.0, 4.0];
        assert!((std(&data, 1) - variance(&data, 1).sqrt()).abs() < 1e-12);
        assert!((std(&data, 1) - (5.0_f64 / 3.0).sqrt()).abs() < 1e-12);
    }

    #[test]
    fn median_unsorted_input() {
        // Unsorted input is sorted internally.
        assert!((median(&[3.0, 1.0, 2.0]).unwrap() - 2.0).abs() < 1e-12);
    }

    #[test]
    fn median_even_odd_1d() {
        assert!((median(&[1.0, 2.0, 3.0]).unwrap() - 2.0).abs() < 1e-12);
        assert!((median(&[1.0, 2.0, 3.0, 4.0]).unwrap() - 2.5).abs() < 1e-12);
    }

    #[test]
    fn median_empty_returns_none() {
        assert!(median(&[]).is_none());
    }

    #[test]
    fn mode_basic_1d() {
        assert!((mode(&[1.0, 2.0, 2.0, 3.0]).unwrap() - 2.0).abs() < 1e-12);
    }

    #[test]
    fn mode_tie_smallest_1d() {
        // tie between 1.0 and 2.0 -> smallest wins
        assert!((mode(&[1.0, 2.0, 1.0, 2.0]).unwrap() - 1.0).abs() < 1e-12);
    }

    #[test]
    fn mode_empty_returns_none() {
        assert!(mode(&[]).is_none());
    }

    #[test]
    fn mode_all_equal_returns_that_value() {
        assert!((mode(&[7.0, 7.0, 7.0]).unwrap() - 7.0).abs() < 1e-12);
    }

    fn transpose(data: &[Vec<f64>]) -> Vec<Vec<f64>> {
        if data.is_empty() {
            return vec![];
        }
        let cols = data[0].len();
        (0..cols)
            .map(|j| data.iter().map(|r| r[j]).collect())
            .collect()
    }

    #[test]
    fn column_sum_basic() {
        let data = vec![vec![1.0, 10.0], vec![3.0, 20.0], vec![5.0, 30.0]];
        let s = column_sum(&data);
        assert!((s[0] - 9.0).abs() < 1e-12);
        assert!((s[1] - 60.0).abs() < 1e-12);
    }

    #[test]
    fn covariance_matrix_identity() {
        let data = vec![
            vec![1.0, 0.0],
            vec![0.0, 1.0],
            vec![-1.0, 0.0],
            vec![0.0, -1.0],
        ];
        // col0: [1,0,-1,0] mean=0, var(ddof=1) = (1+0+1+0)/3 = 2/3
        // col1: [0,1,0,-1] mean=0, var(ddof=1) = 2/3
        let cov = covariance_matrix(&data, 1);
        assert!((cov[0][0] - 2.0 / 3.0).abs() < 1e-9);
        assert!((cov[1][1] - 2.0 / 3.0).abs() < 1e-9);
        assert!((cov[0][1]).abs() < 1e-9);
    }

    #[test]
    fn covariance_matrix_hand_computed() {
        let data = vec![vec![1.0, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]];
        // mean = [3, 4]; centered = [[-2,-2],[0,0],[2,2]]
        // cov[0][0] = (4+0+4)/2 = 4; cov[1][1] = (4+0+4)/2 = 4; cov[0][1] = (4+0+4)/2 = 4
        let cov = covariance_matrix(&data, 1);
        assert!((cov[0][0] - 4.0).abs() < 1e-9);
        assert!((cov[1][1] - 4.0).abs() < 1e-9);
        assert!((cov[0][1] - 4.0).abs() < 1e-9);
    }

    #[test]
    fn covariance_population() {
        let data = vec![vec![1.0, 2.0, 3.0, 4.0]];
        let t = transpose(&data);
        let cov = covariance_matrix(&t, 0);
        // population variance of [1,2,3,4] = 1.25
        assert!((cov[0][0] - 1.25).abs() < 1e-12);
    }

    #[test]
    fn correlation_matrix_identity() {
        let data = vec![
            vec![1.0, 0.0],
            vec![0.0, 1.0],
            vec![-1.0, 0.0],
            vec![0.0, -1.0],
        ];
        let corr = correlation_matrix(&data);
        assert!((corr[0][0] - 1.0).abs() < 1e-9);
        assert!((corr[1][1] - 1.0).abs() < 1e-9);
        assert!((corr[0][1]).abs() < 1e-9);
    }

    #[test]
    fn correlation_perfect_positive() {
        // col1 = 2 * col0 => perfect correlation
        let data = vec![vec![1.0, 2.0], vec![2.0, 4.0], vec![3.0, 6.0]];
        let corr = correlation_matrix(&data);
        assert!((corr[0][1] - 1.0).abs() < 1e-9);
        assert!((corr[0][0] - 1.0).abs() < 1e-9);
    }

    #[test]
    fn correlation_constant_column() {
        let data = vec![vec![1.0, 10.0], vec![2.0, 10.0], vec![3.0, 10.0]];
        let corr = correlation_matrix(&data);
        // col1 is constant -> std=0 -> correlation with it is 0
        assert!((corr[0][1]).abs() < 1e-9);
        assert!((corr[1][0]).abs() < 1e-9);
        assert!((corr[1][1] - 1.0).abs() < 1e-9);
    }

    #[test]
    fn correlation_matrix_flat_matches_nested() {
        let rows = vec![
            vec![1.0, 2.0, 3.0],
            vec![2.0, 4.0, 6.0],
            vec![3.0, 6.0, 9.0],
        ];
        let flat: Vec<f64> = rows.iter().flatten().copied().collect();
        let nested = correlation_matrix(&rows);
        let flat_out = correlation_matrix_flat(&flat, 3, 3);
        for i in 0..3 {
            for j in 0..3 {
                assert!((nested[i][j] - flat_out[i][j]).abs() < 1e-12);
            }
        }
        assert!(correlation_matrix_flat(&[], 0, 0).is_empty());
        assert_eq!(correlation_matrix_flat(&flat, 3, 1).len(), 1);
    }

    #[test]
    fn flat_covariance_matches_nested_and_is_symmetric() {
        // The flat scalar covariance (lower-triangle accumulation + mirror)
        // must agree with the nested full-product path and yield an exactly
        // symmetric matrix. Deterministic pseudo-random centered data.
        let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
        let mut next = || {
            state ^= state >> 12;
            state ^= state << 25;
            state ^= state >> 27;
            (state.wrapping_mul(0x2545_F491_4F6C_DD1D) >> 11) as f64 / (1u64 << 53) as f64 * 2.0
                - 1.0
        };
        let (n, p) = (64, 12);
        let flat: Vec<f64> = (0..n * p).map(|_| next()).collect();
        let nested: Vec<Vec<f64>> = flat.chunks_exact(p).map(|r| r.to_vec()).collect();
        let cov_flat = covariance_centered_flat(&flat, n, p, 1);
        let cov_nested = covariance_centered(&nested, 1);
        for a in 0..p {
            for b in 0..p {
                assert!(
                    (cov_flat[a][b] - cov_nested[a][b]).abs() < 1e-12,
                    "mismatch at ({a},{b})"
                );
                assert_eq!(cov_flat[a][b], cov_flat[b][a]);
            }
        }
    }

    #[test]
    fn mode_matches_small_distinct_and_negative_zero() {
        // Continuous mostly-distinct values: all tied at count 1 -> smallest.
        let v = vec![5.0, 1.0, 4.0, 2.0, 3.0];
        assert!((mode(&v).unwrap() - 1.0).abs() < 1e-12);
        // Explicit negative zero is a distinct bit pattern from +0.0.
        let with_nz = vec![-0.0, 0.0, -0.0];
        assert_eq!(mode(&with_nz).unwrap().to_bits(), (-0.0f64).to_bits());
        assert!(mode(&[f64::NAN, f64::NAN]).unwrap().is_nan());
    }
}