foxhunt/ml/tests/dqn_reward_normalization_test.rs

//! Unit tests for DQN reward normalization (Fix #1)
//!
//! Verifies that reward calculation is normalized by initial capital (10,000)
//! instead of current portfolio value, ensuring BUY and SELL rewards are symmetric.

#[cfg(test)]
mod dqn_reward_normalization_tests {
    use ml::dqn::agent::{TradingAction, TradingState};
    use ml::dqn::reward::{RewardConfig, RewardFunction};
    use rust_decimal::Decimal;

    /// Helper function to create a test state with specific portfolio features
    fn create_test_state(
        portfolio_value: f32,
        position: f32,
        spread: f32,
    ) -> TradingState {
        TradingState {
            price_features: vec![100.0, 100.0, 100.0, 100.0],
            technical_indicators: vec![0.5, 0.5, 0.5, 0.5],
            market_features: vec![spread, 100.0, 0.0, 0.0],
            portfolio_features: vec![portfolio_value, position, spread, 0.0],
        }
    }

    /// Helper function to create a reward function with default config
    fn create_reward_function() -> RewardFunction {
        let config = RewardConfig {
            pnl_weight: Decimal::ONE,
            risk_weight: Decimal::ZERO, // Disable risk penalty for cleaner tests
            cost_weight: Decimal::ZERO, // Disable cost penalty for cleaner tests
            hold_reward: Decimal::try_from(0.01).unwrap(),
            movement_threshold: Decimal::try_from(0.02).unwrap(), // 2% threshold (Fix #3)
        };
        RewardFunction::new(config)
    }

    #[test]
    fn test_buy_sell_reward_symmetry() {
        let recent_actions = vec![]; // No action history for unit test
        let mut reward_fn = create_reward_function();

        // Scenario 1: BUY action with +100 P&L
        // Portfolio value normalized: 0.9 (9,000 / 10,000)
        // After BUY: 1.01 (10,100 / 10,000)
        // P&L = 10,100 - 9,000 = +1,100 (denormalized)
        // But we want +100 P&L normalized, so:
        // current_value = 1.0, next_value = 1.01 (normalized by 10,000)
        let buy_current = create_test_state(1.0, 0.0, 0.001);
        let buy_next = create_test_state(1.01, 1.0, 0.001);

        let buy_reward = reward_fn
            .calculate_reward(TradingAction::Buy, &buy_current, &buy_next, &recent_actions)
            .expect("Buy reward calculation failed");

        // Scenario 2: SELL action with +100 P&L
        // Portfolio value normalized: 1.1 (11,000 / 10,000)
        // After SELL: 1.11 (11,100 / 10,000)
        // P&L = 11,100 - 11,000 = +100 (denormalized)
        let sell_current = create_test_state(1.0, 1.0, 0.001);
        let sell_next = create_test_state(1.01, 0.0, 0.001);

        let sell_reward = reward_fn
            .calculate_reward(TradingAction::Sell, &sell_current, &sell_next, &recent_actions)
            .expect("Sell reward calculation failed");

        // Expected reward: 100 / 10,000 = 0.01
        let expected_reward = Decimal::try_from(0.01).unwrap();

        // Assert BUY and SELL rewards are identical
        assert_eq!(
            buy_reward, sell_reward,
            "BUY and SELL rewards should be identical for same P&L. BUY: {}, SELL: {}",
            buy_reward, sell_reward
        );

        // Assert rewards match expected value (within small tolerance for floating point)
        let tolerance = Decimal::try_from(0.0001).unwrap();
        assert!(
            (buy_reward - expected_reward).abs() < tolerance,
            "BUY reward should be ~0.01. Got: {}",
            buy_reward
        );
        assert!(
            (sell_reward - expected_reward).abs() < tolerance,
            "SELL reward should be ~0.01. Got: {}",
            sell_reward
        );
    }

    #[test]
    fn test_zero_pnl_symmetry() {
        let mut reward_fn = create_reward_function();

        // Scenario 1: BUY with 0 P&L
        let buy_current = create_test_state(1.0, 0.0, 0.001);
        let buy_next = create_test_state(1.0, 1.0, 0.001); // No value change

        let buy_reward = reward_fn
            .calculate_reward(TradingAction::Buy, &buy_current, &buy_next, &recent_actions)
            .expect("Buy reward calculation failed");

        // Scenario 2: SELL with 0 P&L
        let sell_current = create_test_state(1.0, 1.0, 0.001);
        let sell_next = create_test_state(1.0, 0.0, 0.001); // No value change

        let sell_reward = reward_fn
            .calculate_reward(TradingAction::Sell, &sell_current, &sell_next, &recent_actions)
            .expect("Sell reward calculation failed");

        // Both should return 0.0
        assert_eq!(
            buy_reward,
            Decimal::ZERO,
            "BUY with 0 P&L should return 0.0 reward. Got: {}",
            buy_reward
        );
        assert_eq!(
            sell_reward,
            Decimal::ZERO,
            "SELL with 0 P&L should return 0.0 reward. Got: {}",
            sell_reward
        );

        // Assert symmetry
        assert_eq!(
            buy_reward, sell_reward,
            "BUY and SELL should have identical rewards for 0 P&L"
        );
    }

