//! Logit clipping utilities to prevent softmax saturation
//!
//! # Problem
//! Unbounded logits before softmax can cause:
//! - Numerical overflow: `exp(large_value)` -> inf
//! - Probability saturation: exp(-large_value) -> 0.0
//! - Gradient vanishing: d/dx softmax(saturated) ~ 0.0
//!
//! # Solution
//! Clip logits to [-10, 10] before softmax to ensure:
//! - All probabilities > 1e-6 (no complete saturation)
//! - Gradients remain non-zero
//! - Numerical stability maintained
//!
//! # Mathematical Justification
//! ```text
//! Softmax saturation zones:
//!   exp(-44) ~ 0.0  -> probability ~ 0.0 -> gradient ~ 0.0 (BAD)
//!   exp(+44) -> inf  -> probability ~ 1.0 -> gradient ~ 0.0 (BAD)
//!
//! Safe range with clipping [-10, 10]:
//!   exp(-10) ~ 0.000045 -> still trainable (GOOD)
//!   exp(+10) ~ 22026    -> numerically stable (GOOD)
//!   exp(0) = 1.0        -> baseline reference
//! ```

use ml_core::MLError;

/// Default maximum absolute value for logit clipping
pub const DEFAULT_CLIP_MAX: f32 = 10.0;

/// Clip logits to prevent softmax saturation (in-place on host slice).
///
/// # Arguments
/// * `logits` - Mutable slice of logit values (any length)
/// * `min_val` - Minimum value (e.g., -10.0)
/// * `max_val` - Maximum value (e.g., +10.0)
pub fn clip_logits_inplace(logits: &mut [f32], min_val: f32, max_val: f32) {
    for v in logits.iter_mut() {
        *v = v.clamp(min_val, max_val);
    }
}

/// Clip logits to prevent softmax saturation (returns new Vec).
///
/// # Arguments
/// * `logits` - Slice of raw logit values
/// * `min_val` - Minimum value (e.g., -10.0)
/// * `max_val` - Maximum value (e.g., +10.0)
///
/// # Returns
/// Clipped logits in range [`min_val`, `max_val`]
pub fn clip_logits(logits: &[f32], min_val: f32, max_val: f32) -> Vec<f32> {
    logits.iter().map(|&v| v.clamp(min_val, max_val)).collect()
}

/// Clip logits with default range [-10.0, 10.0]
///
/// Convenience function using the recommended safe range.
///
/// # Arguments
/// * `logits` - Slice of raw logit values
///
/// # Returns
/// Clipped logits in range [-10.0, 10.0]
pub fn clip_logits_default(logits: &[f32]) -> Vec<f32> {
    clip_logits(logits, -DEFAULT_CLIP_MAX, DEFAULT_CLIP_MAX)
}

/// Apply clipping before softmax (convenience function)
///
/// Clips logits and applies softmax in one operation.
///
/// # Arguments
/// * `logits` - Slice of raw logit values
/// * `_dim` - Dimension parameter (ignored for 1D; kept for API compat)
///
/// # Returns
/// Softmax probabilities with clipped logits
pub fn softmax_with_clipping(logits: &[f32], _dim: usize) -> Result<Vec<f32>, MLError> {
    let clipped = clip_logits_default(logits);
    crate::softmax::softmax_with_temperature(&clipped, 1.0)
}

/// Clip a 2D batch of logits (row-major). Each row is clipped independently.
///
/// # Arguments
/// * `logits` - Flat row-major f32 data
/// * `rows` - Number of rows
/// * `cols` - Number of columns per row
/// * `min_val` - Minimum clip value
/// * `max_val` - Maximum clip value
///
/// # Returns
/// Clipped logits (same flat layout)
pub fn clip_logits_batch(
    logits: &[f32],
    _rows: usize,
    _cols: usize,
    min_val: f32,
    max_val: f32,
) -> Vec<f32> {
    clip_logits(logits, min_val, max_val)
}

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

    #[test]
    fn test_clip_logits_basic() -> Result<(), MLError> {
        let logits = [44.0_f32, -44.0_f32, 0.0_f32];

        let clipped = clip_logits(&logits, -10.0, 10.0);

        assert_eq!(clipped[0], 10.0);
        assert_eq!(clipped[1], -10.0);
        assert_eq!(clipped[2], 0.0);

        Ok(())
    }

    #[test]
    fn test_clip_logits_default() -> Result<(), MLError> {
        let logits = [100.0_f32, -100.0_f32];

        let clipped = clip_logits_default(&logits);

        assert_eq!(clipped[0], DEFAULT_CLIP_MAX);
        assert_eq!(clipped[1], -DEFAULT_CLIP_MAX);

        Ok(())
    }

    #[test]
    fn test_softmax_with_clipping() -> Result<(), MLError> {
        // Use less extreme logits that still demonstrate clipping
        // After clipping: [10.0, -10.0, 0.0]
        let logits = [15.0_f32, -15.0_f32, 0.0_f32];

        let probs = softmax_with_clipping(&logits, 0)?;

        // All probabilities should be > 0.0 (no complete saturation)
        for &p in &probs {
            assert!(p > 0.0, "Probability {} should be > 0.0 (not completely saturated)", p);
            assert!(p.is_finite(), "Probability {} should be finite", p);
        }

        // Sum should be 1.0
        let sum: f32 = probs.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5, "Sum should be 1.0, got {}", sum);

        Ok(())
    }

    #[test]
    fn test_batch_clipping() -> Result<(), MLError> {
        let logits = [44.0_f32, -44.0_f32, 20.0_f32, -20.0_f32];

        let clipped = clip_logits_batch(&logits, 2, 2, -10.0, 10.0);

        // Verify all values in range
        for &val in &clipped {
            assert!(
                val >= -10.0 && val <= 10.0,
                "Value {} outside [-10, 10]",
                val
            );
        }

        Ok(())
    }
}
