foxhunt/ml/tests/dqn_numerical_stability_test.rs

//! DQN Numerical Stability Tests
//!
//! Tests to expose and validate fixes for:
//! - Unbounded reward accumulation
//! - Q-value explosion (e.g., +24,055 in Trial 3)
//! - Gradient underflow (217 collapses observed)
//! - Missing Huber loss protection

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

/// Helper to create a state vector for testing
fn create_test_state(portfolio_value: f32) -> Vec<f32> {
    let mut state = vec![0.0; 52];
    // Price features (4)
    state[0..4].copy_from_slice(&[100.0, 101.0, 99.0, 100.5]);
    // Technical indicators (16)
    for i in 4..20 {
        state[i] = 0.5;
    }
    // Market features (16)
    for i in 20..36 {
        state[i] = 0.5;
    }
    // Portfolio features (16) - set portfolio value in first position
    state[36] = portfolio_value;
    state[37] = 0.0; // position
    state[38] = 0.001; // spread
    for i in 39..52 {
        state[i] = 0.0;
    }
    state
}

#[test]
fn test_rewards_stay_bounded() -> Result<()> {
    println!("\n=== TEST: Rewards Stay Bounded ===");

    // Create DQN with high penalty to trigger large rewards
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.batch_size = 4;
    config.min_replay_size = 10;

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

    // Populate replay buffer with extreme rewards (simulating unbounded accumulation)
    for i in 0..20 {
        // Create state with extreme portfolio value (200% gain)
        let portfolio_value = 2.0 + (i as f32 * 0.1); // 2.0 to 3.9
        let state = create_test_state(portfolio_value);
        let next_state = create_test_state(portfolio_value + 0.5);

        let action = TradingAction::from_int((i % 3) as u8).unwrap();

        // Simulate extreme rewards that would come from unbounded P&L
        // This mimics the issue found in reward.rs:144-156
        let reward = if i < 10 {
            1.0 + (i as f32 * 0.1) // Rewards > 1.0 (unbounded)
        } else {
            -1.0 - ((i - 10) as f32 * 0.1) // Rewards < -1.0 (unbounded)
        };

        let experience = Experience::new(
            state.clone(),
            action as u8,
            reward,
            next_state.clone(),
            false,
        );

        dqn.store_experience(experience)?;

        println!(
            "Step {}: Reward={:.4} (portfolio_value={:.2})",
            i, reward, portfolio_value
        );

        // NOTE: This test will FAIL until reward clipping is implemented in reward.rs
        // Expected failure: rewards can be > 1.0 or < -1.0 without clipping
        // After fix: rewards should be clamped to [-1.0, +1.0] in calculate_reward()
        if reward.abs() > 1.0 {
            println!("⚠️  UNBOUNDED REWARD DETECTED: {:.4}", reward);
        }
    }

    // Train a few steps to verify rewards stay bounded during training
    for step in 0..10 {
        let (loss, _grad_norm) = dqn.train_step(None)?;
        println!("Train step {}: loss={:.4}", step, loss);

        // Loss should be reasonable (not exploding)
        assert!(loss < 1000.0, "Loss exploded: {}", loss);
    }

    println!("✅ All rewards can be checked for bounds (test demonstrates unbounded issue)");
    Ok(())
}

