foxhunt/ml/tests/bug17_reward_normalization_test.rs

/// Bug #17 P1 Fix: Reward Normalization & Percentage-based P&L Tests
///
/// Tests for the reward normalization system that prevents the positive feedback loop
/// discovered in Bug #17 where amplified rewards were fed back into Sharpe ratio calculation.
///
/// Key improvements tested:
/// 1. RewardNormalizer using Welford's algorithm for online mean/variance
/// 2. Percentage-based P&L (pct_return = (next - current) / current)
/// 3. Defense-in-depth layering (normalize → clamp to [-3, +3])
/// 4. Stationary reward distribution ~N(0,1)

use ml::dqn::reward::{RewardFunction, RewardNormalizer};
use ml::dqn::action_space::{FactoredAction, ExposureLevel, OrderType, Urgency};
use ml::dqn::agent::TradingState;
use ml::dqn::circuit_breaker::CircuitBreakerConfig;
use approx::assert_relative_eq;

#[test]
fn test_reward_normalizer_initialization() {
    // Test that RewardNormalizer initializes with correct default values
    let normalizer = RewardNormalizer::new();

    let (mean, std) = normalizer.get_stats();
    assert_eq!(mean, 0.0, "Initial mean should be 0.0");
    assert_eq!(std, 0.0, "Initial std should be 0.0");
    assert_eq!(normalizer.count(), 0, "Initial count should be 0");
}

#[test]
fn test_welford_algorithm_running_stats() {
    // Test that Welford's algorithm correctly computes running mean/std
    let mut normalizer = RewardNormalizer::new();

    // Add known values: [1.0, 2.0, 3.0, 4.0, 5.0]
    // Expected mean: 3.0, Expected std: sqrt(2.0) ≈ 1.414
    let values = vec![1.0, 2.0, 3.0, 4.0, 5.0];

    for &val in &values {
        normalizer.update(val);
    }

    let (mean, std) = normalizer.get_stats();
    assert_relative_eq!(mean, 3.0, epsilon = 1e-6);
    assert_relative_eq!(std, 1.4142135623730951, epsilon = 1e-6);
    assert_eq!(normalizer.count(), 5);
}

#[test]
fn test_normalization_produces_standard_normal() {
    // Test that normalization produces values with mean ≈ 0, std ≈ 1
    let mut normalizer = RewardNormalizer::new();

    // Simulate a sequence of rewards with known distribution
    // Mean: 10.0, Std: ~3.16
    let raw_rewards: Vec<f64> = vec![
        5.0, 10.0, 15.0, 8.0, 12.0, 7.0, 13.0, 9.0, 11.0, 14.0,
        6.0, 10.0, 15.0, 8.0, 12.0, 7.0, 13.0, 9.0, 11.0, 14.0,
    ];

    // Build up statistics
    for &reward in &raw_rewards {
        normalizer.update(reward);
    }

    // Now normalize a new batch of rewards
    let test_rewards = vec![5.0, 10.0, 15.0];
    let normalized: Vec<f64> = test_rewards
        .iter()
        .map(|&r| normalizer.normalize(r))
        .collect();

    // Check that normalized values are in reasonable range
    for &norm_val in &normalized {
        assert!(
            norm_val.abs() <= 3.0,
            "Normalized value {} should be within [-3, 3]",
            norm_val
        );
    }

    // Check that extreme values get normalized appropriately
    let low_value = 5.0; // ~1 std below mean
    let high_value = 15.0; // ~1 std above mean

    let norm_low = normalizer.normalize(low_value);
    let norm_high = normalizer.normalize(high_value);

    assert!(norm_low < 0.0, "Value below mean should normalize negative");
    assert!(norm_high > 0.0, "Value above mean should normalize positive");
}

#[test]
fn test_percentage_based_pnl_calculation() {
    // Test that percentage-based P&L is calculated correctly
    // pct_return = (next_value - current_value) / current_value

    let current_value = 100_000.0;
    let next_value_up = 102_000.0; // +2% gain
    let next_value_down = 98_000.0; // -2% loss

    let pct_gain = (next_value_up - current_value) / current_value;
    let pct_loss = (next_value_down - current_value) / current_value;

    assert_relative_eq!(pct_gain, 0.02, epsilon = 1e-6);
    assert_relative_eq!(pct_loss, -0.02, epsilon = 1e-6);

    // Test with different portfolio sizes to ensure scale-invariance
    let small_portfolio = 10_000.0;
    let large_portfolio = 1_000_000.0;

    // Both should produce same percentage for same relative change
    let small_next = small_portfolio * 1.02; // +2%
    let large_next = large_portfolio * 1.02; // +2%

    let small_pct = (small_next - small_portfolio) / small_portfolio;
    let large_pct = (large_next - large_portfolio) / large_portfolio;

    assert_relative_eq!(small_pct, large_pct, epsilon = 1e-6);
    assert_relative_eq!(small_pct, 0.02, epsilon = 1e-6);
}

#[test]
fn test_defense_in_depth_clamping() {
    // Test that the defense-in-depth system correctly clamps outliers
    // Even after normalization, values should be clamped to [-3, +3]

    let mut normalizer = RewardNormalizer::new();

    // Build statistics with normal values
    for i in 0..100 {
        normalizer.update(i as f64);
    }

    // Test extreme outlier
    let extreme_value = 1000.0;
    let normalized = normalizer.normalize(extreme_value);

    // Should be clamped after normalization
    let clamped = normalized.clamp(-3.0, 3.0);

    assert!(
        clamped.abs() <= 3.0,
        "Clamped value {} should be within [-3, 3]",
        clamped
    );

    // Test that reasonable values pass through
    let reasonable_value = 50.0;
    let norm_reasonable = normalizer.normalize(reasonable_value);
    let clamped_reasonable = norm_reasonable.clamp(-3.0, 3.0);

    // Should be unchanged by clamping (within bounds)
    assert_relative_eq!(norm_reasonable, clamped_reasonable, epsilon = 1e-6);
}

