//! Gradient Safety Framework
//!
//! This module provides comprehensive gradient safety controls for ML training,
//! including gradient clipping, NaN detection, explosion protection, and
//! mathematical stability guarantees for all optimization steps.

use std;

use std::collections::VecDeque;
use std::sync::Arc;

// candle_core::Tensor removed — gradient safety now operates on f64 norm values.
// GPU gradient operations are handled by cuda_autograd::GpuAdamW.
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};

use super::{MLSafetyError, SafetyResult};

/// Gradient safety configuration
///
/// SAFETY: All values should come from configuration system, not hardcoded defaults
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GradientSafetyConfig {
    /// Maximum L2 norm for gradient clipping
    pub max_gradient_norm: f64,
    /// Minimum gradient norm (detect vanishing gradients)
    pub min_gradient_norm: f64,
    /// Maximum individual gradient value
    pub max_individual_gradient: f64,
    /// Enable gradient clipping by norm
    pub enable_norm_clipping: bool,
    /// Enable gradient clipping by value
    pub enable_value_clipping: bool,
    /// Enable NaN/Infinity detection
    pub enable_nan_detection: bool,
    /// History window size for gradient statistics
    pub gradient_history_size: usize,
    /// Gradient explosion detection threshold (ratio of current to average)
    pub explosion_threshold: f64,
    /// Minimum number of gradients before explosion detection
    pub min_gradient_history: usize,
    /// Enable adaptive gradient scaling
    pub enable_adaptive_scaling: bool,
    /// Learning rate adjustment factor for gradient explosions
    pub lr_adjustment_factor: f64,
    /// Base learning rate for safety calculations
    pub base_learning_rate: f64,
}

impl GradientSafetyConfig {
    /// Create from configuration system - eliminates hardcoded defaults
    pub fn from_config_manager(
        _config_manager: &config::ConfigManager,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        // Use emergency defaults since specific gradient safety configs may not be available
        tracing::warn!("Using emergency gradient safety config defaults - gradient safety configs not available in ServiceConfig");
        Ok(Self::emergency_safe_defaults())
    }

    /// EMERGENCY FALLBACK: Ultra-conservative gradient safety defaults
    ///
    /// WARNING: These are safety-first defaults that may severely limit training
    pub fn emergency_safe_defaults() -> Self {
        tracing::warn!("Using emergency gradient safety defaults - check config system!");
        Self {
            max_gradient_norm: 1.0,       // Very aggressive clipping
            min_gradient_norm: 1e-10,     // Detect very small gradients
            max_individual_gradient: 1.0, // Conservative individual gradient limit
            enable_norm_clipping: true,   // Enable all safety features
            enable_value_clipping: true,
            enable_nan_detection: true,
            gradient_history_size: 10, // Small history to save memory
            explosion_threshold: 2.0,  // Very sensitive explosion detection
            min_gradient_history: 3,   // Minimal history before detection
            enable_adaptive_scaling: true,
            lr_adjustment_factor: 0.1, // Aggressive LR reduction
            base_learning_rate: 1e-6,  // Ultra-conservative learning rate
        }
    }
}

impl Default for GradientSafetyConfig {
    fn default() -> Self {
        // Redirect to emergency safe defaults to prevent hardcoded usage
        Self::emergency_safe_defaults()
    }
}

/// Gradient safety errors
#[derive(Error, Debug)]
pub enum GradientSafetyError {
    #[error("Gradient explosion detected: norm {norm:.3} > threshold {threshold:.3}")]
    GradientExplosion { norm: f64, threshold: f64 },

    #[error("Gradient vanishing detected: norm {norm:.3e} < threshold {threshold:.3e}")]
    GradientVanishing { norm: f64, threshold: f64 },

    #[error("NaN gradient detected in parameter: {parameter}")]
    NaNGradient { parameter: String },

    #[error("Infinite gradient detected in parameter: {parameter}, value: {value}")]
    InfiniteGradient { parameter: String, value: f64 },

    #[error("Gradient value out of bounds: {value} not in [{min}, {max}]")]
    GradientOutOfBounds { value: f64, min: f64, max: f64 },

    #[error("Gradient computation failed: {reason}")]
    ComputationFailed { reason: String },
}

/// Gradient statistics for monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GradientStatistics {
    pub current_norm: f64,
    pub average_norm: f64,
    pub max_norm: f64,
    pub min_norm: f64,
    pub nan_count: u64,
    pub infinity_count: u64,
    pub clipping_count: u64,
    pub explosion_count: u64,
    pub vanishing_count: u64,
    pub adaptive_scaling_factor: f64,
}

