//! Pure mathematical operations for GGUF inference

//!

//! This module contains standalone math functions used by both CPU and GPU

//! inference paths. By extracting these to a shared module, we enable:

//!

//! - Code reuse between `OwnedQuantizedModel` (CPU) and `OwnedQuantizedModelCuda` (GPU)

//! - Easier testing of mathematical correctness

//! - Clear separation of concerns

//!

//! ## Functions

//!

//! - `rms_norm`: RMSNorm normalization (LLaMA, Qwen, Mistral)

//! - `gelu`: GELU activation function

//! - `silu`: SiLU/Swish activation function

//! - `add_bias`: Add bias vector to output

//! - `argmax`: Find index of maximum value

//! - `softmax`: Numerically stable softmax

use trueno::Vector as TruenoVector;

// =============================================================================

// Normalization Operations

// =============================================================================

/// RMSNorm (Root Mean Square Layer Normalization)

///

/// Used by LLaMA, TinyLlama, Qwen, Mistral instead of LayerNorm.

/// Formula: output = x / sqrt(mean(x^2) + eps) * weight

///

/// # Arguments

/// * `input` - Input tensor [seq_len * hidden_dim]

/// * `weight` - Normalization weights [hidden_dim]

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

///

/// # Returns

/// Normalized output [seq_len * hidden_dim]

pub fn rms_norm(input: &[f32], weight: &[f32], eps: f32) -> Vec<f32> {

    let hidden_dim = weight.len();

    let seq_len = input.len() / hidden_dim;

    let mut output = Vec::with_capacity(input.len());

    let weight_vec = TruenoVector::from_slice(weight);

    for i in 0..seq_len {

        let start = i * hidden_dim;

        let end = start + hidden_dim;

        let x = &input[start..end];

        let x_vec = TruenoVector::from_slice(x);

        // SIMD: sum of squares

        let sum_sq = x_vec

            .sum_of_squares()

            .unwrap_or_else(|_| 
x0
.
iter0
().
map0
(|v| 
v0
 * 
v0
).
sum0
::<f32>());

        let mean_sq = sum_sq / hidden_dim as f32;

        let inv_rms = 1.0 / (mean_sq + eps).sqrt();

        // SIMD: scale by inv_rms, then multiply by weight

        match x_vec

            .scale(inv_rms)

            .and_then(|scaled| scaled.mul(&weight_vec))

        {

            Ok(result) => {

                output.extend_from_slice(result.as_slice());

            }

            Err(_) => {

                // Fallback to scalar

                for j in 0..hidden_dim {

                    output.push(x[j] * inv_rms * weight[j]);

                }

            }

        }

    }

    output

}

/// RMSNorm to pre-allocated buffer (zero-allocation path)

///

/// # Arguments

/// * `input` - Input tensor [hidden_dim] (single position)

/// * `weight` - Normalization weights [hidden_dim]

/// * `eps` - Small constant for numerical stability

/// * `output` - Pre-allocated output buffer [hidden_dim]

pub fn rms_norm_into(input: &[f32], weight: &[f32], eps: f32, output: &mut [f32]) {

    let hidden_dim = weight.len();

    let x = &input[..hidden_dim];

    let x_vec = TruenoVector::from_slice(x);

    let weight_vec = TruenoVector::from_slice(weight);

    let sum_sq = x_vec

        .sum_of_squares()

        .unwrap_or_else(|_| x.iter().map(|v| v * v).sum::<f32>());

    let mean_sq = sum_sq / hidden_dim as f32;

    let inv_rms = 1.0 / (mean_sq + eps).sqrt();

    match x_vec

        .scale(inv_rms)

        .and_then(|scaled| scaled.mul(&weight_vec))

    {

        Ok(result) => {

            output[..hidden_dim].copy_from_slice(result.as_slice());

        }

        Err(_) => {

            for j in 0..hidden_dim {

                output[j] = x[j] * inv_rms * weight[j];

            }

        }

    }

}

/// Layer normalization with optional bias

///

/// PMAT-094: This is actually RMSNorm for LLaMA-style models.

/// Kept for API compatibility with models that expect layer_norm signature.

///

/// # Arguments