    #[test]
    fn test_negative_pnl_symmetry() {
        let mut reward_fn = create_reward_function();

        // Scenario 1: BUY with -100 P&L
        // Portfolio drops from 1.0 to 0.99 (normalized)
        // P&L = 9,900 - 10,000 = -100
        let buy_current = create_test_state(1.0, 0.0, 0.001);
        let buy_next = create_test_state(0.99, 1.0, 0.001);

        let buy_reward = reward_fn
            .calculate_reward(TradingAction::Buy, &buy_current, &buy_next, &recent_actions)
            .expect("Buy reward calculation failed");

        // Scenario 2: SELL with -100 P&L
        // Portfolio drops from 1.0 to 0.99 (normalized)
        let sell_current = create_test_state(1.0, 1.0, 0.001);
        let sell_next = create_test_state(0.99, 0.0, 0.001);

        let sell_reward = reward_fn
            .calculate_reward(TradingAction::Sell, &sell_current, &sell_next, &recent_actions)
            .expect("Sell reward calculation failed");

        // Expected reward: -100 / 10,000 = -0.01
        let expected_reward = Decimal::try_from(-0.01).unwrap();

        // Assert BUY and SELL rewards are identical
        assert_eq!(
            buy_reward, sell_reward,
            "BUY and SELL rewards should be identical for same negative P&L. BUY: {}, SELL: {}",
            buy_reward, sell_reward
        );

        // Assert rewards match expected value
        let tolerance = Decimal::try_from(0.0001).unwrap();
        assert!(
            (buy_reward - expected_reward).abs() < tolerance,
            "BUY reward should be ~-0.01. Got: {}",
            buy_reward
        );
        assert!(
            (sell_reward - expected_reward).abs() < tolerance,
            "SELL reward should be ~-0.01. Got: {}",
            sell_reward
        );

        // Verify rewards are negative
        assert!(
            buy_reward < Decimal::ZERO,
            "BUY reward should be negative for losses"
        );
        assert!(
            sell_reward < Decimal::ZERO,
            "SELL reward should be negative for losses"
        );
    }

    #[test]
    fn test_large_pnl_normalization() {
        let mut reward_fn = create_reward_function();

        // Scenario 1: BUY with +1,000 P&L
        // Portfolio increases from 1.0 to 1.1 (normalized by 10,000)
        // P&L = 11,000 - 10,000 = +1,000
        let buy_current = create_test_state(1.0, 0.0, 0.001);
        let buy_next = create_test_state(1.1, 1.0, 0.001);

        let buy_reward = reward_fn
            .calculate_reward(TradingAction::Buy, &buy_current, &buy_next, &recent_actions)
            .expect("Buy reward calculation failed");

        // Scenario 2: SELL with +1,000 P&L
        let sell_current = create_test_state(1.0, 1.0, 0.001);
        let sell_next = create_test_state(1.1, 0.0, 0.001);

        let sell_reward = reward_fn
            .calculate_reward(TradingAction::Sell, &sell_current, &sell_next, &recent_actions)
            .expect("Sell reward calculation failed");

        // Expected reward: 1,000 / 10,000 = 0.1
        let expected_reward = Decimal::try_from(0.1).unwrap();

        // Assert BUY and SELL rewards are identical
        assert_eq!(
            buy_reward, sell_reward,
            "BUY and SELL rewards should be identical for same large P&L. BUY: {}, SELL: {}",
            buy_reward, sell_reward
        );

        // Assert rewards match expected value
        let tolerance = Decimal::try_from(0.001).unwrap();
        assert!(
            (buy_reward - expected_reward).abs() < tolerance,
            "BUY reward should be ~0.1. Got: {}",
            buy_reward
        );
        assert!(
            (sell_reward - expected_reward).abs() < tolerance,
            "SELL reward should be ~0.1. Got: {}",
            sell_reward
        );

        // Verify normalization scales correctly (10% gain → 0.1 reward)
        assert!(
            buy_reward > Decimal::ZERO,
            "Large positive P&L should produce positive reward"
        );
    }

