foxhunt/ml/tests/risk_adjusted_reward_test.rs

//! TDD Tests for Risk-Adjusted Reward Calculation
//!
//! This test suite validates risk-adjusted reward calculation based on Sharpe ratio.
//!
//! Risk-adjusted reward formula:
//! - If pnl_history.len() < 20: return raw pnl_change
//! - Otherwise: sharpe_ratio = mean(pnl) / std(pnl)
//!            risk_adjusted_reward = sharpe_ratio * pnl_change
//!
//! Test Coverage:
//! 1. Stable returns (low volatility) → higher reward multiplier
//! 2. Volatile returns (high volatility) → lower multiplier
//! 3. Insufficient history (<20 samples) → raw pnl_change returned
//! 4. Positive PnL + positive Sharpe → amplified reward
//! 5. Positive PnL + negative Sharpe → penalized reward
//! 6. Zero volatility edge case → raw pnl_change returned
//! 7. Rolling window (last 20 samples) used for calculation
//! 8. Outlier handling → extreme returns don't break calculation
//! 9. Logging validation → Sharpe multiplier logged every 100 steps
//! 10. Comparison test → risk-adjusted > raw in stable periods

use std::collections::VecDeque;

/// Helper struct to simulate PnL history and calculate risk-adjusted rewards
#[derive(Debug, Clone)]
struct RiskAdjustmentCalculator {
    pnl_history: VecDeque<f64>,
    max_history: usize,
    step_count: u32,
}

impl RiskAdjustmentCalculator {
    fn new(max_history: usize) -> Self {
        Self {
            pnl_history: VecDeque::with_capacity(max_history),
            max_history,
            step_count: 0,
        }
    }

    fn add_pnl(&mut self, pnl: f64) {
        if self.pnl_history.len() >= self.max_history {
            self.pnl_history.pop_front();
        }
        self.pnl_history.push_back(pnl);
    }

    fn calculate_risk_adjusted_reward(&self, pnl_change: f64) -> f64 {
        // Insufficient history: return raw pnl_change
        if self.pnl_history.len() < 20 {
            return pnl_change;
        }

        // Calculate mean
        let mean = self.pnl_history.iter().sum::<f64>() / self.pnl_history.len() as f64;

        // Calculate standard deviation
        let variance = self.pnl_history.iter()
            .map(|&x| (x - mean).powi(2))
            .sum::<f64>() / self.pnl_history.len() as f64;
        let std_dev = variance.sqrt();

        // Zero volatility edge case: return raw pnl_change
        if std_dev < 1e-8 {
            return pnl_change;
        }

        // Calculate Sharpe ratio
        let sharpe = mean / std_dev;

        // Return risk-adjusted reward
        sharpe * pnl_change
    }

    fn get_sharpe_ratio(&self) -> Option<f64> {
        if self.pnl_history.len() < 20 {
            return None;
        }

        let mean = self.pnl_history.iter().sum::<f64>() / self.pnl_history.len() as f64;
        let variance = self.pnl_history.iter()
            .map(|&x| (x - mean).powi(2))
            .sum::<f64>() / self.pnl_history.len() as f64;
        let std_dev = variance.sqrt();

        if std_dev < 1e-8 {
            // Zero variance case: return None (Sharpe is undefined)
            return None;
        }

        Some(mean / std_dev)
    }

    fn increment_step(&mut self) {
        self.step_count += 1;
    }

    fn should_log_sharpe(&self) -> bool {
        self.step_count % 100 == 0 && self.step_count > 0
    }

    fn history_len(&self) -> usize {
        self.pnl_history.len()
    }
}

// ============================================================================
// TEST 1: Stable Returns (Low Volatility) → Higher Reward Multiplier
// ============================================================================

