//! Normalization operations for Vector<f32>

//!

//! This module provides normalization methods:

//! - `zscore()` - Z-score normalization (standardization)

//! - `minmax_normalize()` - Min-max normalization to [0, 1]

//! - `layer_norm()` - Layer normalization with learnable parameters

//! - `layer_norm_simple()` - Layer normalization without learnable parameters

//! - `normalize()` - Normalize to unit length (L2 norm = 1)

use crate::{Result, TruenoError, Vector};

impl Vector<f32> {

    /// Z-score normalization (standardization)

    ///

    /// Transforms the vector to have mean = 0 and standard deviation = 1.

    /// Each element is transformed as: z\[i\] = (x\[i\] - μ) / σ

    ///

    /// This is a fundamental preprocessing step in machine learning and statistics,

    /// ensuring features have comparable scales and are centered around zero.

    ///

    /// # Performance

    ///

    /// Uses optimized SIMD implementations via mean() and stddev(), then applies

    /// element-wise operations (sub, scale) which also use SIMD.

    ///

    /// # Examples

    ///

    /// ```

    /// use trueno::Vector;

    ///

    /// let v = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0]);

    /// let z = v.zscore().unwrap();

    ///

    /// // Verify mean ≈ 0

    /// let mean = z.mean().unwrap();

    /// assert!(mean.abs() < 1e-5);

    ///

    /// // Verify stddev ≈ 1

    /// let std = z.stddev().unwrap();

    /// assert!((std - 1.0).abs() < 1e-5);

    /// ```

    ///

    /// # Empty vectors

    ///

    /// Returns EmptyVector error for empty vectors (cannot compute mean/stddev).

    ///

    /// # Division by zero

    ///

    /// Returns DivisionByZero error if the vector has zero standard deviation

    /// (i.e., all elements are identical/constant).

    ///

    /// ```

    /// use trueno::{Vector, TruenoError};

    ///

    /// let v = Vector::from_slice(&[5.0, 5.0, 5.0]); // Constant

    /// assert!(matches!(v.zscore(), Err(TruenoError::DivisionByZero)));

    /// ```

    pub fn zscore(&self) -> Result<Self> {

        if self.as_slice().is_empty() {

            return Err(TruenoError::EmptyVector);

        }

        let mean_val = self.mean()?;

        let std_val = self.stddev()?;

        // Check for zero standard deviation (constant vector)

        if std_val.abs() < 1e-10 {

            return Err(TruenoError::DivisionByZero);

        }

        // Transform: z[i] = (x[i] - μ) / σ

        let inv_std = 1.0 / std_val;

        let data: Vec<f32> = self

            .as_slice()

            .iter()

            .map(|&x| (x - mean_val) * inv_std)

            .collect();

        Ok(Vector::from_vec(data))

    }

    /// Min-max normalization (scaling to [0, 1] range)

    ///

    /// Transforms the vector so that the minimum value becomes 0 and the maximum

    /// value becomes 1, with all other values scaled proportionally.

    /// Formula: x'\[i\] = (x\[i\] - min) / (max - min)

    ///

    /// This is a fundamental preprocessing technique in machine learning, especially

    /// for algorithms sensitive to feature magnitudes (e.g., neural networks, k-NN).

    ///

    /// # Performance

    ///

    /// Uses optimized SIMD implementations via min() and max() operations, then

    /// applies element-wise transformation.

    ///

    /// # Examples

    ///

    /// ```

    /// use trueno::Vector;

    ///

    /// let v = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0]);

    /// let normalized = v.minmax_normalize().unwrap();

    ///

    /// // Verify range [0, 1]

    /// let min = normalized.min().unwrap();

    /// let max = normalized.max().unwrap();

    /// assert!((min - 0.0).abs() < 1e-5);

    /// assert!((max - 1.0).abs() < 1e-5);

    /// ```

    ///

    /// # Empty vectors

    ///

    /// Returns EmptyVector error for empty vectors (cannot compute min/max).

    ///

    /// # Division by zero

    ///

    /// Returns DivisionByZero error if the vector has all identical elements

    /// (i.e., min = max, causing division by zero in the normalization formula).

    ///

    /// ```

    /// use trueno::{Vector, TruenoError};

    ///

    /// let v = Vector::from_slice(&[5.0, 5.0, 5.0]); // Constant

    /// assert!(matches!(v.minmax_normalize(), Err(TruenoError::DivisionByZero)));

    /// ```

    pub fn minmax_normalize(&self) -> Result<Self> {

        if self.as_slice().is_empty() {

            return Err(TruenoError::EmptyVector);

        }

        let min_val = self.min()?;

        let max_val = self.max()?;

        let range = max_val - min_val;

        // Check for zero range (constant vector)

        if range.abs() < 1e-10 {

            return Err(TruenoError::DivisionByZero);

        }

        // Transform: x'[i] = (x[i] - min) / (max - min)

        let inv_range = 1.0 / range;

        let data: Vec<f32> = self

            .as_slice()

            .iter()

            .map(|&x| (x - min_val) * inv_range)

            .collect();

        Ok(Vector::from_vec(data))

    }

