//! Gradient utilities for Candle framework
//!
//! Provides gradient clipping, NaN/Inf detection, and other gradient-related
//! operations that are missing from candle_nn::optim

use candle_core::{backprop::GradStore, Error, Var};

use crate::MLError;

/// Clip gradients by global L2 norm (similar to PyTorch's clip_grad_norm_)
///
/// This function computes the L2 norm of all gradients and scales them
/// proportionally if the total norm exceeds max_norm.
///
/// # Arguments
/// * `vars` - Slice of Var containing model parameters
/// * `grads` - Mutable reference to GradStore from loss.backward()
/// * `max_norm` - Maximum allowed gradient norm
///
/// # Returns
/// A tuple of (actual_norm, clipped_norm) where:
/// - actual_norm: The total gradient norm before clipping
/// - clipped_norm: The gradient norm after clipping (≤ max_norm)
///
/// # Example
/// ```no_run
/// use ml::gradient_utils::clip_grad_norm;
/// use candle_core::Var;
///
/// let vars = vec![]; // Your model variables
/// // ... compute loss ...
/// let mut grads = loss.backward()?;
/// let (actual_norm, clipped_norm) = clip_grad_norm(&vars, &mut grads, 10.0)?;
/// println!("Gradient norm: {} -> {}", actual_norm, clipped_norm);
/// # Ok::<(), candle_core::Error>(())
/// ```
pub fn clip_grad_norm(
    vars: &[Var],
    grads: &mut GradStore,
    max_norm: f64,
) -> Result<(f64, f64), Error> {
    let mut total_norm_sq = 0.0_f64;

    // First pass: Calculate the total L2 norm of all gradients
    for var in vars.iter() {
        if let Some(grad) = grads.get(var) {
            let norm_sq = grad.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
            total_norm_sq += norm_sq;
        }
    }
    let total_norm = total_norm_sq.sqrt();

    // Second pass: Scale gradients if the norm exceeds the maximum
    if total_norm > max_norm {
        let scale_factor = max_norm / (total_norm + 1e-6); // Add epsilon for numerical stability
        for var in vars.iter() {
            // Remove gradient, scale it, and insert back
            if let Some(grad) = grads.remove(var) {
                let scaled_grad = (grad * scale_factor)?;
                grads.insert(var, scaled_grad);
            }
        }
        // Return both actual and clipped norms
        Ok((total_norm, max_norm))
    } else {
        // Norm is already below max_norm, no clipping needed
        Ok((total_norm, total_norm))
    }
}

/// Check if any gradient in the GradStore contains NaN or Inf values.
///
/// Returns `Ok(())` if all gradients are finite, or an error naming
/// the first parameter with non-finite values. Uses `sum_all` as an
/// efficient check: if any element is NaN/Inf, the sum will be non-finite.
///
/// # Arguments
/// * `vars` - Slice of Var containing model parameters
/// * `grads` - Reference to GradStore from loss.backward()
///
/// # Returns
/// `Ok(())` if all gradients are finite, `Err(MLError::TrainingError)` otherwise
pub fn check_gradients_finite(vars: &[Var], grads: &GradStore) -> Result<(), MLError> {
    for (idx, var) in vars.iter().enumerate() {
        if let Some(grad) = grads.get(var) {
            // Sum all elements -- if any element is NaN, the sum will be NaN
            let sum = grad
                .sum_all()
                .and_then(|t| t.to_scalar::<f32>())
                .map_err(|e| {
                    MLError::TrainingError(format!(
                        "Failed to check gradient for param {}: {}",
                        idx, e
                    ))
                })?;

            if !sum.is_finite() {
                return Err(MLError::TrainingError(format!(
                    "NaN/Inf gradient detected in parameter index {}. \
                     Training is numerically unstable -- halting to prevent model corruption. \
                     Consider reducing learning rate or checking input data for anomalies.",
                    idx
                )));
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::{Device, Tensor};

    #[test]
    fn test_finite_gradients_pass() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &Tensor::new(&[1.0_f32, 2.0, 3.0], &device).expect("failed to create tensor"),
        )
        .expect("failed to create var");

        let loss = var
            .mul(&var)
            .expect("mul failed")
            .sum_all()
            .expect("sum failed");
        let grads = loss.backward().expect("backward failed");

        let result = check_gradients_finite(&[var], &grads);
        assert!(result.is_ok(), "Finite gradients should pass check");
    }

    #[test]
    fn test_nan_gradient_detected() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &Tensor::new(&[0.0_f32], &device).expect("failed to create tensor"),
        )
        .expect("failed to create var");

        // 0/0 produces NaN
        let zero = Tensor::new(&[0.0_f32], &device).expect("failed to create zero");
        let nan_result = var.div(&zero).expect("div failed");
        let loss = nan_result.sum_all().expect("sum failed");
        let grads = loss.backward().expect("backward failed");

        let result = check_gradients_finite(&[var], &grads);
        assert!(result.is_err(), "NaN/Inf gradients should be detected");
        let err_msg = format!("{}", result.expect_err("expected error"));
        assert!(
            err_msg.contains("NaN/Inf gradient detected"),
            "Error should mention NaN/Inf: {}",
            err_msg
        );
    }

    #[test]
    fn test_empty_vars_passes() {
        let device = Device::Cpu;
        // Create an empty GradStore by computing backward on a constant
        let loss = Tensor::new(1.0_f32, &device).expect("failed to create tensor");
        let grads = loss.backward().expect("backward failed");

        let vars: Vec<Var> = vec![];
        let result = check_gradients_finite(&vars, &grads);
        assert!(result.is_ok(), "Empty vars should pass");
    }
}