/// * `input` - Input tensor [seq_len * hidden_dim]

/// * `weight` - Normalization weights [hidden_dim]

/// * `bias` - Optional bias [hidden_dim]

/// * `eps` - Small constant for numerical stability

pub fn layer_norm(input: &[f32], weight: &[f32], bias: Option<&[f32]>, eps: f32) -> Vec<f32> {

    let hidden_dim = weight.len();

    let seq_len = input.len() / hidden_dim;

    let mut output = Vec::with_capacity(input.len());

    for i in 0..
seq_len1.36k
 {

        let start = i * hidden_dim;

        let end = start + hidden_dim;

        let x = &input[start..end];

        // RMSNorm: compute root mean square (no mean subtraction!)

        let 
sum_sq1.47k
: 
f321.47k
 = 
x1.47k
.
iter1.47k
().
map1.47k
(|v| v * v).
sum1.47k
();

        let rms = (sum_sq / hidden_dim as f32 + eps).sqrt();

        for j in 0..
hidden_dim1.47k
 {

            let normalized = x[j] / rms;

            let mut val = normalized * weight[j];

            if let Some(
b2
) = bias {

                val += b[j];

            }

            output.push(val);

        }

    }

    output

}

/// Layer normalization to pre-allocated buffer

pub fn layer_norm_into(

    input: &[f32],

    weight: &[f32],

    bias: Option<&[f32]>,

    eps: f32,

    output: &mut [f32],

) {

    let hidden_dim = weight.len();

    let x = &input[..hidden_dim];

    let sum_sq: f32 = x.iter().map(|v| v * v).sum();

    let rms = (sum_sq / hidden_dim as f32 + eps).sqrt();

    for j in 0..hidden_dim {

        let normalized = x[j] / rms;

        output[j] = normalized * weight[j];

        if let Some(b) = bias {

            output[j] += b[j];

        }

    }

}

// =============================================================================

// Activation Functions

// =============================================================================

/// GELU (Gaussian Error Linear Unit) activation

///

/// Approximation: GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))

///

/// # Arguments

/// * `input` - Input tensor (modified in-place)

#[inline]

pub fn gelu(input: &mut [f32]) {

    const SQRT_2_OVER_PI: f32 = 0.797_884_6;

    const C: f32 = 0.044_715;

    for x in 
input717
.
iter_mut717
() {

        let inner = SQRT_2_OVER_PI * (*x + C * *x * *x * *x);

        *x = 0.5 * *x * (1.0 + inner.tanh());

    }

}

/// SiLU (Sigmoid Linear Unit) / Swish activation

///

/// SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))

/// Used in SwiGLU FFN (LLaMA, Mistral, etc.)

///

/// # Arguments

/// * `input` - Input tensor (modified in-place)

#[inline]

pub fn silu(input: &mut [f32]) {

    for x in input.iter_mut() {

        *x = *x * (1.0 / (1.0 + (-*x).exp()));

    }

}

// =============================================================================

// Utility Operations

// =============================================================================

/// Add bias vector to output tensor

///

/// # Arguments

/// * `output` - Output tensor [seq_len * out_dim] (modified in-place)

/// * `bias` - Bias vector [out_dim]

#[inline]

pub fn add_bias(output: &mut [f32], bias: &[f32]) {

    let out_dim = bias.len();

    let seq_len = output.len() / out_dim;

    for s in 0..
seq_len1
 {

        for o in 0..
out_dim2
 {

            output[s * out_dim + o] += bias[o];

        }

    }

}

/// Find index of maximum value (greedy decoding)

///

/// # Arguments

/// * `logits` - Logit values [vocab_size]

///

/// # Returns

/// Index of the maximum value

#[inline]

pub fn argmax(logits: &[f32]) -> u32 {

    let mut max_idx = 0u32;

    let mut max_val = f32::NEG_INFINITY;

    for (i, &val) in 
logits398
.
iter398
().
enumerate398
() {

        if val > max_val {

            max_val = val;

            max_idx = i as u32;

        }

    }

    max_idx

}

/// Numerically stable softmax

///

/// Computes softmax(x) = exp(x - max(x)) / sum(exp(x - max(x)))