    /// Layer normalization with learnable parameters (Issue #61: ML primitives)

    ///

    /// Applies layer normalization: `y = gamma * (x - mean) / sqrt(variance + eps) + beta`

    ///

    /// This is a fundamental normalization technique in transformers and other

    /// modern neural network architectures. Unlike batch normalization, layer norm

    /// normalizes across the feature dimension, making it suitable for sequence models.

    ///

    /// # Arguments

    ///

    /// * `gamma` - Scale parameter (typically learned, initialized to 1.0)

    /// * `beta` - Shift parameter (typically learned, initialized to 0.0)

    /// * `eps` - Small constant for numerical stability (typically 1e-5 or 1e-6)

    ///

    /// # Returns

    ///

    /// Normalized vector with the same shape as input

    ///

    /// # Errors

    ///

    /// Returns `SizeMismatch` if gamma or beta have different lengths than self

    /// Returns `EmptyVector` if input is empty

    ///

    /// # Example

    ///

    /// ```

    /// use trueno::Vector;

    ///

    /// let x = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0]);

    /// let gamma = Vector::from_slice(&[1.0, 1.0, 1.0, 1.0]); // Scale = 1

    /// let beta = Vector::from_slice(&[0.0, 0.0, 0.0, 0.0]);  // Shift = 0

    ///

    /// let y = x.layer_norm(&gamma, &beta, 1e-5).unwrap();

    ///

    /// // Output should be approximately standardized (mean ≈ 0, std ≈ 1)

    /// let mean: f32 = y.as_slice().iter().sum::<f32>() / y.len() as f32;

    /// assert!(mean.abs() < 1e-5);

    /// ```

    ///

    /// # Performance

    ///

    /// Single-pass computation using Welford's algorithm for numerical stability.

    /// Time complexity: O(n), Space complexity: O(n).

    pub fn layer_norm(&self, gamma: &Self, beta: &Self, eps: f32) -> Result<Self> {

        if self.as_slice().is_empty() {

            return Err(TruenoError::EmptyVector);

        }

        if self.len() != gamma.len() {

            return Err(TruenoError::SizeMismatch {

                expected: self.len(),

                actual: gamma.len(),

            });

        }

        if self.len() != beta.len() {

            return Err(TruenoError::SizeMismatch {

                expected: self.len(),

                actual: beta.len(),

            });

        }

        // Compute mean

        let mean_val = self.mean()?;

        // Compute variance: E[(x - mean)^2]

        let variance: f32 = self

            .as_slice()

            .iter()

            .map(|&x| {

                let diff = x - mean_val;

                diff * diff

            })

            .sum::<f32>()

            / self.len() as f32;

        // Compute inverse standard deviation for numerical stability

        let inv_std = 1.0 / (variance + eps).sqrt();

        // Apply normalization: y = gamma * (x - mean) * inv_std + beta

        let data: Vec<f32> = self

            .as_slice()

            .iter()

            .zip(gamma.as_slice().iter())

            .zip(beta.as_slice().iter())

            .map(|((&x, &g), &b)| g * (x - mean_val) * inv_std + b)

            .collect();

        Ok(Vector::from_vec(data))

    }

    /// Layer normalization without learnable parameters

    ///

    /// Simplified version that just standardizes the input: `y = (x - mean) / sqrt(variance + eps)`

    ///

    /// This is equivalent to calling `layer_norm` with gamma=1 and beta=0.

    ///

    /// # Arguments

    ///

    /// * `eps` - Small constant for numerical stability (typically 1e-5)

    ///

    /// # Example

    ///

    /// ```

    /// use trueno::Vector;

    ///

    /// let x = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0]);

    /// let y = x.layer_norm_simple(1e-5).unwrap();

    ///

    /// // Output should be standardized

    /// let mean: f32 = y.as_slice().iter().sum::<f32>() / y.len() as f32;

    /// assert!(mean.abs() < 1e-5);

    /// ```

    pub fn layer_norm_simple(&self, eps: f32) -> Result<Self> {

        if self.as_slice().is_empty() {

            return Err(TruenoError::EmptyVector);

        }

        let mean_val = self.mean()?;

        // Compute variance

        let variance: f32 = self

            .as_slice()

            .iter()

            .map(|&x| {

                let diff = x - mean_val;

                diff * diff

            })

            .sum::<f32>()

            / self.len() as f32;

        let inv_std = 1.0 / (variance + eps).sqrt();

        let data: Vec<f32> = self

            .as_slice()

            .iter()

            .map(|&x| (x - mean_val) * inv_std)

            .collect();

        Ok(Vector::from_vec(data))

    }

    /// Normalize the vector to unit length (L2 norm = 1)

    ///

    /// Returns a new vector in the same direction but with magnitude 1.

    ///

    /// # Errors

    ///