/// Helper to create a TradingState with specified portfolio value
fn create_state_with_portfolio(portfolio_value: f64) -> TradingState {
    TradingState {
        price_features: vec![100.0, 100.0, 100.0, 100.0], // OHLC
        technical_indicators: vec![0.5; 121],             // 121 technical indicators
        market_features: vec![],                          // Empty
        portfolio_features: vec![portfolio_value as f32, 0.0, 0.001], // [value, position, spread]
    }
}

#[test]
fn test_reward_function_integration_with_normalization() -> Result<(), Box<dyn std::error::Error>> {
    // Integration test: RewardFunction with normalization enabled

    // Create RewardFunction with normalization enabled
    let circuit_breaker_config = CircuitBreakerConfig::default();
    let config = RewardFunction::builder()
        .pnl_weight(1.0)
        .hold_penalty_weight(0.01)
        .activity_bonus_weight(0.0)
        .use_percentage_pnl(true) // Enable percentage-based P&L
        .enable_normalization(true) // Enable normalization
        .circuit_breaker_config(circuit_breaker_config)
        .build()?;

    let mut reward_fn = RewardFunction::new(config);

    // Simulate a sequence of portfolio value changes
    let portfolio_values = vec![
        100_000.0, // Initial
        102_000.0, // +2%
        101_000.0, // -0.98%
        103_000.0, // +1.98%
        100_000.0, // -2.91%
        105_000.0, // +5%
    ];

    let mut rewards = Vec::new();
    let action = FactoredAction {
        exposure: ExposureLevel::Long50,
        order: OrderType::Market,
        urgency: Urgency::Normal,
    };
    let recent_actions = vec![action; 10]; // Dummy recent actions

    for i in 0..portfolio_values.len() - 1 {
        let current_state = create_state_with_portfolio(portfolio_values[i]);
        let next_state = create_state_with_portfolio(portfolio_values[i + 1]);

        // Calculate reward (which should use percentage-based P&L internally)
        let reward = reward_fn.calculate_reward(
            action,
            &current_state,
            &next_state,
            &recent_actions,
        )?;

        // Convert Decimal to f64 for comparison
        let reward_f64: f64 = reward.try_into()?;
        rewards.push(reward_f64);
    }

    // All rewards should be within reasonable bounds after normalization
    for (i, &reward) in rewards.iter().enumerate() {
        assert!(
            reward.abs() <= 3.0,
            "Reward {} at step {} should be clamped to [-3, 3]",
            reward,
            i
        );
    }

    // Rewards should not all be identical (diversity check)
    let first_reward = rewards[0];
    let all_same = rewards.iter().all(|&r| (r - first_reward).abs() < 1e-6);
    assert!(
        !all_same,
        "Rewards should vary based on portfolio performance"
    );

    Ok(())
}

#[test]
fn test_normalization_disabled_backward_compatibility() -> Result<(), Box<dyn std::error::Error>> {
    // Test that normalization can be disabled for backward compatibility

    let circuit_breaker_config = CircuitBreakerConfig::default();
    let config = RewardFunction::builder()
        .pnl_weight(1.0)
        .hold_penalty_weight(0.01)
        .activity_bonus_weight(0.0)
        .use_percentage_pnl(false) // Use absolute P&L
        .enable_normalization(false) // Disable normalization
        .circuit_breaker_config(circuit_breaker_config)
        .build()?;

    let mut reward_fn = RewardFunction::new(config);

    let action = FactoredAction {
        exposure: ExposureLevel::Long50,
        order: OrderType::Market,
        urgency: Urgency::Normal,
    };
    let recent_actions = vec![action; 10];

    let current_state = create_state_with_portfolio(100_000.0);
    let next_state = create_state_with_portfolio(102_000.0);

    // Calculate a reward
    let reward = reward_fn.calculate_reward(
        action,
        &current_state,
        &next_state,
        &recent_actions,
    )?;

    // Convert to f64
    let reward_f64: f64 = reward.try_into()?;

    // Without normalization, reward should still be clamped to [-1, 1]
    assert!(
        reward_f64.abs() <= 1.0,
        "Without normalization, reward should be clamped to [-1, 1]"
    );

    Ok(())
}

#[test]
fn test_normalizer_handles_edge_cases() {
    // Test edge cases: single value, identical values, zero std

    let mut normalizer = RewardNormalizer::new();

    // Edge case 1: First value should return unchanged (count < 2)
    normalizer.update(5.0);
    let norm1 = normalizer.normalize(5.0);
    assert_eq!(norm1, 5.0, "First value should return unchanged");

    // Edge case 2: Identical values (zero std)
    let mut norm_identical = RewardNormalizer::new();
    for _ in 0..10 {
        norm_identical.update(10.0);
    }

    let norm_val = norm_identical.normalize(10.0);
    assert_eq!(norm_val, 10.0, "Zero std should return value unchanged");

    // Edge case 3: Very small std (near epsilon)
    let mut norm_small_std = RewardNormalizer::new();
    for i in 0..100 {
        // Values very close together
        norm_small_std.update(10.0 + (i as f64) * 0.0001);
    }

    let (mean, std) = norm_small_std.get_stats();
    assert!(mean > 0.0, "Mean should be computed");
    assert!(std >= 0.0, "Std should be non-negative");
}