foxhunt/ml/tests/dqn_reward_calculation_test.rs

//! Comprehensive DQN Reward Function Tests
//!
//! This test suite validates the DQN reward calculation function to prevent regression
//! and ensure correct behavior across various market scenarios.
//!
//! Test Coverage:
//! 1. Price increase scenarios (positive rewards)
//! 2. Price decrease scenarios (negative rewards)
//! 3. Flat market (zero reward)
//! 4. Large price movements (reward clamping)
//! 5. Reward normalization (proportional scaling)
//! 6. Non-constant rewards (variance validation)
//! 7. Symmetry (balanced positive/negative rewards)

use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

/// Helper function to create a default DQN trainer for testing
fn create_test_trainer() -> DQNTrainer {
    let hyperparams = DQNHyperparameters::conservative();
    DQNTrainer::new(hyperparams).expect("Failed to create test trainer")
}

/// Helper function to calculate reward from price change
///
/// This implements the expected reward calculation logic:
/// reward = (next_price - current_price) / normalization_factor
/// where normalization_factor = 10.0 (to normalize typical ES futures movements)
fn calculate_price_reward(current: f64, next: f64) -> f32 {
    let price_change = next - current;
    // Normalize by 10.0 (typical ES futures tick movements)
    // Clamp to [-1.0, 1.0] range
    (price_change / 10.0).clamp(-1.0, 1.0) as f32
}

/// Test 1: Reward for price increase
///
/// Validates that price increases generate positive rewards
/// Example: ES futures rising from 5900.0 to 5914.25 should yield positive reward
#[test]
fn test_reward_price_increase() {
    let current = 5900.0;
    let next = 5914.25;

    let reward = calculate_price_reward(current, next);

    // Verify reward is positive
    assert!(
        reward > 0.0,
        "Reward should be positive for price increase. Got: {}",
        reward
    );

    // Verify reward is within valid range [0.0, 1.0]
    assert!(
        reward <= 1.0,
        "Reward should be clamped to 1.0. Got: {}",
        reward
    );

    assert!(
        reward >= 0.0,
        "Reward should be non-negative for price increase. Got: {}",
        reward
    );

    // Expected: (5914.25 - 5900.0) / 10.0 = 1.425, clamped to 1.0
    let expected = 1.0;
    assert!(
        (reward - expected).abs() < 0.01,
        "Expected reward ~{}, got {}",
        expected,
        reward
    );
}

/// Test 2: Reward for price decrease
///
/// Validates that price decreases generate negative rewards
/// Example: ES futures falling from 5914.25 to 5900.0 should yield negative reward
#[test]
fn test_reward_price_decrease() {
    let current = 5914.25;
    let next = 5900.0;

    let reward = calculate_price_reward(current, next);

    // Verify reward is negative
    assert!(
        reward < 0.0,
        "Reward should be negative for price decrease. Got: {}",
        reward
    );

    // Verify reward is within valid range [-1.0, 0.0]
    assert!(
        reward >= -1.0,
        "Reward should be clamped to -1.0. Got: {}",
        reward
    );

    assert!(
        reward <= 0.0,
        "Reward should be non-positive for price decrease. Got: {}",
        reward
    );

    // Expected: (5900.0 - 5914.25) / 10.0 = -1.425, clamped to -1.0
    let expected = -1.0;
    assert!(
        (reward - expected).abs() < 0.01,
        "Expected reward ~{}, got {}",
        expected,
        reward
    );
}

/// Test 3: Reward for flat market
///
/// Validates that no price change generates near-zero reward
/// Example: Price staying at 5900.0 should yield ~0.0 reward
#[test]
fn test_reward_flat_market() {
    let current = 5900.0;
    let next = 5900.0;

    let reward = calculate_price_reward(current, next);

    // Verify reward is approximately zero (within tolerance)
    assert!(
        reward.abs() < 0.01,
        "Reward should be near zero for flat market. Got: {}",
        reward
    );

    // Expected: (5900.0 - 5900.0) / 10.0 = 0.0
    assert_eq!(reward, 0.0, "Flat market should produce exactly 0.0 reward");
}

/// Test 4: Reward clamping for large price movements
///
/// Validates that extremely large price movements are clamped to [-1.0, 1.0]
/// This prevents unbounded rewards from distorting training
#[test]
fn test_reward_large_move() {
    // Large upward movement (50 points)
    let current_up = 5900.0;
    let next_up = 5950.0;
    let reward_up = calculate_price_reward(current_up, next_up);

    // Should clamp to 1.0
    assert!(
        (reward_up - 1.0).abs() < 0.01,
        "Large upward move should clamp to 1.0. Got: {}",
        reward_up
    );

    // Large downward movement (50 points)
    let current_down = 5950.0;
    let next_down = 5900.0;
    let reward_down = calculate_price_reward(current_down, next_down);

    // Should clamp to -1.0
    assert!(
        (reward_down - (-1.0)).abs() < 0.01,
        "Large downward move should clamp to -1.0. Got: {}",
        reward_down
    );

    // Verify symmetry
    assert!(
        (reward_up + reward_down).abs() < 0.01,
        "Large moves should be symmetric: up={}, down={}",
        reward_up,
        reward_down
    );
}

