gbrt_rs/utils/

math.rs

//! Mathematical utilities for statistical and linear algebra operations.
//!
//! This module provides a comprehensive set of mathematical functions and traits for:
//! - Statistical operations (mean, variance, quantiles, etc.)
//! - Vector mathematics (dot products, norms, normalization)
//! - Matrix operations (covariance, correlation, axis-wise statistics)
//!
//! # Design
//!
//! The module uses **trait-based polymorphism** to support multiple data types:
//! - [`Statistics`] trait for statistical measures
//! - [`VectorMath`] trait for vector operations
//! - [`MatrixMath`] trait for matrix operations
//!
//! Implementations are provided for both `Vec<f64>` and `ndarray` types, enabling
//! seamless integration with different data representations.
//!
//! # Error Handling
//!
//! All operations return [`MathResult<T>`], a type alias for `Result<T, MathError>`.
//! This provides robust error handling for edge cases like empty inputs, dimension
//! mismatches, and numerical instability.

use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use std::collections::HashMap;
use thiserror::Error;

/// Errors that can occur during mathematical computations.
///
/// These errors cover invalid inputs, numerical instabilities, dimension mismatches,
/// and other mathematical operation failures.
#[derive(Error, Debug)]
pub enum MathError {
    /// Input values are invalid for the operation.
    #[error("Invalid input: {0}")]
    InvalidInput(String),
    
    /// Numerical computation failed (e.g., overflow, underflow).
    #[error("Numerical error: {0}")]
    NumericalError(String),
    
    /// Array dimensions don't match the expected shape.
    #[error("Dimension mismatch: expected {expected}, got {actual}")]
    DimensionMismatch { expected: String, actual: String },
    
    /// Input data is empty.
    #[error("Empty input data")]
    EmptyInput,
    
    /// Division by zero occurred.
    #[error("Division by zero")]
    DivisionByZero,
}

/// Result type for mathematical operations.
///
/// This is a convenient type alias for `Result<T, MathError>`.
pub type MathResult<T> = std::result::Result<T, MathError>;

/// Statistical operations for numeric collections.
///
/// This trait provides descriptive statistics including central tendency,
/// dispersion, shape measures, and quantile calculations.
pub trait Statistics {
    /// Computes the arithmetic mean (average).
    fn mean(&self) -> MathResult<f64>;
    
    /// Computes sample variance (with Bessel's correction, divides by n-1).
    fn variance(&self) -> MathResult<f64>;
    
    /// Computes sample standard deviation (square root of variance).
    fn std_dev(&self) -> MathResult<f64>;
    
    /// Finds the minimum value.
    fn min(&self) -> MathResult<f64>;
    
    /// Finds the maximum value.
    fn max(&self) -> MathResult<f64>;
    
    /// Computes the q-th quantile (0 ≤ q ≤ 1).
    ///
    /// Uses linear interpolation for non-integer quantile positions.
    fn quantile(&self, q: f64) -> MathResult<f64>;
    
    /// Computes skewness (measure of asymmetry).
    ///
    /// Returns 0 for symmetric distributions, positive for right-skewed,
    /// negative for left-skewed.
    fn skewness(&self) -> MathResult<f64>;
    
    /// Computes excess kurtosis (measure of tail heaviness).
    ///
    /// Normal distributions have kurtosis = 0. Positive values indicate
    /// heavier tails, negative values indicate lighter tails.
    fn kurtosis(&self) -> MathResult<f64>;
}

/// Vector mathematical operations.
///
/// Provides linear algebra operations for 1D numeric arrays.
pub trait VectorMath {
    /// Computes dot product with another vector.
    fn dot(&self, other: &[f64]) -> MathResult<f64>;
    
    /// Computes Lp norm (p > 0).
    ///
    /// - p = 1: Manhattan norm
    /// - p = 2: Euclidean norm (default)
    fn norm(&self, p: f64) -> MathResult<f64>;
    
    /// Normalizes vector to unit length (L2 norm = 1).
    fn normalize(&self) -> MathResult<Vec<f64>>;
    
    /// Standardizes to zero mean and unit variance.
    fn standardize(&self) -> MathResult<Vec<f64>>;
    
    /// Clips values to [min, max] range.
    fn clip(&self, min: f64, max: f64) -> Vec<f64>;
    
    /// Element-wise addition with another vector.
    fn add(&self, other: &[f64]) -> MathResult<Vec<f64>>;
    
