//! Gradient Clipping Integration Test
//!
//! Tests DQN with gradient clipping enabled to ensure:
//! 1. Training completes without errors
//! 2. Gradient norms are tracked correctly
//! 3. Loss remains bounded (doesn't explode)

use anyhow::Result;
use ml::dqn::dqn::{WorkingDQN, WorkingDQNConfig};
use ml::dqn::Experience;

/// Test that DQN training works with gradient clipping enabled
#[test]
fn test_dqn_with_gradient_clipping() -> Result<()> {
    // Create DQN with gradient clipping enabled
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.gradient_clip_norm = Some(1.0); // Enable clipping with max_norm=1.0
    config.batch_size = 32;
    config.min_replay_size = 32;

    let mut dqn = WorkingDQN::new(config)?;

    // Fill replay buffer with dummy experiences
    for _ in 0..100 {
        let state: Vec<f32> = (0..32).map(|i| (i as f32) * 0.1).collect();
        let next_state: Vec<f32> = (0..32).map(|i| (i as f32) * 0.1 + 0.01).collect();

        let experience = Experience::new(
            state, 0,   // action
            1.0, // reward
            next_state, false, // done
        );

        dqn.store_experience(experience)?;
    }

    // Train for a few steps and verify it works
    let mut losses = Vec::new();
    for _ in 0..10 {
        let (loss, grad_norm) = dqn.train_step(None)?;
        losses.push(loss);

        // Verify loss is finite
        assert!(loss.is_finite(), "Loss should be finite, got: {}", loss);
        assert!(loss >= 0.0, "Loss should be non-negative, got: {}", loss);

        // Verify gradient norm is tracked (if clipping is enabled, it should be > 0)
        // Note: grad_norm is 0.0 if clipping is disabled
        if grad_norm > 0.0 {
            println!(
                "Step with clipping: loss={:.4}, grad_norm={:.4}",
                loss, grad_norm
            );
        }
    }

    // Verify training progressed (loss should change)
    let loss_variance = losses
        .iter()
        .map(|&l| (l - losses.iter().sum::<f32>() / losses.len() as f32).powi(2))
        .sum::<f32>()
        / losses.len() as f32;

    assert!(
        loss_variance > 1e-10,
        "Loss should vary during training, got variance: {:.6}",
        loss_variance
    );

    println!("✅ DQN with gradient clipping trained successfully");
    println!("   Losses: {:?}", losses);
    println!("   Variance: {:.6}", loss_variance);

    Ok(())
}

/// Test that DQN training works without gradient clipping
#[test]
fn test_dqn_without_gradient_clipping() -> Result<()> {
    // Create DQN with gradient clipping disabled
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.gradient_clip_norm = None; // Disable clipping
    config.batch_size = 32;
    config.min_replay_size = 32;

    let mut dqn = WorkingDQN::new(config)?;

    // Fill replay buffer with dummy experiences
    for _ in 0..100 {
        let state: Vec<f32> = (0..32).map(|i| (i as f32) * 0.1).collect();
        let next_state: Vec<f32> = (0..32).map(|i| (i as f32) * 0.1 + 0.01).collect();

        let experience = Experience::new(
            state, 0,   // action
            1.0, // reward
            next_state, false, // done
        );

        dqn.store_experience(experience)?;
    }

    // Train for a few steps
    for _ in 0..10 {
        let (loss, grad_norm) = dqn.train_step(None)?;

        // Verify loss is finite
        assert!(loss.is_finite(), "Loss should be finite, got: {}", loss);
        assert!(loss >= 0.0, "Loss should be non-negative, got: {}", loss);

        // Verify gradient norm is 0.0 when clipping is disabled
        assert_eq!(
            grad_norm, 0.0,
            "Gradient norm should be 0.0 when clipping is disabled"
        );
    }

    println!("✅ DQN without gradient clipping trained successfully");

    Ok(())
}