/// Test 5: Reward normalization and proportionality
///
/// Validates that rewards scale proportionally with price movements
/// Ensures reward function is not constant or saturated
#[test]
fn test_reward_normalization() {
    let current = 5900.0;

    // 10-point move (should be 1.0 after clamping)
    let next_10 = current + 10.0;
    let reward_10 = calculate_price_reward(current, next_10);
    assert!(
        (reward_10 - 1.0).abs() < 0.01,
        "10-point move: expected ~1.0, got {}",
        reward_10
    );

    // 5-point move (should be 0.5)
    let next_5 = current + 5.0;
    let reward_5 = calculate_price_reward(current, next_5);
    assert!(
        (reward_5 - 0.5).abs() < 0.01,
        "5-point move: expected ~0.5, got {}",
        reward_5
    );

    // 1-point move (should be 0.1)
    let next_1 = current + 1.0;
    let reward_1 = calculate_price_reward(current, next_1);
    assert!(
        (reward_1 - 0.1).abs() < 0.01,
        "1-point move: expected ~0.1, got {}",
        reward_1
    );

    // Verify proportionality: reward_10 ≈ 2 * reward_5 ≈ 10 * reward_1
    assert!(
        (reward_10 / reward_5 - 2.0).abs() < 0.1,
        "Proportionality check 10/5 failed: {} / {} ≠ 2.0",
        reward_10,
        reward_5
    );

    assert!(
        (reward_10 / reward_1 - 10.0).abs() < 0.2,
        "Proportionality check 10/1 failed: {} / {} ≠ 10.0",
        reward_10,
        reward_1
    );
}

/// Test 6: Reward variance (non-constant)
///
/// Validates that rewards vary across different price pairs
/// This is critical - if rewards are constant, the DQN cannot learn
#[test]
fn test_reward_not_constant() {
    let mut rewards = Vec::with_capacity(100);

    // Test 100 different price pairs with varied movements
    for i in 0..100 {
        let current = 5900.0 + (i as f64 * 0.5);
        let next = current + ((i % 20) as f64 - 10.0); // Oscillating movements
        let reward = calculate_price_reward(current, next);
        rewards.push(reward);
    }

    // Calculate statistics
    let mean = rewards.iter().sum::<f32>() / rewards.len() as f32;
    let variance = rewards.iter().map(|r| (r - mean).powi(2)).sum::<f32>() / rewards.len() as f32;
    let std = variance.sqrt();

    // Verify standard deviation is significant
    assert!(
        std > 0.1,
        "Reward std deviation too low ({:.6}), indicates constant rewards",
        std
    );

    // Verify min ≠ max (rewards actually vary)
    let min_reward = rewards.iter().cloned().fold(f32::INFINITY, f32::min);
    let max_reward = rewards.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

    assert_ne!(
        min_reward, max_reward,
        "Min and max rewards are equal ({} == {}), indicates constant rewards",
        min_reward, max_reward
    );

    // Verify we have both positive and negative rewards
    let has_positive = rewards.iter().any(|&r| r > 0.0);
    let has_negative = rewards.iter().any(|&r| r < 0.0);

    assert!(
        has_positive && has_negative,
        "Should have both positive and negative rewards (pos: {}, neg: {})",
        has_positive,
        has_negative
    );

    println!(
        "Reward statistics: mean={:.4}, std={:.4}, min={:.4}, max={:.4}",
        mean, std, min_reward, max_reward
    );
}

/// Test 7: Reward symmetry
///
/// Validates that equal magnitude movements in opposite directions
/// produce equal magnitude rewards with opposite signs
#[test]
fn test_reward_symmetry() {
    let current = 5900.0;

    // +10 point move
    let next_up = current + 10.0;
    let reward_up = calculate_price_reward(current, next_up);

    // -10 point move
    let next_down = current - 10.0;
    let reward_down = calculate_price_reward(current, next_down);

    // Verify rewards are opposite (symmetric)
    assert!(
        (reward_up + reward_down).abs() < 0.01,
        "Symmetric moves should produce opposite rewards: up={}, down={}",
        reward_up,
        reward_down
    );

    assert!(
        (reward_up - (-reward_down)).abs() < 0.01,
        "Reward magnitudes should be equal: |{}| ≠ |{}|",
        reward_up,
        reward_down
    );

    // Test multiple magnitudes
    for magnitude in &[1.0, 2.5, 5.0, 7.5, 10.0, 15.0, 20.0] {
        let up = calculate_price_reward(current, current + magnitude);
        let down = calculate_price_reward(current, current - magnitude);

        assert!(
            (up + down).abs() < 0.01,
            "Symmetry broken at magnitude {}: up={}, down={}",
            magnitude,
            up,
            down
        );
    }
}