    /// Element-wise subtraction with another vector.
    fn subtract(&self, other: &[f64]) -> MathResult<Vec<f64>>;
    
    /// Scalar multiplication.
    fn multiply(&self, scalar: f64) -> Vec<f64>;
}

/// Matrix mathematical operations.
///
/// Provides linear algebra and statistics for 2D arrays.
pub trait MatrixMath {
    /// Computes mean along specified axis.
    ///
    /// - axis = 0: column means
    /// - axis = 1: row means
    fn mean_along_axis(&self, axis: usize) -> MathResult<Array1<f64>>;

    /// Computes variance along specified axis (sample variance).
    fn variance_along_axis(&self, axis: usize) -> MathResult<Array1<f64>>;

    /// Computes covariance matrix.
    fn covariance_matrix(&self) -> MathResult<Array2<f64>>;

    /// Computes correlation matrix.
    fn correlation_matrix(&self) -> MathResult<Array2<f64>>;
}

// Vec<f64> implementation of Statistics trait
impl Statistics for Vec<f64> {
    fn mean(&self) -> MathResult<f64> {
        compute_mean(self)
    }
    
    fn variance(&self) -> MathResult<f64> {
        compute_variance(self)
    }
    
    fn std_dev(&self) -> MathResult<f64> {
        compute_std_dev(self)
    }
    
    fn min(&self) -> MathResult<f64> {
        if self.is_empty() {
            return Err(MathError::EmptyInput);
        }
        Ok(self.iter().fold(f64::INFINITY, |a, &b| a.min(b)))
    }
    
    fn max(&self) -> MathResult<f64> {
        if self.is_empty() {
            return Err(MathError::EmptyInput);
        }
        Ok(self.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)))
    }
    
    fn quantile(&self, q: f64) -> MathResult<f64> {
        if !(0.0..=1.0).contains(&q) {
            return Err(MathError::InvalidInput(
                format!("Quantile must be between 0 and 1, got {}", q)
            ));
        }
        
        if self.is_empty() {
            return Err(MathError::EmptyInput);
        }
        
        let mut sorted = self.clone();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
        
        let index = q * (sorted.len() - 1) as f64;
        let lower_index = index.floor() as usize;
        let upper_index = index.ceil() as usize;
        
        if lower_index == upper_index {
            Ok(sorted[lower_index])
        } else {
            let weight = index - lower_index as f64;
            Ok(sorted[lower_index] * (1.0 - weight) + sorted[upper_index] * weight)
        }
    }
    
    fn skewness(&self) -> MathResult<f64> {
        let mean = self.mean()?;
        let std_dev = self.std_dev()?;
        
        if std_dev == 0.0 {
            return Ok(0.0);
        }
        
        let n = self.len() as f64;
        let sum_cubed_deviations: f64 = self.iter()
            .map(|&x| ((x - mean) / std_dev).powi(3))
            .sum();
        
        Ok(sum_cubed_deviations / n)
    }
    
    fn kurtosis(&self) -> MathResult<f64> {
        let mean = self.mean()?;
        let std_dev = self.std_dev()?;
        
        if std_dev == 0.0 {
            return Ok(0.0);
        }
        
        let n = self.len() as f64;
        let sum_fourth_deviations: f64 = self.iter()
            .map(|&x| ((x - mean) / std_dev).powi(4))
            .sum();
        
        Ok(sum_fourth_deviations / n - 3.0) // excess kurtosis
    }
}