/// Test that stable returns (low volatility) produce a higher Sharpe multiplier
/// and thus amplify the reward signal.
///
/// Scenario: PnL changes of +0.5 each, very consistent, low std dev
/// Expected: Sharpe > 1.0, so reward > pnl_change
#[test]
fn test_sharpe_reward_with_stable_returns() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add 20 stable PnL samples with slight variation (not identical)
    // Values oscillate slightly around 0.5 to give low (but non-zero) variance
    for i in 0..20 {
        let noise = if i % 2 == 0 { 0.01 } else { -0.01 };
        calc.add_pnl(0.5 + noise);
    }

    assert_eq!(calc.history_len(), 20, "Should have 20 samples");

    // Now calculate reward for +1.0 PnL change
    let pnl_change = 1.0;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);
    let sharpe = calc.get_sharpe_ratio().unwrap();

    // Verify Sharpe calculation
    assert!(sharpe > 0.0, "Sharpe should be positive for positive stable returns");
    assert!(sharpe > 1.0, "Sharpe should be > 1.0 for very stable returns (low std)");

    // Verify reward amplification
    // mean ≈ 0.5, std ≈ 0.01, sharpe = 0.5/0.01 = 50 (very large)
    // reward = sharpe * pnl_change = 50 * 1.0 = 50 >> pnl_change
    println!("Stable returns - Sharpe: {}, Reward: {}, Raw PnL: {}", sharpe, reward, pnl_change);
}

// ============================================================================
// TEST 2: Volatile Returns (High Volatility) → Lower Multiplier
// ============================================================================

/// Test that volatile returns (high volatility) produce a lower Sharpe multiplier
/// and thus reduce the reward signal.
///
/// Scenario: PnL changes oscillate between -1.0 and +1.0, high std dev
/// Expected: Sharpe < 1.0, so reward < pnl_change (for positive PnL)
#[test]
fn test_sharpe_reward_with_volatile_returns() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add 20 volatile PnL samples (oscillating)
    for i in 0..20 {
        let pnl = if i % 2 == 0 { 1.0 } else { -1.0 };
        calc.add_pnl(pnl);
    }

    assert_eq!(calc.history_len(), 20, "Should have 20 samples");

    // Calculate Sharpe ratio (mean should be ~0, std should be ~1.0)
    let sharpe = calc.get_sharpe_ratio();
    assert!(sharpe.is_some(), "Should have valid Sharpe");

    let sharpe_val = sharpe.unwrap();
    // With mean=0 and std=1, Sharpe = 0
    assert!(sharpe_val.abs() < 0.1, "Sharpe should be near 0 for oscillating returns");

    // Calculate risk-adjusted reward for +1.0 PnL change
    let pnl_change = 1.0;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);

    // reward should be approximately 0 (Sharpe ≈ 0)
    assert!(reward.abs() < 0.2, "Reward should be reduced by low Sharpe");

    println!(
        "Volatile returns - Sharpe: {}, Reward: {}, Raw PnL: {}",
        sharpe_val, reward, pnl_change
    );
}

// ============================================================================
// TEST 3: Insufficient History (<20 Samples) → Raw PnL Returned
// ============================================================================

/// Test that with insufficient history (< 20 samples), raw pnl_change is returned.
///
/// Scenario: Only 15 PnL samples accumulated
/// Expected: calculate_risk_adjusted_reward returns raw pnl_change unchanged
#[test]
fn test_sharpe_reward_insufficient_history() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add only 15 samples
    for i in 0..15 {
        calc.add_pnl((i as f64) * 0.1);
    }

    assert_eq!(calc.history_len(), 15, "Should have exactly 15 samples");

    // Sharpe should not be available
    let sharpe = calc.get_sharpe_ratio();
    assert!(sharpe.is_none(), "Sharpe should not be available with < 20 samples");

    // Calculate risk-adjusted reward
    let pnl_change = 2.5;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);

    // Should return raw pnl_change
    assert_eq!(
        reward, pnl_change,
        "With insufficient history, reward should equal raw pnl_change"
    );
}

// ============================================================================
// TEST 4: Positive PnL + Positive Sharpe → Amplified Reward
// ============================================================================

/// Test that positive PnL combined with positive Sharpe ratio produces an amplified reward.
///
/// Scenario: Consistent positive returns (mean = +0.5, low std)
/// Expected: Sharpe > 1.0, reward > pnl_change
#[test]
fn test_sharpe_reward_positive_pnl_positive_sharpe() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add 20 samples: positive with low volatility
    // Series: 0.4, 0.5, 0.6, 0.5, 0.4, 0.5, ... (oscillates slightly around 0.5)
    for i in 0..20 {
        let base = 0.5;
        let noise = (i as f64 % 3.0 - 1.0) * 0.05;
        calc.add_pnl(base + noise);
    }

    assert_eq!(calc.history_len(), 20, "Should have 20 samples");

    let sharpe = calc.get_sharpe_ratio().unwrap();
    assert!(sharpe > 0.0, "Sharpe should be positive");

    // Positive PnL change
    let pnl_change = 0.8;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);

    // With positive Sharpe > 1, reward should be > pnl_change
    if sharpe > 1.0 {
        assert!(
            reward > pnl_change,
            "Positive Sharpe > 1.0 should amplify reward: {} > {}",
            reward,
            pnl_change
        );
    }

    println!(
        "Positive PnL + Positive Sharpe - Sharpe: {}, Reward: {}, Raw PnL: {}",
        sharpe, reward, pnl_change
    );
}