/// Test 8: Edge case - very small price movements
///
/// Validates handling of sub-tick movements (< 0.01 points)
#[test]
fn test_reward_small_movements() {
    let current = 5900.0;

    // Sub-tick movement (0.001 points)
    let next = current + 0.001;
    let reward = calculate_price_reward(current, next);

    // Should be very small but non-zero
    assert!(
        reward > 0.0,
        "Small positive movement should yield positive reward"
    );

    assert!(
        reward < 0.001,
        "Small movement should yield small reward, got {}",
        reward
    );

    // Expected: 0.001 / 10.0 = 0.0001
    let expected = 0.0001;
    assert!(
        (reward - expected).abs() < 0.00001,
        "Expected ~{}, got {}",
        expected,
        reward
    );
}

/// Test 9: Edge case - exact boundary conditions
///
/// Validates reward values at exact clamping boundaries
#[test]
fn test_reward_boundary_conditions() {
    let current = 5900.0;

    // Exactly at upper boundary (+10 points)
    let next_upper = current + 10.0;
    let reward_upper = calculate_price_reward(current, next_upper);
    assert_eq!(
        reward_upper, 1.0,
        "Upper boundary should be exactly 1.0, got {}",
        reward_upper
    );

    // Exactly at lower boundary (-10 points)
    let next_lower = current - 10.0;
    let reward_lower = calculate_price_reward(current, next_lower);
    assert_eq!(
        reward_lower, -1.0,
        "Lower boundary should be exactly -1.0, got {}",
        reward_lower
    );

    // Just below upper boundary (+9.99 points)
    let next_below_upper = current + 9.99;
    let reward_below_upper = calculate_price_reward(current, next_below_upper);
    assert!(
        reward_below_upper < 1.0 && reward_below_upper > 0.99,
        "Just below upper boundary should be < 1.0, got {}",
        reward_below_upper
    );
}

/// Test 10: Integration test - reward calculation in trainer context
///
/// Validates that the DQN trainer can be instantiated and used
/// (actual reward calculation requires full training context)
#[test]
fn test_trainer_integration() {
    let trainer = create_test_trainer();

    // Verify trainer was created successfully
    // (actual reward calculation is tested through the helper function above)
    drop(trainer);
}

/// Test 11: Reward distribution analysis
///
/// Statistical analysis of reward distribution across realistic price scenarios
#[test]
fn test_reward_distribution() {
    let mut rewards = Vec::new();

    // Simulate realistic ES futures price movements
    // Typical daily range: 20-50 points, with most movements < 10 points
    let scenarios = vec![
        (5900.0, 5900.25), // +0.25 tick
        (5900.0, 5900.50), // +0.50 tick
        (5900.0, 5901.0),  // +1 point
        (5900.0, 5902.0),  // +2 points
        (5900.0, 5905.0),  // +5 points
        (5900.0, 5910.0),  // +10 points (boundary)
        (5900.0, 5915.0),  // +15 points (clamped)
        (5900.0, 5899.75), // -0.25 tick
        (5900.0, 5899.50), // -0.50 tick
        (5900.0, 5899.0),  // -1 point
        (5900.0, 5898.0),  // -2 points
        (5900.0, 5895.0),  // -5 points
        (5900.0, 5890.0),  // -10 points (boundary)
        (5900.0, 5885.0),  // -15 points (clamped)
    ];

    for (current, next) in scenarios {
        let reward = calculate_price_reward(current, next);
        rewards.push(reward);
    }

    // Verify we have a good distribution
    let positive_count = rewards.iter().filter(|&&r| r > 0.0).count();
    let negative_count = rewards.iter().filter(|&&r| r < 0.0).count();
    let zero_count = rewards.iter().filter(|&&r| r == 0.0).count();

    assert_eq!(
        positive_count, 7,
        "Expected 7 positive rewards, got {}",
        positive_count
    );

    assert_eq!(
        negative_count, 7,
        "Expected 7 negative rewards, got {}",
        negative_count
    );

    assert_eq!(
        zero_count, 0,
        "Expected 0 zero rewards (no flat scenarios), got {}",
        zero_count
    );
}

/// Test 12: Precision validation
///
/// Ensures reward calculation maintains precision for small values
#[test]
fn test_reward_precision() {
    let current = 5900.0;

    // Test incrementally small movements
    let test_cases = vec![
        (0.01, 0.001), // 0.01 point move
        (0.05, 0.005), // 0.05 point move
        (0.10, 0.010), // 0.10 point move
        (0.25, 0.025), // 0.25 point move
        (0.50, 0.050), // 0.50 point move
    ];

    for (movement, expected_reward) in test_cases {
        let next = current + movement;
        let reward = calculate_price_reward(current, next);

        assert!(
            (reward - expected_reward).abs() < 0.0001,
            "Precision error for {}-point move: expected {}, got {}",
            movement,
            expected_reward,
            reward
        );
    }
}