// ndarray::Array1<f64> implementation of Statistics trait
impl Statistics for Array1<f64> {
    fn mean(&self) -> MathResult<f64> {
        compute_mean(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array data".to_string()))?)
    }
    
    fn variance(&self) -> MathResult<f64> {
        compute_variance(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array data".to_string()))?)
    }
    
    fn std_dev(&self) -> MathResult<f64> {
        compute_std_dev(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array data".to_string()))?)
    }
    
    fn min(&self) -> MathResult<f64> {
        if self.is_empty() {
            return Err(MathError::EmptyInput);
        }
        Ok(self.iter().fold(f64::INFINITY, |a, &b| a.min(b)))
    }
    
    fn max(&self) -> MathResult<f64> {
        if self.is_empty() {
            return Err(MathError::EmptyInput);
        }
        Ok(self.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)))
    }
    
    fn quantile(&self, q: f64) -> MathResult<f64> {
        compute_quantile(&self.to_vec(), q)
    }
    
    fn skewness(&self) -> MathResult<f64> {
        compute_skewness(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array data".to_string()))?)
    }
    
    fn kurtosis(&self) -> MathResult<f64> {
        compute_kurtosis(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array data".to_string()))?)
    }
}

// ndarray::ArrayView1<f64> implementation of Statistics trait
impl Statistics for ArrayView1<'_, f64> {
    fn mean(&self) -> MathResult<f64> {
        compute_mean(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array view".to_string()))?)
    }
    
    fn variance(&self) -> MathResult<f64> {
        compute_variance(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array view".to_string()))?)
    }
    
    fn std_dev(&self) -> MathResult<f64> {
        compute_std_dev(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array view".to_string()))?)
    }
    
    fn min(&self) -> MathResult<f64> {
        if self.is_empty() {
            return Err(MathError::EmptyInput);
        }
        Ok(self.iter().fold(f64::INFINITY, |a, &b| a.min(b)))
    }
    
    fn max(&self) -> MathResult<f64> {
        if self.is_empty() {
            return Err(MathError::EmptyInput);
        }
        Ok(self.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)))
    }
    
    fn quantile(&self, q: f64) -> MathResult<f64> {
        compute_quantile(&self.to_vec(), q)
    }
    
    fn skewness(&self) -> MathResult<f64> {
        compute_skewness(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array view".to_string()))?)
    }
    
    fn kurtosis(&self) -> MathResult<f64> {
        compute_kurtosis(self.as_slice()
            .ok_or_else(|| MathError::NumericalError("Non-contiguous array view".to_string()))?)
    }
}

// Helper functions for skewness and kurtosis
fn compute_skewness(values: &[f64]) -> MathResult<f64> {
    let mean = compute_mean(values)?;
    let std_dev = compute_std_dev(values)?;
    
    if std_dev == 0.0 {
        return Ok(0.0);
    }
    
    let n = values.len() as f64;
    let sum_cubed_deviations: f64 = values.iter()
        .map(|&x| ((x - mean) / std_dev).powi(3))
        .sum();
    
    Ok(sum_cubed_deviations / n)
}

fn compute_kurtosis(values: &[f64]) -> MathResult<f64> {
    let mean = compute_mean(values)?;
    let std_dev = compute_std_dev(values)?;
    
    if std_dev == 0.0 {
        return Ok(0.0);
    }
    
    let n = values.len() as f64;
    let sum_fourth_deviations: f64 = values.iter()
        .map(|&x| ((x - mean) / std_dev).powi(4))
        .sum();
    
    Ok(sum_fourth_deviations / n - 3.0)
}

// Helper function for quantile calculation
fn compute_quantile(values: &[f64], q: f64) -> MathResult<f64> {
    if !(0.0..=1.0).contains(&q) {
        return Err(MathError::InvalidInput(
            format!("Quantile must be between 0 and 1, got {}", q)
        ));
    }
    
    if values.is_empty() {
        return Err(MathError::EmptyInput);
    }
    
    let mut sorted = values.to_vec();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
    
    let index = q * (sorted.len() - 1) as f64;
    let lower_index = index.floor() as usize;
    let upper_index = index.ceil() as usize;
    
    if lower_index == upper_index {
        Ok(sorted[lower_index])
    } else {
        let weight = index - lower_index as f64;
        Ok(sorted[lower_index] * (1.0 - weight) + sorted[upper_index] * weight)
    }
}

// Vec<f64> implementation of VectorMath trait
impl VectorMath for Vec<f64> {
    fn dot(&self, other: &[f64]) -> MathResult<f64> {
        if self.len() != other.len() {
            return Err(MathError::DimensionMismatch {
                expected: self.len().to_string(),
                actual: other.len().to_string(),
            });
        }
        
        Ok(self.iter()
            .zip(other.iter())
            .map(|(&a, &b)| a * b)
            .sum())
    }
    
    fn norm(&self, p: f64) -> MathResult<f64> {
        if p <= 0.0 {
            return Err(MathError::InvalidInput(
                format!("p must be positive, got {}", p)
            ));
        }
        
        if p == 1.0 {
            Ok(l1_norm(self))
        } else if p == 2.0 {
            Ok(l2_norm(self))
        } else {
            let sum: f64 = self.iter()
                .map(|&x| x.abs().powf(p))
                .sum();
            Ok(sum.powf(1.0 / p))
        }
    }
    
    fn normalize(&self) -> MathResult<Vec<f64>> {
        normalize_vector(self)
    }
    
    fn standardize(&self) -> MathResult<Vec<f64>> {
        standardize_vector(self)
    }
    
    fn clip(&self, min: f64, max: f64) -> Vec<f64> {
        clip_values(self, min, max)
    }
    
    fn add(&self, other: &[f64]) -> MathResult<Vec<f64>> {
        if self.len() != other.len() {
            return Err(MathError::DimensionMismatch {
                expected: self.len().to_string(),
                actual: other.len().to_string(),
            });
        }
        
        Ok(self.iter()
            .zip(other.iter())
            .map(|(&a, &b)| a + b)
            .collect())
    }
    
    fn subtract(&self, other: &[f64]) -> MathResult<Vec<f64>> {
        if self.len() != other.len() {
            return Err(MathError::DimensionMismatch {
                expected: self.len().to_string(),
                actual: other.len().to_string(),
            });
        }
        
        Ok(self.iter()
            .zip(other.iter())
            .map(|(&a, &b)| a - b)
            .collect())
    }
    
    fn multiply(&self, scalar: f64) -> Vec<f64> {
        self.iter().map(|&x| x * scalar).collect()
    }
}

// ndarray::Array2<f64> implementation of MatrixMath trait
impl MatrixMath for Array2<f64> {
    fn mean_along_axis(&self, axis: usize) -> MathResult<Array1<f64>> {
        if axis > 1 {
            return Err(MathError::InvalidInput(
                format!("Axis must be 0 or 1, got {}", axis)
            ));
        }
        
        match axis {
            0 => {
                // Mean of each column
                let n_rows = self.nrows() as f64;
                Ok(self.mean_axis(ndarray::Axis(0)).unwrap() / n_rows)
            }
            1 => {
                // Mean of each row
                let n_cols = self.ncols() as f64;
                Ok(self.mean_axis(ndarray::Axis(1)).unwrap() / n_cols)
            }
            _ => unreachable!(),
        }
    }
    
    fn variance_along_axis(&self, axis: usize) -> MathResult<Array1<f64>> {
        if axis > 1 {
            return Err(MathError::InvalidInput(
                format!("Axis must be 0 or 1, got {}", axis)
            ));
        }
        
        let means = self.mean_along_axis(axis)?;
        let n = if axis == 0 { self.nrows() } else { self.ncols() } as f64;
        
        match axis {
            0 => {
                // Variance of each column
                let variances: Vec<f64> = (0..self.ncols())
                    .map(|col| {
                        let col_data = self.column(col);
                        let mean = means[col];
                        col_data.iter()
                            .map(|&x| (x - mean).powi(2))
                            .sum::<f64>() / (n - 1.0)
                    })
                    .collect();
                Ok(Array1::from(variances))
            }
            1 => {
                // Variance of each row
                let variances: Vec<f64> = (0..self.nrows())
                    .map(|row| {
                        let row_data = self.row(row);
                        let mean = means[row];
                        row_data.iter()
                            .map(|&x| (x - mean).powi(2))
                            .sum::<f64>() / (n - 1.0)
                    })
                    .collect();
                Ok(Array1::from(variances))
            }
            _ => unreachable!(),
        }
    }
    
    fn covariance_matrix(&self) -> MathResult<Array2<f64>> {
        let n_samples = self.nrows() as f64;
        if n_samples <= 1.0 {
            return Err(MathError::InvalidInput(
                "Need at least 2 samples for covariance".to_string()
            ));
        }
        
        let means = self.mean_along_axis(0)?;
        let mut covariance = Array2::zeros((self.ncols(), self.ncols()));
        
        for i in 0..self.ncols() {
            for j in 0..self.ncols() {
                let cov: f64 = (0..self.nrows())
                    .map(|k| (self[[k, i]] - means[i]) * (self[[k, j]] - means[j]))
                    .sum::<f64>() / (n_samples - 1.0);
                covariance[[i, j]] = cov;
            }
        }
        
        Ok(covariance)
    }
    
    fn correlation_matrix(&self) -> MathResult<Array2<f64>> {
        let covariance = self.covariance_matrix()?;
        let std_devs = self.variance_along_axis(0)?.mapv(|v| v.sqrt());
        
        let mut correlation = Array2::zeros(covariance.raw_dim());
        
        for i in 0..covariance.nrows() {
            for j in 0..covariance.ncols() {
                if std_devs[i] == 0.0 || std_devs[j] == 0.0 {
                    correlation[[i, j]] = 0.0;
                } else {
                    correlation[[i, j]] = covariance[[i, j]] / (std_devs[i] * std_devs[j]);
                }
            }
        }
        
        Ok(correlation)
    }
}

// Basic statistical functions

/// Computes the arithmetic mean of a slice.
///
/// # Arguments
///
/// * `values` - Slice of numeric values
///
/// # Returns
///
/// The average value
///
/// # Errors
///
/// Returns [`MathError::EmptyInput`] if `values` is empty
pub fn compute_mean(values: &[f64]) -> MathResult<f64> {
    if values.is_empty() {
        return Err(MathError::EmptyInput);
    }
    
    let sum: f64 = values.iter().sum();
    Ok(sum / values.len() as f64)
}

/// Computes sample variance (with Bessel's correction).
///
/// Divides by `n-1` rather than `n` to provide an unbiased estimator.
///
/// # Arguments
///
/// * `values` - Slice of numeric values
///
/// # Returns
///
/// Sample variance
///
/// # Errors
///
/// Returns [`MathError::InvalidInput`] if fewer than 2 values are provided
pub fn compute_variance(values: &[f64]) -> MathResult<f64> {
    if values.len() < 2 {
        return Err(MathError::InvalidInput(
            "Need at least 2 values for variance".to_string()
        ));
    }
    
    let mean = compute_mean(values)?;
    let sum_sq_diff: f64 = values.iter()
        .map(|&x| (x - mean).powi(2))
        .sum();
    
    Ok(sum_sq_diff / (values.len() - 1) as f64)
}

/// Computes sample standard deviation.
///
/// Returns the square root of the sample variance.
///
/// # Arguments
///
/// * `values` - Slice of numeric values
///
/// # Returns
///
/// Sample standard deviation
pub fn compute_std_dev(values: &[f64]) -> MathResult<f64> {
    let variance = compute_variance(values)?;
    Ok(variance.sqrt())
}

/// Applies softmax function to convert logits to probabilities.
///
/// For numerical stability, subtracts the maximum value before exponentiation.
///
/// # Arguments
///
/// * `values` - Slice of logits (unnormalized log probabilities)
///
/// # Returns
///
/// Vector of probabilities that sum to 1.0
///
/// # Errors
///
/// Returns [`MathError::EmptyInput`] if `values` is empty
pub fn softmax(values: &[f64]) -> MathResult<Vec<f64>> {
    if values.is_empty() {
        return Err(MathError::EmptyInput);
    }
    
    // For numerical stability, subtract max value
    let max_val = values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
    let exps: Vec<f64> = values.iter()
        .map(|&x| (x - max_val).exp())
        .collect();
    
    let sum_exps: f64 = exps.iter().sum();
    
    Ok(exps.iter().map(|&x| x / sum_exps).collect())
}

/// Applies sigmoid (logistic) function.
///
/// Maps any real number to (0, 1).
///
/// # Arguments
///
/// * `x` - Input value
///
/// # Returns
///
/// Sigmoid(x) = 1 / (1 + e^(-x))
pub fn sigmoid(x: f64) -> f64 {
    1.0 / (1.0 + (-x).exp())
}

/// Computes logit (inverse of sigmoid).
///
/// # Arguments
///
/// * `p` - Probability in (0, 1)
///
/// # Returns
///
/// logit(p) = ln(p / (1-p))
///
/// # Errors
///
/// Returns [`MathError::InvalidInput`] if `p` is not in (0, 1)
pub fn logit(p: f64) -> MathResult<f64> {
    if p <= 0.0 || p >= 1.0 {
        return Err(MathError::InvalidInput(
            format!("Probability must be in (0, 1), got {}", p)
        ));
    }
    
    Ok((p / (1.0 - p)).ln())
}

/// Computes L1 norm (Manhattan distance).
///
/// Sum of absolute values.
///
/// # Arguments
///
/// * `values` - Slice of numeric values
///
/// # Returns
///
/// L1 norm
pub fn l1_norm(values: &[f64]) -> f64 {
    values.iter().map(|&x| x.abs()).sum()
}

/// Computes L2 norm (Euclidean norm).
///
/// Square root of sum of squares.
///
/// # Arguments
///
/// * `values` - Slice of numeric values
///
/// # Returns
///
/// L2 norm
pub fn l2_norm(values: &[f64]) -> f64 {
    values.iter().map(|&x| x.powi(2)).sum::<f64>().sqrt()
}

/// Computes Euclidean distance between two vectors.
///
/// # Arguments
///
/// * `a` - First vector
/// * `b` - Second vector
///
/// # Returns
///
/// Euclidean distance
///
/// # Errors
///
/// Returns [`MathError::DimensionMismatch`] if vectors have different lengths
pub fn euclidean_distance(a: &[f64], b: &[f64]) -> MathResult<f64> {
    if a.len() != b.len() {
        return Err(MathError::DimensionMismatch {
            expected: a.len().to_string(),
            actual: b.len().to_string(),
        });
    }
    
    let sum_sq: f64 = a.iter()
        .zip(b.iter())
        .map(|(&x, &y)| (x - y).powi(2))
        .sum();
    
    Ok(sum_sq.sqrt())
}

/// Clips values to a specified range.
///
/// Values below `min` become `min`, values above `max` become `max`.
///
/// # Arguments
///
/// * `values` - Input values
/// * `min` - Minimum bound
/// * `max` - Maximum bound
///
/// # Returns
///
/// Clipped values
pub fn clip_values(values: &[f64], min: f64, max: f64) -> Vec<f64> {
    values.iter()
        .map(|&x| x.max(min).min(max))
        .collect()
}

/// Normalizes vector to unit length (L2 norm = 1).
///
/// # Arguments
///
/// * `values` - Input vector
///
/// # Returns
///
/// Normalized vector
///
/// # Errors
///
/// Returns [`MathError::NumericalError`] if vector has zero norm
pub fn normalize_vector(values: &[f64]) -> MathResult<Vec<f64>> {
    let norm = l2_norm(values);
    if norm == 0.0 {
        return Err(MathError::NumericalError("Cannot normalize zero vector".to_string()));
    }
    
    Ok(values.iter().map(|&x| x / norm).collect())
}

/// Standardizes vector to zero mean and unit variance.
///
/// # Arguments
///
/// * `values` - Input vector
///
/// # Returns
///
/// Standardized vector: (x - mean) / std_dev
///
/// # Errors
///
/// Returns [`MathError::NumericalError`] if standard deviation is zero
pub fn standardize_vector(values: &[f64]) -> MathResult<Vec<f64>> {
    let mean = compute_mean(values)?;
    let std_dev = compute_std_dev(values)?;
    
    if std_dev == 0.0 {
        return Err(MathError::NumericalError("Cannot standardize constant vector".to_string()));
    }
    
    Ok(values.iter().map(|&x| (x - mean) / std_dev).collect())
}

/// Computes weighted mean.
///
/// # Arguments
///
/// * `values` - Values to average
/// * `weights` - Weights for each value
///
/// # Returns
///
/// Weighted mean: Σ(values × weights) / Σ(weights)
///
/// # Errors
///
/// Returns [`MathError::DimensionMismatch`] if lengths differ,
/// [`MathError::EmptyInput`] if empty, or [`MathError::DivisionByZero`] if weights sum to zero
pub fn weighted_mean(values: &[f64], weights: &[f64]) -> MathResult<f64> {
    if values.len() != weights.len() {
        return Err(MathError::DimensionMismatch {
            expected: values.len().to_string(),
            actual: weights.len().to_string(),
        });
    }
    
    if values.is_empty() {
        return Err(MathError::EmptyInput);
    }
    
    let weighted_sum: f64 = values.iter()
        .zip(weights.iter())
        .map(|(&x, &w)| x * w)
        .sum();
    
    let weight_sum: f64 = weights.iter().sum();
    
    if weight_sum == 0.0 {
        return Err(MathError::DivisionByZero);
    }
    
    Ok(weighted_sum / weight_sum)
}

/// Computes Root Mean Square Error (RMSE).
///
/// # Arguments
///
/// * `predictions` - Model predictions
/// * `targets` - True target values
///
/// # Returns
///
/// RMSE = √[Σ(pred - target)² / n]
///
/// # Errors
///
/// Returns [`MathError::DimensionMismatch`] if lengths differ,
/// or [`MathError::EmptyInput`] if empty
pub fn rmse(predictions: &[f64], targets: &[f64]) -> MathResult<f64> {
    if predictions.len() != targets.len() {
        return Err(MathError::DimensionMismatch {
            expected: predictions.len().to_string(),
            actual: targets.len().to_string(),
        });
    }
    
    if predictions.is_empty() {
        return Err(MathError::EmptyInput);
    }
    
    let mse: f64 = predictions.iter()
        .zip(targets.iter())
        .map(|(&pred, &target)| (pred - target).powi(2))
        .sum::<f64>() / predictions.len() as f64;
    
    Ok(mse.sqrt())
}