// ============================================================================
// TEST 5: Positive PnL + Negative Sharpe → Penalized Reward
// ============================================================================

/// Test that positive PnL combined with negative Sharpe ratio produces a penalized (negative) reward.
///
/// Scenario: Mean negative returns with occasional large gains
/// Expected: Sharpe < 0, even positive PnL becomes negative after adjustment
#[test]
fn test_sharpe_reward_positive_pnl_negative_sharpe() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add 20 samples: negative mean with one outlier
    // Series: -0.3, -0.3, -0.3, ..., -0.3 (19 times), then +5.0 (1 time)
    for i in 0..20 {
        if i == 19 {
            calc.add_pnl(5.0);
        } else {
            calc.add_pnl(-0.3);
        }
    }

    assert_eq!(calc.history_len(), 20, "Should have 20 samples");

    let sharpe = calc.get_sharpe_ratio().unwrap();
    // mean = (-0.3 * 19 + 5.0) / 20 = 4.3 / 20 = 0.215
    // This is actually positive! Let's use different values

    // Recreate with negative mean
    let mut calc2 = RiskAdjustmentCalculator::new(50);
    for i in 0..20 {
        if i == 0 {
            calc2.add_pnl(3.0);
        } else {
            calc2.add_pnl(-0.5);
        }
    }

    let sharpe2 = calc2.get_sharpe_ratio().unwrap();
    // mean = (3.0 - 0.5*19) / 20 = (3.0 - 9.5) / 20 = -6.5/20 = -0.325 (negative!)

    assert!(
        sharpe2 < 0.0,
        "Sharpe should be negative when mean returns are negative"
    );

    // Now test: positive pnl_change but negative sharpe
    let pnl_change = 0.5;
    let reward = calc2.calculate_risk_adjusted_reward(pnl_change);

    // reward should be negative: negative_sharpe * positive_pnl = negative_reward
    assert!(
        reward < 0.0,
        "Positive PnL with negative Sharpe should produce negative reward: {}",
        reward
    );

    println!(
        "Positive PnL + Negative Sharpe - Sharpe: {}, Reward: {}, Raw PnL: {}",
        sharpe2, reward, pnl_change
    );
}

// ============================================================================
// TEST 6: Zero Volatility Edge Case → Raw PnL Returned
// ============================================================================

/// Test that zero volatility (all identical returns) returns raw pnl_change.
///
/// Scenario: PnL history is constant [0.5, 0.5, 0.5, ..., 0.5]
/// Expected: std = 0, sharpe undefined, returns raw pnl_change
#[test]
fn test_sharpe_reward_zero_volatility() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add 20 identical samples (zero volatility)
    for _ in 0..20 {
        calc.add_pnl(0.5);
    }

    assert_eq!(calc.history_len(), 20, "Should have 20 samples");

    // Sharpe should be undefined (None) due to zero std dev
    let sharpe = calc.get_sharpe_ratio();
    assert!(sharpe.is_none(), "Sharpe should be None for zero volatility");

    // Calculate risk-adjusted reward
    let pnl_change = 1.5;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);

    // Should return raw pnl_change due to zero std dev
    assert_eq!(
        reward, pnl_change,
        "Zero volatility should return raw pnl_change unchanged"
    );
}

// ============================================================================
// TEST 7: Rolling Window (Last 20 Samples) Used for Calculation
// ============================================================================