impl Default for GradientStatistics {
    fn default() -> Self {
        Self {
            current_norm: 0.0,
            average_norm: 0.0,
            max_norm: 0.0,
            min_norm: f64::INFINITY,
            nan_count: 0,
            infinity_count: 0,
            clipping_count: 0,
            explosion_count: 0,
            vanishing_count: 0,
            adaptive_scaling_factor: 1.0,
        }
    }
}

/// Safe gradient computation and clipping system
#[derive(Debug)]
pub struct GradientSafetyManager {
    config: GradientSafetyConfig,
    statistics: Arc<RwLock<GradientStatistics>>,
    gradient_history: Arc<RwLock<VecDeque<f64>>>,
    current_learning_rate: Arc<RwLock<f64>>,
    original_learning_rate: f64,
}

impl GradientSafetyManager {
    /// Create new gradient safety manager
    pub fn new(config: GradientSafetyConfig, learning_rate: f64) -> Self {
        let history_size = config.gradient_history_size;
        Self {
            config,
            statistics: Arc::new(RwLock::new(GradientStatistics::default())),
            gradient_history: Arc::new(RwLock::new(VecDeque::with_capacity(history_size))),
            current_learning_rate: Arc::new(RwLock::new(learning_rate)),
            original_learning_rate: learning_rate,
        }
    }

    /// Process gradient norms with comprehensive safety checks.
    ///
    /// Takes per-parameter gradient norms (computed on GPU) and validates them.
    /// Actual gradient tensor manipulation is handled by `cuda_autograd::GpuAdamW`.
    pub async fn process_gradient_norms(
        &self,
        gradient_norms: &[f64],
        parameter_names: &[String],
    ) -> SafetyResult<()> {
        let mut stats = self.statistics.write().await;
        stats.current_norm = 0.0;

        let mut total_norm_squared = 0.0;
        for (i, &norm) in gradient_norms.iter().enumerate() {
            let default_name = format!("param_{}", i);
            let param_name = parameter_names
                .get(i)
                .map(|s| s.as_str())
                .unwrap_or(&default_name);

            if self.config.enable_nan_detection {
                if norm.is_nan() {
                    stats.nan_count += 1;
                    return Err(MLSafetyError::from(
                        GradientSafetyError::NaNGradient {
                            parameter: param_name.to_owned(),
                        },
                    )
                    .into());
                }
                if norm.is_infinite() {
                    stats.infinity_count += 1;
                    return Err(MLSafetyError::from(
                        GradientSafetyError::InfiniteGradient {
                            parameter: param_name.to_owned(),
                            value: norm,
                        },
                    )
                    .into());
                }
            }
            total_norm_squared += norm * norm;
        }

        let total_norm = total_norm_squared.sqrt();
        stats.current_norm = total_norm;

        self.update_gradient_statistics(total_norm, &mut stats).await;

        let anomaly_result = self.detect_gradient_anomalies(total_norm, &mut stats).await;

        if self.config.enable_adaptive_scaling {
            self.update_learning_rate(&mut stats).await;
        }

        anomaly_result?;

        debug!(
            "Processed {} gradient norms safely. Total norm: {:.6}",
            gradient_norms.len(),
            stats.current_norm
        );

        Ok(())
    }

    /// Update gradient statistics and history
    async fn update_gradient_statistics(&self, current_norm: f64, stats: &mut GradientStatistics) {
        // Update norm statistics
        stats.max_norm = stats.max_norm.max(current_norm);
        stats.min_norm = stats.min_norm.min(current_norm);

        // Update gradient history
        let mut history = self.gradient_history.write().await;
        history.push_back(current_norm);

        if history.len() > self.config.gradient_history_size {
            history.pop_front();
        }

        // Compute average norm
        if !history.is_empty() {
            stats.average_norm = history.iter().sum::<f64>() / history.len() as f64;
        }
    }

