//! Comprehensive tests for NovoGrad optimizer
//!
//! NovoGrad: Layer-wise gradient normalization with adaptive second moment averaging
//! Key features tested:
//! - Layer-wise gradient normalization: g_norm = g / ||g||
//! - Adaptive second moment: v_t = β2 * v_{t-1} + (1-β2) * ||g||²
//! - Momentum with normalized gradients: m_t = β1 * m_{t-1} + (lr / √(v_t + ε)) * g_norm
//! - Weight update: θ_t+1 = θ_t - m_t - λ * θ_t (with weight decay)
//! - Bias correction for early iterations
//! - Gradient clipping support

#[cfg(test)]
mod tests {
    use super::super::{NovoGradOptimizer, NovoGradConfig, Optimizer};
    use crate::{Result, TransformerError};
    use rtx_tensor::{Tensor, Device, DType};
    use std::collections::HashMap;

    // Test configuration defaults
    fn default_novograd_config() -> NovoGradConfig {
        NovoGradConfig {
            learning_rate: 1e-3,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.0,
            grad_averaging: false,
        }
    }

    #[test]
    fn test_novograd_config_validation() {
        // Valid config should pass
        let config = default_novograd_config();
        assert!(NovoGradOptimizer::new(config).is_ok());

        // Invalid learning rate
        let mut invalid_config = default_novograd_config();
        invalid_config.learning_rate = -0.1;
        assert!(NovoGradOptimizer::new(invalid_config).is_err());

        // Invalid beta1
        let mut invalid_config = default_novograd_config();
        invalid_config.beta1 = 1.1;
        assert!(NovoGradOptimizer::new(invalid_config).is_err());

        // Invalid beta2
        let mut invalid_config = default_novograd_config();
        invalid_config.beta2 = -0.1;
        assert!(NovoGradOptimizer::new(invalid_config).is_err());

        // Invalid eps
        let mut invalid_config = default_novograd_config();
        invalid_config.eps = -1e-8;
        assert!(NovoGradOptimizer::new(invalid_config).is_err());

        // Invalid weight decay
        let mut invalid_config = default_novograd_config();
        invalid_config.weight_decay = -0.5;
        assert!(NovoGradOptimizer::new(invalid_config).is_err());
    }

    #[test]
    fn test_novograd_parameter_creation() -> Result<()> {
        let config = default_novograd_config();
        let optimizer = NovoGradOptimizer::new(config)?;

        assert_eq!(optimizer.learning_rate(), 1e-3);
        assert_eq!(optimizer.beta1(), 0.95);
        assert_eq!(optimizer.beta2(), 0.98);
        assert_eq!(optimizer.eps(), 1e-8);
        assert_eq!(optimizer.weight_decay(), 0.0);
        assert_eq!(optimizer.grad_averaging(), false);
        assert_eq!(optimizer.optimizer_type(), "NovoGrad");

        Ok(())
    }