/// Test that only the last 20 samples are used for Sharpe calculation.
///
/// Scenario: Add 30 samples, verify that first 10 are not used
/// Expected: Sharpe calculated from samples 10-29, not 0-19
#[test]
fn test_sharpe_reward_history_rolling_window() {
    let mut calc = RiskAdjustmentCalculator::new(30);

    // Add 10 early samples with very high values
    for _ in 0..10 {
        calc.add_pnl(100.0);
    }

    // Add 20 normal samples with value 0.5
    for _ in 0..20 {
        calc.add_pnl(0.5);
    }

    assert_eq!(calc.history_len(), 30, "Should have 30 samples");

    // The rolling window should only use the last 20
    let sharpe = calc.get_sharpe_ratio().unwrap();

    // If first 10 were included: mean ≈ 5.5 (very high)
    // If only last 20 used: mean ≈ 0.5
    // Let's verify the second case
    let expected_mean_if_rolling = 0.5;
    let tolerance = 0.01;

    // Recalculate to verify
    let sum_last_20: f64 = calc.pnl_history.iter().skip(10).take(20).sum();
    let mean_of_last_20 = sum_last_20 / 20.0;

    assert!(
        (mean_of_last_20 - expected_mean_if_rolling).abs() < tolerance,
        "Last 20 samples should have mean ~0.5, got {}",
        mean_of_last_20
    );

    println!(
        "Rolling window test - Sharpe: {}, Mean of last 20: {}",
        sharpe, mean_of_last_20
    );
}

// ============================================================================
// TEST 8: Outlier Handling → Extreme Returns Don't Break Calculation
// ============================================================================

/// Test that outliers (extreme returns) don't cause NaN/Inf in calculation.
///
/// Scenario: 19 normal samples + 1 extreme outlier
/// Expected: Calculation completes without NaN/Inf, produces valid result
#[test]
fn test_sharpe_reward_outlier_handling() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add 19 normal samples
    for _ in 0..19 {
        calc.add_pnl(0.1);
    }

    // Add 1 extreme outlier
    calc.add_pnl(1000.0);

    assert_eq!(calc.history_len(), 20, "Should have 20 samples");

    let sharpe = calc.get_sharpe_ratio();
    assert!(sharpe.is_some(), "Sharpe should be valid despite outlier");

    let sharpe_val = sharpe.unwrap();
    assert!(sharpe_val.is_finite(), "Sharpe should be finite, not NaN/Inf");
    assert!(sharpe_val > 0.0, "Sharpe should be positive (mean > 0)");

    // Calculate reward
    let pnl_change = 10.0;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);

    assert!(reward.is_finite(), "Reward should be finite, not NaN/Inf");

    println!(
        "Outlier handling - Sharpe: {}, Reward: {}, Raw PnL: {}",
        sharpe_val, reward, pnl_change
    );
}

// ============================================================================
// TEST 9: Sharpe Multiplier Logging (Every 100 Steps)
// ============================================================================

/// Test that Sharpe multiplier is logged at the correct frequency (every 100 steps).
///
/// Scenario: Increment step counter and check logging flag
/// Expected: Logging triggers at steps 100, 200, 300, etc.
#[test]
fn test_sharpe_reward_logging() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add 20 samples to enable Sharpe calculation
    for _ in 0..20 {
        calc.add_pnl(0.5);
    }

    // Test logging at various step counts
    let mut logged_at = Vec::new();

    for step in 1..=350 {
        calc.step_count = step;
        if calc.should_log_sharpe() {
            logged_at.push(step);
        }
    }

    // Verify logging happens at 100, 200, 300
    assert_eq!(
        logged_at,
        vec![100, 200, 300],
        "Sharpe should be logged at steps 100, 200, 300"
    );

    // Verify specific steps don't log
    calc.step_count = 99;
    assert!(!calc.should_log_sharpe(), "Step 99 should not trigger logging");

    calc.step_count = 100;
    assert!(calc.should_log_sharpe(), "Step 100 should trigger logging");

    calc.step_count = 101;
    assert!(!calc.should_log_sharpe(), "Step 101 should not trigger logging");
}

// ============================================================================
// TEST 10: Comparison Test → Risk-Adjusted > Raw in Stable Periods
// ============================================================================