    /// Detect gradient explosion and vanishing gradients
    async fn detect_gradient_anomalies(
        &self,
        current_norm: f64,
        stats: &mut GradientStatistics,
    ) -> SafetyResult<()> {
        // Check for gradient explosion
        if current_norm > self.config.max_gradient_norm {
            stats.explosion_count += 1;
            warn!(
                "Gradient explosion detected: norm {:.6} > {:.6}",
                current_norm, self.config.max_gradient_norm
            );

            return Err(MLSafetyError::from(GradientSafetyError::GradientExplosion {
                norm: current_norm,
                threshold: self.config.max_gradient_norm,
            })
            .into());
        }

        // Check for gradient vanishing
        if current_norm < self.config.min_gradient_norm {
            stats.vanishing_count += 1;
            warn!(
                "Gradient vanishing detected: norm {:.3e} < {:.3e}",
                current_norm, self.config.min_gradient_norm
            );

            return Err(MLSafetyError::from(GradientSafetyError::GradientVanishing {
                norm: current_norm,
                threshold: self.config.min_gradient_norm,
            })
            .into());
        }

        // Check for relative explosion (compared to history average)
        let history = self.gradient_history.read().await;
        if history.len() >= self.config.min_gradient_history {
            let avg_norm = stats.average_norm;
            if avg_norm > 0.0 && current_norm > avg_norm * self.config.explosion_threshold {
                stats.explosion_count += 1;
                warn!(
                    "Relative gradient explosion: {:.6} / {:.6} = {:.3}x average",
                    current_norm,
                    avg_norm,
                    current_norm / avg_norm
                );

                return Err(MLSafetyError::from(GradientSafetyError::GradientExplosion {
                    norm: current_norm,
                    threshold: avg_norm * self.config.explosion_threshold,
                })
                .into());
            }
        }

        Ok(())
    }

    /// Compute the scale factor for gradient clipping.
    ///
    /// Returns `min(max_norm / (total_norm + eps), 1.0)`.
    /// Actual clipping is applied by `cuda_autograd::GpuAdamW` on GPU.
    pub fn clip_scale_factor(&self, total_norm: f64) -> f64 {
        if total_norm > self.config.max_gradient_norm {
            self.config.max_gradient_norm / (total_norm + 1e-6)
        } else {
            1.0
        }
    }

    /// Update learning rate based on gradient behavior
    async fn update_learning_rate(&self, stats: &mut GradientStatistics) {
        let mut current_lr = self.current_learning_rate.write().await;
        let old_lr = *current_lr;

        // Adaptive scaling based on gradient statistics
        if stats.explosion_count > 0 {
            // Reduce learning rate after gradient explosion
            *current_lr *= self.config.lr_adjustment_factor;
            stats.adaptive_scaling_factor = *current_lr / self.original_learning_rate;
            info!(
                "Reduced learning rate due to gradient explosion: {:.6} -> {:.6}",
                old_lr, *current_lr
            );
        } else if stats.vanishing_count > 0 {
            // Increase learning rate for vanishing gradients (but cap it)
            let increase_factor = 1.0 / self.config.lr_adjustment_factor;
            *current_lr = (*current_lr * increase_factor).min(self.original_learning_rate * 2.0);
            stats.adaptive_scaling_factor = *current_lr / self.original_learning_rate;
            info!(
                "Increased learning rate due to vanishing gradients: {:.6} -> {:.6}",
                old_lr, *current_lr
            );
        } else {
            // Gradients are healthy, no learning rate adjustment needed
        }
    }

    /// Get current learning rate
    pub async fn get_current_learning_rate(&self) -> f64 {
        *self.current_learning_rate.read().await
    }

    /// Get gradient statistics
    pub async fn get_statistics(&self) -> GradientStatistics {
        self.statistics.read().await.clone()
    }

    /// Reset statistics and history
    pub async fn reset_statistics(&self) {
        let mut stats = self.statistics.write().await;
        *stats = GradientStatistics::default();

        let mut history = self.gradient_history.write().await;
        history.clear();

        let mut lr = self.current_learning_rate.write().await;
        *lr = self.original_learning_rate;

        info!("Reset gradient safety statistics");
    }

    /// Emergency gradient reset — resets all statistics.
    ///
    /// Callers should zero out their gradient buffers separately.
    pub async fn emergency_gradient_reset(&self) {
        warn!("Emergency gradient reset activated — resetting safety statistics");
        self.reset_statistics().await;
    }
}