    #[test]
    fn test_novograd_gradient_normalization() -> Result<()> {
        let config = NovoGradConfig {
            learning_rate: 0.01,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.0,
            grad_averaging: false,
        };

        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        // Create test parameter and gradient with known norm
        let param = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], &[4], &device)?;
        let grad = Tensor::from_data(vec![3.0, 4.0, 0.0, 0.0], &[4], &device)?; // ||grad|| = 5.0

        // First step - should normalize gradient by its norm
        let updated_param1 = optimizer.step_param("test_param", &param, &grad)?;
        
        // Verify state was created
        assert!(optimizer.has_state("test_param"));
        assert_eq!(optimizer.get_step_count("test_param")?, 1);

        // For first step with normalized gradient: g_norm = [0.6, 0.8, 0.0, 0.0]
        // Second moment: v_1 = (1-0.98) * 25 = 0.5
        // Momentum: m_1 = (0.01 / √(0.5 + 1e-8)) * [0.6, 0.8, 0.0, 0.0]
        let param_data = param.to_cpu()?;
        let updated_data = updated_param1.to_cpu()?;
        
        // Parameters should be updated in proportion to normalized gradient
        let update_ratio = (param_data[1] - updated_data[1]) / (param_data[0] - updated_data[0]);
        assert!((update_ratio - 4.0/3.0).abs() < 1e-6, 
                "Update should be proportional to normalized gradient: expected 4/3, got {}", update_ratio);

        Ok(())
    }

    #[test]
    fn test_novograd_momentum_accumulation() -> Result<()> {
        let config = NovoGradConfig {
            learning_rate: 0.01,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.0,
            grad_averaging: false,
        };

        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![1.0], &[1], &device)?;
        let grad = Tensor::from_data(vec![2.0], &[1], &device)?; // ||grad|| = 2.0, normalized = 1.0

        // Take multiple steps with same gradient
        let mut current_param = param.clone();
        let mut momentum_effects = Vec::new();
        
        for step in 0..5 {
            let prev_data = current_param.to_cpu()?;
            current_param = optimizer.step_param("param", &current_param, &grad)?;
            let new_data = current_param.to_cpu()?;
            
            let update_magnitude = (prev_data[0] - new_data[0]).abs();
            momentum_effects.push(update_magnitude);
            
            assert_eq!(optimizer.get_step_count("param")?, step + 1);
        }

        // Momentum should build up over iterations for consistent gradient direction
        // Later updates should be influenced by accumulated momentum
        assert!(momentum_effects.len() == 5, "Should have 5 momentum measurements");

        Ok(())
    }

    #[test]
    fn test_novograd_second_moment_averaging() -> Result<()> {
        let config = NovoGradConfig {
            learning_rate: 0.01,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.0,
            grad_averaging: false,
        };

        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![1.0], &[1], &device)?;
        
        // Use different gradient magnitudes to test second moment averaging
        let gradients = vec![
            Tensor::from_data(vec![1.0], &[1], &device)?, // ||g|| = 1.0
            Tensor::from_data(vec![2.0], &[1], &device)?, // ||g|| = 2.0
            Tensor::from_data(vec![0.5], &[1], &device)?, // ||g|| = 0.5
        ];

        let mut current_param = param;
        for (step, grad) in gradients.iter().enumerate() {
            current_param = optimizer.step_param("param", &current_param, grad)?;
            assert_eq!(optimizer.get_step_count("param")?, step as i64 + 1);
        }

        // The optimizer should have adapted to different gradient norms
        assert!(optimizer.has_state("param"));

        Ok(())
    }

    #[test]
    fn test_novograd_bias_correction() -> Result<()> {
        let config = NovoGradConfig {
            learning_rate: 0.1,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.0,
            grad_averaging: false,
        };

        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![1.0], &[1], &device)?;
        let grad = Tensor::from_data(vec![1.0], &[1], &device)?;

        // First few steps should show bias correction effects
        let mut update_magnitudes = Vec::new();
        let mut current_param = param;

        for _step in 0..10 {
            let prev_data = current_param.to_cpu()?;
            current_param = optimizer.step_param("param", &current_param, &grad)?;
            let new_data = current_param.to_cpu()?;
            
            let update_magnitude = (prev_data[0] - new_data[0]).abs();
            update_magnitudes.push(update_magnitude);
        }

        // Early steps should have larger updates due to bias correction
        // Later steps should stabilize
        assert!(update_magnitudes[0] > update_magnitudes[9] * 0.5,
                "Early steps should benefit from bias correction");

        Ok(())
    }

    #[test]
    fn test_novograd_multi_parameter() -> Result<()> {
        let config = default_novograd_config();
        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        // Create multiple parameters with different layer characteristics
        let weight1 = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], &[2, 2], &device)?;
        let weight2 = Tensor::from_data(vec![0.1, 0.2], &[2], &device)?;
        let bias = Tensor::from_data(vec![0.01], &[1], &device)?;

        // Create gradients with different norms (simulating different layers)
        let weight1_grad = Tensor::from_data(vec![0.1, -0.1, 0.2, -0.2], &[2, 2], &device)?; // norm ≈ 0.316
        let weight2_grad = Tensor::from_data(vec![1.0, -1.0], &[2], &device)?; // norm ≈ 1.414
        let bias_grad = Tensor::from_data(vec![0.01], &[1], &device)?; // norm = 0.01

        // Store gradients and process
        let mut gradients = HashMap::new();
        gradients.insert("weight1".to_string(), weight1_grad);
        gradients.insert("weight2".to_string(), weight2_grad);
        gradients.insert("bias".to_string(), bias_grad);

        optimizer.set_gradients(gradients)?;
        let parameter_updates = optimizer.step(0.01)?;

        // Verify updates exist for all parameters
        assert!(parameter_updates.contains_key("weight1"));
        assert!(parameter_updates.contains_key("weight2"));
        assert!(parameter_updates.contains_key("bias"));

        // Verify shapes match
        assert_eq!(parameter_updates["weight1"].shape().dims(), &[2, 2]);
        assert_eq!(parameter_updates["weight2"].shape().dims(), &[2]);
        assert_eq!(parameter_updates["bias"].shape().dims(), &[1]);

        // Verify state was created for all parameters
        assert!(optimizer.has_state("weight1"));
        assert!(optimizer.has_state("weight2"));
        assert!(optimizer.has_state("bias"));

        Ok(())
    }

    #[test]
    fn test_novograd_weight_decay() -> Result<()> {
        let config = NovoGradConfig {
            learning_rate: 0.1,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.01,
            grad_averaging: false,
        };

        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![2.0], &[1], &device)?; // Non-zero parameter
        let grad = Tensor::from_data(vec![0.0], &[1], &device)?; // Zero gradient

        // With zero gradient, only weight decay should affect parameter
        let updated_param = optimizer.step_param("param", &param, &grad)?;
        let param_data = param.to_cpu()?;
        let updated_data = updated_param.to_cpu()?;

        // Weight decay should reduce parameter magnitude: θ_new = θ_old - λ * θ_old
        let expected = param_data[0] * (1.0 - 0.01);
        assert!((updated_data[0] - expected).abs() < 1e-6, 
                "Weight decay should reduce parameter: expected {}, got {}", expected, updated_data[0]);

        Ok(())
    }

    #[test]
    fn test_novograd_gradient_averaging() -> Result<()> {
        let config = NovoGradConfig {
            learning_rate: 0.01,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.0,
            grad_averaging: true,
        };

        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![1.0, 2.0], &[2], &device)?;
        let grad = Tensor::from_data(vec![3.0, 4.0], &[2], &device)?; // ||grad|| = 5.0

        let updated_param = optimizer.step_param("param", &param, &grad)?;
        
        // With gradient averaging enabled, behavior should differ from standard NovoGrad
        assert!(optimizer.has_state("param"));
        assert_eq!(optimizer.get_step_count("param")?, 1);

        Ok(())
    }

    #[test]
    fn test_novograd_learning_rate_updates() -> Result<()> {
        let config = default_novograd_config();
        let mut optimizer = NovoGradOptimizer::new(config)?;

        assert_eq!(optimizer.learning_rate(), 1e-3);

        // Update learning rate
        optimizer.set_learning_rate(0.01)?;
        assert_eq!(optimizer.learning_rate(), 0.01);

        // Invalid learning rate should error
        assert!(optimizer.set_learning_rate(-0.1).is_err());

        Ok(())
    }

    #[test]
    fn test_novograd_gradient_clipping_compatibility() -> Result<()> {
        let config = default_novograd_config();
        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        // Create large gradients that would need clipping
        let large_grad1 = Tensor::from_data(vec![10.0, 20.0], &[2], &device)?;
        let large_grad2 = Tensor::from_data(vec![30.0], &[1], &device)?;

        let mut gradients = HashMap::new();
        gradients.insert("param1".to_string(), large_grad1);
        gradients.insert("param2".to_string(), large_grad2);

        optimizer.set_gradients(gradients)?;

        // Should be able to compute gradient norm before clipping
        let total_norm_squared = 10.0*10.0 + 20.0*20.0 + 30.0*30.0; // 100 + 400 + 900 = 1400
        let expected_norm = total_norm_squared.sqrt(); // ≈ 37.42

        // The optimizer should handle large gradients without issues
        let parameter_updates = optimizer.step(0.01)?;
        assert_eq!(parameter_updates.len(), 2);

        Ok(())
    }

    #[test]
    fn test_novograd_zero_gradient_handling() -> Result<()> {
        let config = default_novograd_config();
        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![1.0], &[1], &device)?;
        let zero_grad = Tensor::from_data(vec![0.0], &[1], &device)?;

        // Zero gradient should not cause errors (gradient norm will be 0)
        let updated_param = optimizer.step_param("param", &param, &zero_grad)?;
        
        let param_data = param.to_cpu()?;
        let updated_data = updated_param.to_cpu()?;
        
        // With zero gradient and no weight decay, parameter should be unchanged
        // (normalized gradient is 0/0 which should be handled gracefully)
        assert!((updated_data[0] - param_data[0]).abs() < 1e-6, 
                "Zero gradient should not change parameter significantly");

        Ok(())
    }

    #[test]
    fn test_novograd_convergence_quadratic() -> Result<()> {
        let config = NovoGradConfig {
            learning_rate: 0.1,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.0,
            grad_averaging: false,
        };

        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        // Optimize f(x) = (x - 3)^2, minimum at x = 3
        let target = 3.0;
        let mut param = Tensor::from_data(vec![0.0], &[1], &device)?;
        let mut loss_history = Vec::new();

        for _step in 0..50 {
            let param_data = param.to_cpu()?;
            let current_value = param_data[0];
            let loss = (current_value - target).powi(2);
            loss_history.push(loss);

            // Gradient of (x - 3)^2 is 2(x - 3)
            let gradient_value = 2.0 * (current_value - target);
            let grad = Tensor::from_data(vec![gradient_value], &[1], &device)?;

            param = optimizer.step_param("param", &param, &grad)?;
        }

        let final_data = param.to_cpu()?;
        let final_value = final_data[0];
        let final_loss = (final_value - target).powi(2);

        // Should converge reasonably close to target
        assert!(
            (final_value - target).abs() < 0.5,
            "Should converge close to target, got {} vs {}",
            final_value, target
        );

        // Loss should decrease significantly
        assert!(
            final_loss < loss_history[0] * 0.1,
            "Loss should decrease significantly from {} to {}",
            loss_history[0], final_loss
        );

        Ok(())
    }

    #[test]
    fn test_novograd_state_management() -> Result<()> {
        let config = default_novograd_config();
        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![1.0, 2.0, 3.0], &[3], &device)?;
        let grad = Tensor::from_data(vec![0.1, 0.2, 0.3], &[3], &device)?;

        // Initialize state
        optimizer.step_param("param", &param, &grad)?;
        
        // NovoGrad should store momentum and second moment (per-layer)
        assert!(optimizer.has_state("param"));

        // Reset state and verify cleanup
        optimizer.reset_state("param")?;
        assert!(!optimizer.has_state("param"));

        // Reset all state
        optimizer.step_param("param1", &param, &grad)?;
        optimizer.step_param("param2", &param, &grad)?;
        optimizer.reset_all_state();
        assert!(!optimizer.has_state("param1"));
        assert!(!optimizer.has_state("param2"));

        Ok(())
    }

    #[test]
    fn test_novograd_gradient_shapes() -> Result<()> {
        let config = default_novograd_config();
        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![1.0, 2.0], &[2], &device)?;
        let wrong_grad = Tensor::from_data(vec![0.1], &[1], &device)?;

        // Shape mismatch should error
        let result = optimizer.step_param("param", &param, &wrong_grad);
        assert!(result.is_err());

        Ok(())
    }

    #[test]
    fn test_novograd_deterministic_behavior() -> Result<()> {
        let config = default_novograd_config();
        let device = Device::cpu();

        // Run same optimization twice
        let mut results1 = Vec::new();
        let mut results2 = Vec::new();

        for trial in 0..2 {
            let mut optimizer = NovoGradOptimizer::new(config.clone())?;
            let param = Tensor::from_data(vec![1.0], &[1], &device)?;
            let grad = Tensor::from_data(vec![0.5], &[1], &device)?;

            let mut current_param = param;
            for _step in 0..5 {
                current_param = optimizer.step_param("param", &current_param, &grad)?;
                let value = current_param.to_cpu()?[0];
                
                if trial == 0 {
                    results1.push(value);
                } else {
                    results2.push(value);
                }
            }
        }

        // Results should be identical (deterministic)
        assert_eq!(results1.len(), results2.len());
        for (r1, r2) in results1.iter().zip(results2.iter()) {
            assert!(
                (r1 - r2).abs() < 1e-6,
                "Results should be deterministic: {} vs {}",
                r1, r2
            );
        }

        Ok(())
    }

    #[test]
    fn test_novograd_large_batch_stability() -> Result<()> {
        // NovoGrad is particularly designed for large batch training
        let config = NovoGradConfig {
            learning_rate: 0.01,
            beta1: 0.95,
            beta2: 0.98,
            eps: 1e-8,
            weight_decay: 0.0,
            grad_averaging: false,
        };

        let mut optimizer = NovoGradOptimizer::new(config)?;
        let device = Device::cpu();

        // Simulate large batch gradients (high variance)
        let param = Tensor::from_data(vec![1.0, 2.0, 3.0], &[3], &device)?;
        
        let batch_gradients = vec![
            Tensor::from_data(vec![1.0, 0.1, 0.01], &[3], &device)?,     // Different scales
            Tensor::from_data(vec![0.01, 1.0, 0.1], &[3], &device)?,
            Tensor::from_data(vec![0.1, 0.01, 1.0], &[3], &device)?,
        ];

        let mut current_param = param;
        for (step, grad) in batch_gradients.iter().enumerate() {
            let prev_param = current_param.clone();
            current_param = optimizer.step_param("param", &current_param, grad)?;
            
            // Should remain stable and not explode
            let param_data = current_param.to_cpu()?;
            for &value in param_data.iter() {
                assert!(value.is_finite(), "Parameter should remain finite at step {}", step);
                assert!(value.abs() < 100.0, "Parameter should not explode at step {}", step);
            }
        }

        Ok(())
    }
}