#[test]
fn test_q_values_clamped() -> Result<()> {
    println!("\n=== TEST: Q-Values Stay Clamped ===");

    // Recreate Trial 3 conditions (penalty=2.0, high volatility)
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.batch_size = 8;
    config.min_replay_size = 20;
    config.epsilon_start = 0.1; // Low exploration
    config.epsilon_end = 0.01;
    config.epsilon_decay = 0.995;

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

    // Populate replay buffer with high-reward experiences (simulating Q-explosion scenario)
    for i in 0..50 {
        let portfolio_value = 1.0 + (i as f32 * 0.02); // Gradual increase
        let state = create_test_state(portfolio_value);
        let next_state = create_test_state(portfolio_value + 0.01);

        let action = TradingAction::from_int((i % 3) as u8).unwrap();

        // Use rewards that could cause Q-value explosion (±0.5 range)
        let reward = ((i as f32 % 10.0) - 5.0) / 10.0; // Range: [-0.5, +0.4]

        let experience = Experience::new(
            state.clone(),
            action as u8,
            reward,
            next_state.clone(),
            false,
        );

        dqn.store_experience(experience)?;
    }

    // Train for 100 steps and monitor Q-values
    let mut max_q_seen = 0.0_f32;
    let mut min_q_seen = 0.0_f32;

    for step in 0..100 {
        let (loss, _grad_norm) = dqn.train_step(None)?;

        // Get Q-values for a sample state
        let sample_state = create_test_state(1.0);
        let state_tensor = candle_core::Tensor::from_vec(
            sample_state.clone(),
            (1, config.state_dim),
            dqn.device(),
        )?;

        let q_values = dqn.forward(&state_tensor)?;
        let q_vec = q_values.to_vec2::<f32>()?;
        let q_buy = q_vec[0][0];
        let q_sell = q_vec[0][1];
        let q_hold = q_vec[0][2];

        max_q_seen = max_q_seen.max(q_buy).max(q_sell).max(q_hold);
        min_q_seen = min_q_seen.min(q_buy).min(q_sell).min(q_hold);

        if step % 10 == 0 {
            println!(
                "Step {}: Q=[{:.2}, {:.2}, {:.2}], loss={:.4}",
                step, q_buy, q_sell, q_hold, loss
            );
        }

        // Check for Q-value explosion (like +24,055 in Trial 3)
        // NOTE: This test will FAIL until Q-value clamping is implemented
        assert!(
            q_buy.abs() < 1000.0,
            "Q-value explosion detected: Q_BUY={} at step {}",
            q_buy,
            step
        );
        assert!(
            q_sell.abs() < 1000.0,
            "Q-value explosion detected: Q_SELL={} at step {}",
            q_sell,
            step
        );
        assert!(
            q_hold.abs() < 1000.0,
            "Q-value explosion detected: Q_HOLD={} at step {}",
            q_hold,
            step
        );

        // Check for sudden jumps (>100 in magnitude)
        if step > 0 {
            let prev_state_tensor = candle_core::Tensor::from_vec(
                sample_state.clone(),
                (1, config.state_dim),
                dqn.device(),
            )?;
            let prev_q = dqn.forward(&prev_state_tensor)?;
            let prev_q_vec = prev_q.to_vec2::<f32>()?;

            let jump = (q_buy - prev_q_vec[0][0]).abs();
            assert!(
                jump < 100.0,
                "Sudden Q-value jump detected: {} at step {}",
                jump,
                step
            );
        }
    }

    println!("✅ Q-values stayed within [-1000, +1000] bounds");
    println!("   Max Q: {:.2}, Min Q: {:.2}", max_q_seen, min_q_seen);
    Ok(())
}

#[test]
fn test_gradient_norms_reasonable() -> Result<()> {
    println!("\n=== TEST: Gradient Norms Stay Reasonable ===");

    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.batch_size = 8;
    config.min_replay_size = 20;

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

    // Populate replay buffer
    for i in 0..30 {
        let state = create_test_state(1.0);
        let next_state = create_test_state(1.0 + (i as f32 * 0.01));
        let action = TradingAction::from_int((i % 3) as u8).unwrap();
        let reward = ((i as f32 % 10.0) - 5.0) / 20.0; // Range: [-0.25, +0.2]

        let experience = Experience::new(state, action as u8, reward, next_state, false);
        dqn.store_experience(experience)?;
    }

    // Train and monitor gradient norms
    let mut underflow_count = 0;
    let mut overflow_count = 0;

    for step in 0..100 {
        let (loss, grad_norm) = dqn.train_step(None)?;

        if step % 10 == 0 {
            println!(
                "Step {}: grad_norm={:.6}, loss={:.4}",
                step, grad_norm, loss
            );
        }

        // Check for underflow (FP32 threshold ~1e-38, practical threshold 1e-6)
        if grad_norm < 1e-5 {
            underflow_count += 1;
            println!(
                "⚠️  Gradient underflow at step {}: norm={:.2e}",
                step, grad_norm
            );
        }

        // Check for overflow (gradient clipping should prevent this)
        if grad_norm > 100.0 {
            overflow_count += 1;
            println!(
                "⚠️  Gradient overflow at step {}: norm={:.2e}",
                step, grad_norm
            );
        }

        // Assert gradients are in reasonable range
        assert!(
            grad_norm >= 1e-6 && grad_norm <= 100.0,
            "Gradient norm out of range: {} at step {}",
            grad_norm,
            step
        );
    }

    // Allow up to 5% underflow rate (5 out of 100 steps)
    let underflow_rate = (underflow_count as f32 / 100.0) * 100.0;
    println!(
        "Gradient underflow rate: {:.1}% ({}/100 steps)",
        underflow_rate, underflow_count
    );

    assert!(
        underflow_rate < 5.0,
        "Too many gradient underflows: {:.1}% (expected <5%)",
        underflow_rate
    );

    assert_eq!(
        overflow_count, 0,
        "Gradient overflows detected: {}",
        overflow_count
    );

    println!("✅ Gradient norms stayed in reasonable range [1e-6, 100.0]");
    Ok(())
}

