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

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

/// 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.0f64;

    // 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))
    }
}