    /// Returns `TruenoError::DivisionByZero` if the vector has zero norm (cannot normalize zero vector).

    ///

    /// # Examples

    ///

    /// ```

    /// use trueno::Vector;

    ///

    /// let v = Vector::from_slice(&[3.0, 4.0]);

    /// let unit = v.normalize().unwrap();

    ///

    /// // Result is [0.6, 0.8] (a unit vector)

    /// assert!((unit.as_slice()[0] - 0.6).abs() < 1e-5);

    /// assert!((unit.as_slice()[1] - 0.8).abs() < 1e-5);

    ///

    /// // Verify it's a unit vector (norm = 1)

    /// assert!((unit.norm_l2().unwrap() - 1.0).abs() < 1e-5);

    /// ```

    ///

    /// # Zero Vector Error

    ///

    /// ```

    /// use trueno::{Vector, TruenoError};

    ///

    /// let v = Vector::from_slice(&[0.0, 0.0]);

    /// assert!(matches!(v.normalize(), Err(TruenoError::DivisionByZero)));

    /// ```

    pub fn normalize(&self) -> Result<Vector<f32>> {

        let norm = self.norm_l2()?;

        // Check for zero or near-zero norm (cannot normalize zero vector)

        if norm.abs() < 1e-10 {

            return Err(TruenoError::DivisionByZero);

        }

        // Divide each element by the norm using scalar multiplication

        // This avoids creating an intermediate vector

        self.scale(1.0 / norm)

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test_zscore_basic() {

        let v = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0]);

        let z = v.zscore().unwrap();

        // Mean should be ~0

        let mean = z.mean().unwrap();

        assert!(mean.abs() < 1e-5);

        // Stddev should be ~1

        let std = z.stddev().unwrap();

        assert!((std - 1.0).abs() < 1e-5);

    }

    #[test]

    fn test_zscore_empty() {

        let v: Vector<f32> = Vector::from_slice(&[]);

        assert!(matches!(v.zscore(), Err(TruenoError::EmptyVector)));

    }

    #[test]

    fn test_zscore_constant() {

        let v = Vector::from_slice(&[5.0, 5.0, 5.0]);

        assert!(matches!(v.zscore(), Err(TruenoError::DivisionByZero)));

    }

    #[test]

    fn test_minmax_normalize_basic() {

        let v = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0]);

        let normalized = v.minmax_normalize().unwrap();

        assert!((normalized.min().unwrap() - 0.0).abs() < 1e-5);

        assert!((normalized.max().unwrap() - 1.0).abs() < 1e-5);

    }

    #[test]

    fn test_minmax_normalize_empty() {

        let v: Vector<f32> = Vector::from_slice(&[]);

        assert!(matches!(v.minmax_normalize(), Err(TruenoError::EmptyVector)));

    }

    #[test]

    fn test_minmax_normalize_constant() {

        let v = Vector::from_slice(&[5.0, 5.0, 5.0]);

        assert!(matches!(

            v.minmax_normalize(),

            Err(TruenoError::DivisionByZero)

        ));

    }

    #[test]

    fn test_layer_norm() {

        let x = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0]);

        let gamma = Vector::from_slice(&[1.0, 1.0, 1.0, 1.0]);

        let beta = Vector::from_slice(&[0.0, 0.0, 0.0, 0.0]);

        let y = x.layer_norm(&gamma, &beta, 1e-5).unwrap();

        // Mean should be ~0

        let mean: f32 = y.as_slice().iter().sum::<f32>() / y.len() as f32;

        assert!(mean.abs() < 1e-5);

    }

    #[test]

    fn test_layer_norm_size_mismatch() {

        let x = Vector::from_slice(&[1.0, 2.0, 3.0]);

        let gamma = Vector::from_slice(&[1.0, 1.0]); // Wrong size

        let beta = Vector::from_slice(&[0.0, 0.0, 0.0]);

        assert!(matches!(

            x.layer_norm(&gamma, &beta, 1e-5),

            Err(TruenoError::SizeMismatch { .. })

        ));

    }

    #[test]

    fn test_layer_norm_simple() {

        let x = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0]);

        let y = x.layer_norm_simple(1e-5).unwrap();

        let mean: f32 = y.as_slice().iter().sum::<f32>() / y.len() as f32;

        assert!(mean.abs() < 1e-5);

    }

    #[test]

    fn test_normalize_unit_vector() {

        let v = Vector::from_slice(&[3.0, 4.0]);

        let unit = v.normalize().unwrap();

        assert!((unit.as_slice()[0] - 0.6).abs() < 1e-5);

        assert!((unit.as_slice()[1] - 0.8).abs() < 1e-5);

        assert!((unit.norm_l2().unwrap() - 1.0).abs() < 1e-5);

    }

    #[test]

    fn test_normalize_zero_vector() {

        let v = Vector::from_slice(&[0.0, 0.0]);

        assert!(matches!(v.normalize(), Err(TruenoError::DivisionByZero)));

    }

}