#[test]
fn test_no_nan_or_inf_in_training() -> Result<()> {
    println!("\n=== TEST: No NaN or Inf During Training ===");

    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.batch_size = 8;
    config.min_replay_size = 20;

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

    // Populate replay buffer
    for i in 0..30 {
        let state = create_test_state(1.0);
        let next_state = create_test_state(1.0 + (i as f32 * 0.01));
        let action = TradingAction::from_int((i % 3) as u8).unwrap();
        let reward = ((i as f32 % 10.0) - 5.0) / 20.0; // Range: [-0.25, +0.2]

        let experience = Experience::new(
            state.clone(),
            action as u8,
            reward,
            next_state.clone(),
            false,
        );
        dqn.store_experience(experience)?;
    }

    // Train and check for NaN/Inf
    for step in 0..100 {
        let (loss, grad_norm) = dqn.train_step(None)?;

        // Check loss for NaN/Inf
        assert!(!loss.is_nan(), "Loss is NaN at step {}", step);
        assert!(!loss.is_infinite(), "Loss is Inf at step {}", step);

        // Check gradient norm for NaN/Inf
        assert!(!grad_norm.is_nan(), "Gradient norm is NaN at step {}", step);
        assert!(
            !grad_norm.is_infinite(),
            "Gradient norm is Inf at step {}",
            step
        );

        // Check Q-values for NaN/Inf
        let sample_state = create_test_state(1.0);
        let state_tensor =
            candle_core::Tensor::from_vec(sample_state, (1, config.state_dim), dqn.device())?;

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

        for (i, &q) in q_vec[0].iter().enumerate() {
            assert!(!q.is_nan(), "Q-value[{}] is NaN at step {}", i, step);
            assert!(!q.is_infinite(), "Q-value[{}] is Inf at step {}", i, step);
        }

        if step % 20 == 0 {
            println!(
                "Step {}: loss={:.4}, grad_norm={:.4}, Q=[{:.2}, {:.2}, {:.2}]",
                step, loss, grad_norm, q_vec[0][0], q_vec[0][1], q_vec[0][2]
            );
        }
    }

    println!("✅ No NaN or Inf values detected during 100 training steps");
    Ok(())
}

#[test]
fn test_huber_loss_protection() -> Result<()> {
    println!("\n=== TEST: Huber Loss Provides Adequate Protection ===");

    // Test with current delta=1.0 (should fail with large TD errors)
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.batch_size = 8;
    config.min_replay_size = 20;
    config.use_huber_loss = true;
    config.huber_delta = 1.0; // Current value (too small)

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

    // Create high-reward scenario (large TD errors)
    for i in 0..30 {
        let portfolio_value = 1.0 + (i as f32 * 0.05); // 5% growth per step
        let state = create_test_state(portfolio_value);
        let next_state = create_test_state(portfolio_value + 0.1);

        let action = TradingAction::from_int((i % 3) as u8).unwrap();
        let reward = 0.5 + ((i as f32 % 10.0) / 10.0); // Range: [0.5, 1.4]

        let experience = Experience::new(state, action as u8, reward, next_state, false);
        dqn.store_experience(experience)?;
    }

    // Train and monitor loss magnitude
    let mut max_loss = 0.0_f32;

    for step in 0..50 {
        let (loss, _grad_norm) = dqn.train_step(None)?;
        max_loss = max_loss.max(loss);

        if step % 10 == 0 {
            println!("Step {}: loss={:.4}, max_loss={:.4}", step, loss, max_loss);
        }

        // With delta=1.0, loss should stay bounded for moderate TD errors
        // NOTE: This test validates that Huber loss provides *some* protection
        // but may still show instability with extreme TD errors (>10)
        assert!(
            loss < 10000.0,
            "Loss exploded despite Huber loss: {} at step {}",
            loss,
            step
        );
    }

    println!(
        "✅ Huber loss (delta={}) kept loss bounded (max: {:.4})",
        config.huber_delta, max_loss
    );
    println!("   NOTE: Increasing delta to 10.0 recommended for better stability");
    Ok(())
}