/// Test that risk-adjusted rewards are higher than raw rewards in stable periods.
///
/// Scenario: Simulate a stable trading period with consistent positive returns
/// Expected: risk_adjusted_reward > pnl_change for all trades in this period
#[test]
fn test_sharpe_reward_comparison_to_baseline() {
    let mut calc = RiskAdjustmentCalculator::new(100);

    // Phase 1: Build up history with stable +0.5 returns (low volatility)
    for i in 0..20 {
        let noise = if i % 2 == 0 { 0.01 } else { -0.01 };
        calc.add_pnl(0.5 + noise);
    }

    // Phase 2: Test several trades during stable period
    let test_trades = vec![0.1, 0.2, 0.5, 0.8, 1.0];
    let mut amplified_count = 0;

    for &pnl_change in &test_trades {
        let raw_reward = pnl_change;
        let adjusted_reward = calc.calculate_risk_adjusted_reward(pnl_change);

        // In stable period (low volatility), high Sharpe > 1, so adjusted > raw
        let sharpe = calc.get_sharpe_ratio().unwrap();
        if sharpe > 1.0 {
            if adjusted_reward > raw_reward {
                amplified_count += 1;
            }
        }

        println!(
            "Trade: raw={}, adjusted={}, sharpe={}, ratio={}",
            raw_reward,
            adjusted_reward,
            sharpe,
            adjusted_reward / raw_reward
        );

        // Add this trade to history
        calc.add_pnl(pnl_change);
    }

    // Verify at least some trades were amplified
    assert!(
        amplified_count > 0,
        "Risk-adjusted rewards should amplify at least some trades in stable periods"
    );
}

// ============================================================================
// BONUS TESTS (Beyond Requirements)
// ============================================================================

/// Test negative PnL with positive Sharpe → negative reward (still penalized)
#[test]
fn test_sharpe_reward_negative_pnl_positive_sharpe() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add 20 positive returns with slight variation (Sharpe > 0)
    for i in 0..20 {
        let noise = if i % 2 == 0 { 0.01 } else { -0.01 };
        calc.add_pnl(0.6 + noise);
    }

    let sharpe = calc.get_sharpe_ratio();
    assert!(sharpe.is_some() && sharpe.unwrap() > 0.0);

    // Negative PnL change
    let pnl_change = -0.5;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);

    // negative pnl_change * positive sharpe = negative reward
    assert!(reward < 0.0, "Negative PnL should produce negative reward");
    assert!(
        reward < pnl_change,
        "Sharpe amplification should make loss worse"
    );
}

/// Test history size exactly at boundary (20 samples)
#[test]
fn test_sharpe_reward_exactly_twenty_samples() {
    let mut calc = RiskAdjustmentCalculator::new(50);

    // Add exactly 20 samples
    for i in 0..20 {
        calc.add_pnl((i as f64) * 0.05);
    }

    assert_eq!(calc.history_len(), 20, "Should have exactly 20 samples");

    // Sharpe should be available
    let sharpe = calc.get_sharpe_ratio();
    assert!(sharpe.is_some(), "Sharpe should be available at boundary");

    // Reward should not be raw
    let pnl_change = 0.5;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);

    // Verify it's adjusted (not equal to raw unless Sharpe ≈ 1.0)
    if let Some(s) = sharpe {
        if (s - 1.0).abs() > 0.01 {
            assert_ne!(
                reward, pnl_change,
                "Reward should be adjusted when Sharpe ≠ 1.0"
            );
        }
    }
}

/// Test extremely large history (stress test)
#[test]
fn test_sharpe_reward_large_history() {
    let mut calc = RiskAdjustmentCalculator::new(1000);

    // Add 500 samples with varying returns
    for i in 0..500 {
        let pnl = (i as f64 % 10.0 - 5.0) * 0.1; // Varies from -0.5 to +0.45
        calc.add_pnl(pnl);
    }

    assert_eq!(calc.history_len(), 500, "Should have 500 samples");

    // Sharpe should be computable
    let sharpe = calc.get_sharpe_ratio();
    assert!(sharpe.is_some(), "Sharpe should be available");
    assert!(sharpe.unwrap().is_finite(), "Sharpe should be finite");

    // Reward calculation should not crash or overflow
    let pnl_change = 2.0;
    let reward = calc.calculate_risk_adjusted_reward(pnl_change);

    assert!(reward.is_finite(), "Reward should be finite");
}

/// Test rapid history turnover (rolling window stress test)
#[test]
fn test_sharpe_reward_rapid_history_turnover() {
    let mut calc = RiskAdjustmentCalculator::new(20);

    // Add 40 samples rapidly (causes rolling window updates)
    for i in 0..40 {
        let pnl = if i < 20 { 1.0 } else { -1.0 };
        calc.add_pnl(pnl);

        // Every 5 samples, calculate reward
        if i % 5 == 0 && i >= 20 {
            let reward = calc.calculate_risk_adjusted_reward(0.5);
            assert!(reward.is_finite(), "Reward should be finite at step {}", i);
        }
    }

    assert_eq!(calc.history_len(), 20, "Should only keep last 20 samples");
}