foxhunt/ml/tests/dqn_q_value_statistics_test.rs

//! Q-Value Statistics Tests for DQN (WAVE 3 AGENT 3)
//!
//! Tests for Q-value statistics calculation (mean, std, range) to monitor
//! training stability and detect divergence.

use candle_core::{Device, Tensor};
use ml::dqn::{Experience, WorkingDQN, WorkingDQNConfig};

/// Test 1: Q-value statistics calculation correctness
#[test]
fn test_q_value_statistics_calculation() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 10;
    config.batch_size = 10;
    config.state_dim = 128;

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

    // Add diverse experiences to populate replay buffer
    for i in 0..50 {
        let experience = Experience::new(
            vec![i as f32 * 0.01; 128],
            (i % 3) as u8,
            (i as f32) * 0.1,
            vec![(i + 1) as f32 * 0.01; 128],
            false,
        );
        dqn.store_experience(experience)?;
    }

    // Train to produce non-trivial Q-values
    for _ in 0..5 {
        let _ = dqn.train_step(None)?;
    }

    // Sample multiple states to compute statistics
    let sample_size = 10;
    let mut all_q_values = Vec::new();

    for i in 0..sample_size {
        let test_state = Tensor::from_vec(vec![i as f32 * 0.05; 128], (1, 128), &device)?;

        let q_values = dqn.forward(&test_state)?;
        let q_vec = q_values.to_vec2::<f32>()?;

        // Collect all 3 Q-values (BUY, SELL, HOLD)
        for &q in &q_vec[0] {
            all_q_values.push(q as f64);
        }
    }

    // Calculate expected statistics manually
    let mean = all_q_values.iter().sum::<f64>() / all_q_values.len() as f64;

    let variance = all_q_values
        .iter()
        .map(|&x| {
            let diff = x - mean;
            diff * diff
        })
        .sum::<f64>()
        / all_q_values.len() as f64;
    let std = variance.sqrt();

    let min_q = all_q_values
        .iter()
        .copied()
        .min_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap();
    let max_q = all_q_values
        .iter()
        .copied()
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap();
    let range = max_q - min_q;

    // Verify statistics are reasonable
    assert!(mean.is_finite(), "Mean should be finite");
    assert!(std.is_finite(), "Std should be finite");
    assert!(std >= 0.0, "Std should be non-negative");
    assert!(range >= 0.0, "Range should be non-negative");
    assert!(range >= std, "Range should be >= std");

    println!(
        "Q-value statistics: mean={:.4}, std={:.4}, range={:.4}",
        mean, std, range
    );
    println!(
        "Q-values sample: {:?}",
        &all_q_values[..3.min(all_q_values.len())]
    );

    Ok(())
}

/// Test 2: Q-value variance detection for unstable training
#[test]
fn test_high_variance_detection() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 10;
    config.batch_size = 10;
    config.state_dim = 128;
    config.learning_rate = 0.1; // High LR to induce instability

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

    // Add experiences with extreme reward variance
    for i in 0..50 {
        let extreme_reward = if i % 2 == 0 { 100.0 } else { -100.0 };
        let experience = Experience::new(
            vec![i as f32 * 0.01; 128],
            (i % 3) as u8,
            extreme_reward,
            vec![(i + 1) as f32 * 0.01; 128],
            false,
        );
        dqn.store_experience(experience)?;
    }

    // Train with high variance data
    for _ in 0..20 {
        let _ = dqn.train_step(None)?;
    }

    // Compute Q-value statistics
    let mut all_q_values = Vec::new();
    for i in 0..10 {
        let test_state = Tensor::from_vec(vec![i as f32 * 0.05; 128], (1, 128), &device)?;

        let q_values = dqn.forward(&test_state)?;
        let q_vec = q_values.to_vec2::<f32>()?;

        for &q in &q_vec[0] {
            all_q_values.push(q as f64);
        }
    }

    let mean = all_q_values.iter().sum::<f64>() / all_q_values.len() as f64;
    let variance = all_q_values
        .iter()
        .map(|&x| (x - mean) * (x - mean))
        .sum::<f64>()
        / all_q_values.len() as f64;
    let std = variance.sqrt();

    // With extreme rewards and high LR, variance should be detectable
    assert!(
        std.is_finite(),
        "Std should be finite even with extreme training"
    );

    println!("High variance training: std={:.4}, mean={:.4}", std, mean);

    // Test passes if we can compute statistics without panic
    // Actual variance may vary but should be measurable
    Ok(())
}