// Convert gradient safety errors to ML safety errors
impl From<GradientSafetyError> for MLSafetyError {
    fn from(err: GradientSafetyError) -> Self {
        match err {
            GradientSafetyError::GradientExplosion { norm, threshold } => {
                MLSafetyError::MathSafety {
                    reason: format!(
                        "Gradient explosion: norm {:.3} > threshold {:.3}",
                        norm, threshold
                    ),
                }
            },
            GradientSafetyError::GradientVanishing { norm, threshold } => {
                MLSafetyError::MathSafety {
                    reason: format!(
                        "Gradient vanishing: norm {:.3e} < threshold {:.3e}",
                        norm, threshold
                    ),
                }
            },
            GradientSafetyError::NaNGradient { parameter } => MLSafetyError::InvalidFloat {
                operation: format!("Gradient computation for parameter: {}", parameter),
            },
            GradientSafetyError::InfiniteGradient { parameter, value } => {
                MLSafetyError::InvalidFloat {
                    operation: format!(
                        "Gradient computation for parameter {}: {}",
                        parameter, value
                    ),
                }
            },
            GradientSafetyError::GradientOutOfBounds { value, min, max } => {
                MLSafetyError::PredictionOutOfBounds { value, min, max }
            },
            GradientSafetyError::ComputationFailed { reason } => {
                MLSafetyError::MathSafety { reason }
            },
        }
    }
}

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

    fn create_test_manager() -> GradientSafetyManager {
        let config = GradientSafetyConfig::emergency_safe_defaults();
        let learning_rate = config.base_learning_rate;
        GradientSafetyManager::new(config, learning_rate)
    }

    #[tokio::test]
    async fn test_normal_gradient_norms() {
        let manager = create_test_manager();

        // Normal gradient norms (below threshold)
        let norms = vec![0.1, 0.2];
        let param_names = vec!["weight".to_owned(), "bias".to_owned()];

        let result = manager.process_gradient_norms(&norms, &param_names).await;
        assert!(result.is_ok());

        let stats = manager.get_statistics().await;
        assert!(stats.current_norm > 0.0);
        assert_eq!(stats.nan_count, 0);
        assert_eq!(stats.infinity_count, 0);
    }

    #[tokio::test]
    async fn test_gradient_explosion_detection() {
        let config_safe = GradientSafetyConfig::emergency_safe_defaults();
        let manager =
            GradientSafetyManager::new(config_safe.clone(), config_safe.base_learning_rate);

        // Large gradient norm that should trigger explosion detection
        let norms = vec![15.0]; // norm > max_gradient_norm (1.0)
        let param_names = vec!["weight".to_owned()];

        let result = manager.process_gradient_norms(&norms, &param_names).await;
        assert!(result.is_err());

        let stats = manager.get_statistics().await;
        assert_eq!(stats.explosion_count, 1);
    }

    #[tokio::test]
    async fn test_nan_norm_detection() {
        let manager = create_test_manager();

        let norms = vec![f64::NAN];
        let param_names = vec!["weight".to_owned()];

        let result = manager.process_gradient_norms(&norms, &param_names).await;
        assert!(result.is_err());

        let stats = manager.get_statistics().await;
        assert_eq!(stats.nan_count, 1);
    }

    #[tokio::test]
    async fn test_infinity_norm_detection() {
        let manager = create_test_manager();

        let norms = vec![f64::INFINITY];
        let param_names = vec!["weight".to_owned()];

        let result = manager.process_gradient_norms(&norms, &param_names).await;
        assert!(result.is_err());

        let stats = manager.get_statistics().await;
        assert_eq!(stats.infinity_count, 1);
    }

    #[tokio::test]
    async fn test_learning_rate_adaptation() {
        let config_safe = GradientSafetyConfig::emergency_safe_defaults();
        let manager =
            GradientSafetyManager::new(config_safe.clone(), config_safe.base_learning_rate);

        let initial_lr = manager.get_current_learning_rate().await;
        assert_eq!(initial_lr, config_safe.base_learning_rate);

        // Trigger explosion via large norm
        let norms = vec![10.0];
        let param_names = vec!["weight".to_owned()];

        let _ = manager.process_gradient_norms(&norms, &param_names).await;

        let new_lr = manager.get_current_learning_rate().await;
        assert!(new_lr < initial_lr);
    }

    #[tokio::test]
    async fn test_emergency_reset() {
        let manager = create_test_manager();
        manager.emergency_gradient_reset().await;

        let stats = manager.get_statistics().await;
        assert_eq!(stats.current_norm, 0.0);
    }

    #[test]
    fn test_clip_scale_factor() {
        let manager = create_test_manager();
        // Norm below max: scale = 1.0
        assert!((manager.clip_scale_factor(0.5) - 1.0).abs() < 1e-6);
        // Norm above max: scale < 1.0
        assert!(manager.clip_scale_factor(10.0) < 1.0);
    }
}