    #[test]
    fn test_hold_action_reward_flat_market() {
        let mut reward_fn = create_reward_function();

        // HOLD action in FLAT market (price change < 2%)
        // Price moves from 100.0 to 100.5 (0.5% change)
        let current = TradingState {
            price_features: vec![100.0, 100.0, 100.0, 100.0],
            technical_indicators: vec![0.5, 0.5, 0.5, 0.5],
            market_features: vec![0.001, 100.0, 0.0, 0.0],
            portfolio_features: vec![1.0, 0.5, 0.001, 0.0],
        };
        let next = TradingState {
            price_features: vec![100.5, 100.5, 100.5, 100.5], // 0.5% price change
            technical_indicators: vec![0.5, 0.5, 0.5, 0.5],
            market_features: vec![0.001, 100.5, 0.0, 0.0],
            portfolio_features: vec![1.0, 0.5, 0.001, 0.0],
        };

        let hold_reward = reward_fn
            .calculate_reward(TradingAction::Hold, &current, &next, &recent_actions)
            .expect("Hold reward calculation failed");

        // Expected: +0.002 reward for HOLD in flat market
        let expected_hold_reward = Decimal::try_from(0.002).unwrap();

        assert_eq!(
            hold_reward, expected_hold_reward,
            "HOLD action in flat market (<2% price change) should return +0.002 reward. Got: {}",
            hold_reward
        );
    }

    #[test]
    fn test_hold_action_reward_trending_market() {
        let mut reward_fn = create_reward_function();

        // HOLD action in TRENDING market (price change >= 2%)
        // Price moves from 100.0 to 105.0 (5% change)
        let current = TradingState {
            price_features: vec![100.0, 100.0, 100.0, 100.0],
            technical_indicators: vec![0.5, 0.5, 0.5, 0.5],
            market_features: vec![0.001, 100.0, 0.0, 0.0],
            portfolio_features: vec![1.0, 0.5, 0.001, 0.0],
        };
        let next = TradingState {
            price_features: vec![105.0, 105.0, 105.0, 105.0], // 5% price change
            technical_indicators: vec![0.5, 0.5, 0.5, 0.5],
            market_features: vec![0.001, 105.0, 0.0, 0.0],
            portfolio_features: vec![1.0, 0.5, 0.001, 0.0],
        };

        let hold_reward = reward_fn
            .calculate_reward(TradingAction::Hold, &current, &next, &recent_actions)
            .expect("Hold reward calculation failed");

        // Expected: -0.001 penalty for HOLD in trending market
        let expected_hold_penalty = Decimal::try_from(-0.001).unwrap();

        assert_eq!(
            hold_reward, expected_hold_penalty,
            "HOLD action in trending market (>=2% price change) should return -0.001 penalty. Got: {}",
            hold_reward
        );
    }

    let recent_actions = vec![]; // No action history for unit test
   #[test]
    fn test_normalization_eliminates_portfolio_value_bias() {
        let mut reward_fn = create_reward_function();

        // Test 1: Same absolute P&L (+100) with different portfolio sizes
        // Small portfolio: 5,000 → 5,100
        let small_current = create_test_state(0.5, 0.0, 0.001);
        let small_next = create_test_state(0.51, 1.0, 0.001);

        let small_reward = reward_fn
            .calculate_reward(TradingAction::Buy, &small_current, &small_next, &recent_actions)
            .expect("Small portfolio reward failed");

        // Large portfolio: 20,000 → 20,100
        let large_current = create_test_state(2.0, 0.0, 0.001);
        let large_next = create_test_state(2.01, 1.0, 0.001);

        let large_reward = reward_fn
            .calculate_reward(TradingAction::Buy, &large_current, &large_next, &recent_actions)
            .expect("Large portfolio reward failed");

        // Both should get identical rewards (+100 / 10,000 = 0.01)
        // because normalization uses CONSTANT denominator (10,000), not current value
        let tolerance = Decimal::try_from(0.0001).unwrap();
        assert!(
            (small_reward - large_reward).abs() < tolerance,
            "Rewards should be identical regardless of portfolio size. Small: {}, Large: {}",
            small_reward,
            large_reward
        );

        // Expected reward: 100 / 10,000 = 0.01
        let expected = Decimal::try_from(0.01).unwrap();
        assert!(
            (small_reward - expected).abs() < tolerance,
            "Small portfolio reward should be ~0.01. Got: {}",
            small_reward
        );
        assert!(
            (large_reward - expected).abs() < tolerance,
            "Large portfolio reward should be ~0.01. Got: {}",
            large_reward
        );
    }

    #[test]
    fn test_percentage_based_normalization() {
        let mut reward_fn = create_reward_function();

        // Test that normalization works as percentage of initial capital
        // 5% gain = 500 / 10,000 = 0.05 reward
        let current = create_test_state(1.0, 0.0, 0.001);
        let next = create_test_state(1.05, 1.0, 0.001);

        let reward = reward_fn
            .calculate_reward(TradingAction::Buy, &current, &next, &recent_actions)
            .expect("Reward calculation failed");

        let expected = Decimal::try_from(0.05).unwrap();
        let tolerance = Decimal::try_from(0.001).unwrap();

        assert!(
            (reward - expected).abs() < tolerance,
            "5% gain (500/10,000) should produce 0.05 reward. Got: {}",
            reward
        );
    }
}