/// Test 3: Q-value range detection for divergence
#[test]
fn test_q_value_range_tracking() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 10;
    config.batch_size = 10;
    config.state_dim = 128;

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

    // Add normal experiences
    for i in 0..50 {
        let experience = Experience::new(
            vec![i as f32 * 0.01; 128],
            (i % 3) as u8,
            (i as f32) * 0.1,
            vec![(i + 1) as f32 * 0.01; 128],
            false,
        );
        dqn.store_experience(experience)?;
    }

    // Train normally
    for _ in 0..10 {
        let _ = dqn.train_step(None)?;
    }

    // Compute Q-value range
    let mut all_q_values = Vec::new();
    for i in 0..10 {
        let test_state = Tensor::from_vec(vec![i as f32 * 0.05; 128], (1, 128), &device)?;

        let q_values = dqn.forward(&test_state)?;
        let q_vec = q_values.to_vec2::<f32>()?;

        for &q in &q_vec[0] {
            all_q_values.push(q as f64);
        }
    }

    let min_q = all_q_values
        .iter()
        .copied()
        .min_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap();
    let max_q = all_q_values
        .iter()
        .copied()
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap();
    let range = max_q - min_q;

    // Range should be reasonable (not extreme)
    assert!(range >= 0.0, "Range should be non-negative");
    assert!(range.is_finite(), "Range should be finite");

    // With normal training, range should be bounded
    // (This is a sanity check, not a strict requirement)
    println!(
        "Q-value range: {:.4} (min={:.4}, max={:.4})",
        range, min_q, max_q
    );

    Ok(())
}

/// Test 4: Statistics with zero Q-values (edge case)
#[test]
fn test_statistics_with_zero_q_values() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 4;
    config.batch_size = 4;
    config.state_dim = 128;

    let dqn = WorkingDQN::new(config)?;
    let device = dqn.device().clone();

    // Freshly initialized network should have Q-values near zero
    let test_state = Tensor::from_vec(vec![0.0_f32; 128], (1, 128), &device)?;

    let q_values = dqn.forward(&test_state)?;
    let q_vec = q_values.to_vec2::<f32>()?;

    // Compute statistics
    let q_doubles: Vec<f64> = q_vec[0].iter().map(|&x| x as f64).collect();
    let mean = q_doubles.iter().sum::<f64>() / q_doubles.len() as f64;
    let variance = q_doubles
        .iter()
        .map(|&x| (x - mean) * (x - mean))
        .sum::<f64>()
        / q_doubles.len() as f64;
    let std = variance.sqrt();

    let min_q = q_doubles
        .iter()
        .copied()
        .min_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap();
    let max_q = q_doubles
        .iter()
        .copied()
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap();
    let range = max_q - min_q;

    // All statistics should be finite
    assert!(mean.is_finite(), "Mean should be finite");
    assert!(std.is_finite(), "Std should be finite");
    assert!(range.is_finite(), "Range should be finite");

    // With initialization, Q-values should be small
    assert!(mean.abs() < 10.0, "Initial mean should be small");

    println!(
        "Initial Q-values: mean={:.4}, std={:.4}, range={:.4}",
        mean, std, range
    );

    Ok(())
}

/// Test 5: Statistics remain stable across multiple training steps
#[test]
fn test_statistics_stability_across_training() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 10;
    config.batch_size = 10;
    config.state_dim = 128;
    config.learning_rate = 0.001; // Conservative LR

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

    // Add consistent experiences
    for i in 0..50 {
        let experience = Experience::new(
            vec![i as f32 * 0.01; 128],
            (i % 3) as u8,
            1.0, // Consistent reward
            vec![(i + 1) as f32 * 0.01; 128],
            false,
        );
        dqn.store_experience(experience)?;
    }

    let mut prev_std = 0.0;
    let mut std_changes = Vec::new();

    // Train and track statistics evolution
    for step in 0..10 {
        let _ = dqn.train_step(None)?;

        // Compute statistics every step
        let mut all_q_values = Vec::new();
        for i in 0..5 {
            let test_state = Tensor::from_vec(vec![i as f32 * 0.05; 128], (1, 128), &device)?;

            let q_values = dqn.forward(&test_state)?;
            let q_vec = q_values.to_vec2::<f32>()?;

            for &q in &q_vec[0] {
                all_q_values.push(q as f64);
            }
        }

        let mean = all_q_values.iter().sum::<f64>() / all_q_values.len() as f64;
        let variance = all_q_values
            .iter()
            .map(|&x| (x - mean) * (x - mean))
            .sum::<f64>()
            / all_q_values.len() as f64;
        let std = variance.sqrt();

        if step > 0 {
            std_changes.push((std - prev_std).abs());
        }
        prev_std = std;

        assert!(std.is_finite(), "Std should remain finite at step {}", step);
    }

    // Statistics should evolve smoothly (no wild jumps)
    let max_change = std_changes
        .iter()
        .copied()
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap_or(0.0);

    println!("Max std change across training: {:.4}", max_change);
    println!("Std changes: {:?}", std_changes);

    // Test passes if statistics remain finite and measurable
    Ok(())
}