///

/// # Arguments

/// * `logits` - Input logits (modified in-place to probabilities)

pub fn softmax(logits: &mut [f32]) {

    // Find max for numerical stability

    let max_val = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

    // Compute exp(x - max) and sum

    let mut sum = 0.0f32;

    for x in 
logits2
.
iter_mut2
() {

        *x = (*x - max_val).exp();

        sum += *x;

    }

    // Normalize

    let inv_sum = 1.0 / sum;

    for x in 
logits2
.
iter_mut2
() {

        *x *= inv_sum;

    }

}

// =============================================================================

// Tests

// =============================================================================

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test_gelu_zero() {

        let mut input = vec![0.0];

        gelu(&mut input);

        assert!((input[0] - 0.0).abs() < 1e-6);

    }

    #[test]

    fn test_gelu_positive() {

        let mut input = vec![1.0];

        gelu(&mut input);

        // GELU(1) ≈ 0.8413

        assert!((input[0] - 0.8413).abs() < 0.01);

    }

    #[test]

    fn test_silu_zero() {

        let mut input = vec![0.0];

        silu(&mut input);

        assert!((input[0] - 0.0).abs() < 1e-6);

    }

    #[test]

    fn test_silu_positive() {

        let mut input = vec![1.0];

        silu(&mut input);

        // SiLU(1) = 1 * sigmoid(1) ≈ 0.7311

        assert!((input[0] - 0.7311).abs() < 0.01);

    }

    #[test]

    fn test_argmax() {

        let logits = vec![0.1, 0.5, 0.3, 0.9, 0.2];

        assert_eq!(argmax(&logits), 3);

    }

    #[test]

    fn test_argmax_negative() {

        let logits = vec![-1.0, -0.5, -2.0];

        assert_eq!(argmax(&logits), 1);

    }

    #[test]

    fn test_softmax_uniform() {

        let mut logits = vec![0.0, 0.0, 0.0];

        softmax(&mut logits);

        for &
p3
 in &logits {

            assert!((p - 1.0 / 3.0).abs() < 1e-6);

        }

    }

    #[test]

    fn test_softmax_sums_to_one() {

        let mut logits = vec![1.0, 2.0, 3.0];

        softmax(&mut logits);

        let sum: f32 = logits.iter().sum();

        assert!((sum - 1.0).abs() < 1e-6);

    }

    #[test]

    fn test_add_bias() {

        let mut output = vec![1.0, 2.0, 3.0, 4.0];

        let bias = vec![0.1, 0.2];

        add_bias(&mut output, &bias);

        assert!((output[0] - 1.1).abs() < 1e-6);

        assert!((output[1] - 2.2).abs() < 1e-6);

        assert!((output[2] - 3.1).abs() < 1e-6);

        assert!((output[3] - 4.2).abs() < 1e-6);

    }

    #[test]

    fn test_rms_norm_unit_weight() {

        let input = vec![1.0, 2.0, 3.0, 4.0];

        let weight = vec![1.0, 1.0, 1.0, 1.0];

        let output = rms_norm(&input, &weight, 1e-5);

        // Check output is normalized

        let 
sum_sq1
: 
f321
 = 
output.iter()1
.
map1
(|x| x * x).
sum1
();

        let rms = (sum_sq / 4.0).sqrt();

        assert!((rms - 1.0).abs() < 0.1); // Approximately unit RMS

    }

    #[test]

    fn test_layer_norm_no_bias() {

        let input = vec![1.0, 2.0, 3.0, 4.0];

        let weight = vec![1.0, 1.0, 1.0, 1.0];

        let output = layer_norm(&input, &weight, None, 1e-5);

        assert_eq!(output.len(), 4);

    }

    #[test]

    fn test_layer_norm_with_bias() {

        let input = vec![1.0, 2.0];

        let weight = vec![1.0, 1.0];

        let bias = vec![0.5, 0.5];

        let output = layer_norm(&input, &weight, Some(&bias), 1e-5);

        assert_eq!(output.len(), 2);

        // Output should have bias added

        assert!(output[0] > 0.0);

        assert!(output[1] > 0.0);

    }

}