foxhunt/ml/tests/dqn_zero_price_fix_test.rs

//! Test suite for zero price error fix in calculate_hold_reward
//!
//! Tests that HOLD reward calculation correctly handles log returns
//! instead of treating them as raw prices (which caused division by zero).

use ml::dqn::agent::{TradingAction, TradingState};
use ml::dqn::reward::{calculate_batch_rewards, RewardConfig, RewardFunction};
use rust_decimal::Decimal;

fn create_test_config() -> RewardConfig {
    RewardConfig {
        pnl_weight: Decimal::ONE,
        risk_weight: Decimal::try_from(0.1).unwrap(),
        cost_weight: Decimal::try_from(0.05).unwrap(),
        hold_reward: Decimal::try_from(0.001).unwrap(),
        movement_threshold: Decimal::try_from(0.02).unwrap(),
        hold_penalty_weight: Decimal::try_from(0.5).unwrap(),
        diversity_weight: Decimal::try_from(-0.1).unwrap(),
    }
}

#[test]
fn test_hold_reward_with_zero_log_return() {
    // Test that zero log returns don't crash (stable prices)
    // Velocity-based: volatility = |next_log_return| = 0.0 < 0.02 threshold
    // Expected: hold_reward (0.001) granted
    let config = create_test_config();
    let mut reward_fn = RewardFunction::new(config.clone());

    let current_state = TradingState {
        price_features: vec![0.0, 100.0, 100.0, 100.0], // Zero log return (not used in velocity calc)
        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.0, 0.0, 0.0],
    };

    let next_state = TradingState {
        price_features: vec![0.0, 100.0, 100.0, 100.0], // Zero log return (stable price)
        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.0, 0.0, 0.0],
    };

    let recent_actions = vec![TradingAction::Buy, TradingAction::Sell, TradingAction::Hold];

    // Should NOT crash with zero log return
    let reward = reward_fn.calculate_reward(
        TradingAction::Hold,
        &current_state,
        &next_state,
        &recent_actions,
    );

    assert!(
        reward.is_ok(),
        "Should handle zero log returns without crashing, got: {:?}",
        reward
    );

    // Should return hold_reward (low volatility)
    let reward_value = reward.unwrap();
    println!("Zero log return reward: {}", reward_value);
    // Velocity = |0.0| = 0.0 < 0.02, so should grant hold_reward (0.001)
    // Note: diversity penalty (-0.1) may also be applied due to low entropy
    assert!(
        reward_value <= config.hold_reward,
        "Low volatility should result in hold_reward or less (with diversity penalty), got: {}",
        reward_value
    );
}

#[test]
fn test_hold_reward_high_volatility() {
    // Test high volatility triggers penalty
    // Velocity-based: volatility = |0.05| = 0.05 > 0.02 threshold
    // Expected: -hold_penalty_weight (-0.5) applied
    let config = create_test_config();
    let mut reward_fn = RewardFunction::new(config.clone());

    let current_state = TradingState {
        price_features: vec![0.0, 100.0, 100.0, 100.0], // Not used in velocity calc
        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.0, 0.0, 0.0],
    };

    let next_state = TradingState {
        price_features: vec![0.05, 100.0, 100.0, 100.0], // 5% log return (> 0.02 threshold)
        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.0, 0.0, 0.0],
    };

    let recent_actions = vec![TradingAction::Buy, TradingAction::Sell, TradingAction::Hold];

    let reward = reward_fn.calculate_reward(
        TradingAction::Hold,
        &current_state,
        &next_state,
        &recent_actions,
    );

    assert!(
        reward.is_ok(),
        "Should handle high volatility, got: {:?}",
        reward
    );

    let reward_value = reward.unwrap();
    println!("High volatility reward: {}", reward_value);

    // Velocity = |0.05| = 0.05 > 0.02, so should apply penalty (-0.5)
    // With diversity penalty (-0.1), total = -0.5 - 0.1 = -0.6
    assert!(
        reward_value < Decimal::ZERO,
        "High volatility should trigger penalty, got: {}",
        reward_value
    );
}

#[test]
fn test_hold_reward_negative_log_return() {
    // Test negative log returns (price decrease)
    // Velocity-based: volatility = |-0.03| = 0.03 > 0.02 threshold
    // Expected: penalty triggered (large downward move)
    let config = create_test_config();
    let mut reward_fn = RewardFunction::new(config.clone());

    let current_state = TradingState {
        price_features: vec![0.01, 100.0, 100.0, 100.0], // Not used in velocity calc
        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.0, 0.0, 0.0],
    };

    let next_state = TradingState {
        price_features: vec![-0.03, 100.0, 100.0, 100.0], // -3% log return (downward move)
        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.0, 0.0, 0.0],
    };

    let recent_actions = vec![TradingAction::Buy, TradingAction::Sell, TradingAction::Hold];

    let reward = reward_fn.calculate_reward(
        TradingAction::Hold,
        &current_state,
        &next_state,
        &recent_actions,
    );

    assert!(
        reward.is_ok(),
        "Should handle negative log returns, got: {:?}",
        reward
    );

    let reward_value = reward.unwrap();
    println!("Negative log return reward: {}", reward_value);

    // Velocity = |-0.03| = 0.03 > 0.02, so penalty should be applied
    // Large downward move should NOT be rewarded
    assert!(
        reward_value < Decimal::ZERO,
        "Large price decrease should trigger penalty, got: {}",
        reward_value
    );
}

#[test]
fn test_batch_rewards_with_mixed_log_returns() {
    // Test batch processing with mixed log return scenarios
    // Velocity-based logic:
    //   - Sample 1: volatility = |0.001| < 0.02 → hold_reward (0.001)
    //   - Sample 2: volatility = |0.05| > 0.02 → penalty (-0.5)
    let config = create_test_config();
    let mut reward_fn = RewardFunction::new(config.clone());

    let current_states = vec![
        TradingState {
            price_features: vec![0.0, 100.0, 100.0, 100.0], // Not used in velocity calc
            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.0, 0.0, 0.0],
        },
        TradingState {
            price_features: vec![0.0, 100.0, 100.0, 100.0], // Not used in velocity calc
            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.0, 0.0, 0.0],
        },
    ];

    let next_states = vec![
        TradingState {
            price_features: vec![0.001, 100.0, 100.0, 100.0], // Low volatility (0.1%)
            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.0, 0.0, 0.0],
        },
        TradingState {
            price_features: vec![0.05, 100.0, 100.0, 100.0], // High volatility (5%)
            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.0, 0.0, 0.0],
        },
    ];

    let actions = vec![TradingAction::Hold, TradingAction::Hold];
    let recent_actions = vec![TradingAction::Buy, TradingAction::Sell, TradingAction::Hold];

    let rewards = calculate_batch_rewards(
        &mut reward_fn,
        &actions,
        &current_states,
        &next_states,
        &recent_actions,
    );

    assert!(
        rewards.is_ok(),
        "Batch rewards should handle mixed log returns, got: {:?}",
        rewards
    );

    let reward_values = rewards.unwrap();
    assert_eq!(reward_values.len(), 2);
    println!("Batch rewards: {:?}", reward_values);

    // First HOLD (low volatility) should be less negative than second (high volatility)
    // Sample 1: |0.001| < 0.02 → reward (0.001 or less with diversity penalty)
    // Sample 2: |0.05| > 0.02 → penalty (-0.5 or less with diversity penalty)
    assert!(
        reward_values[0] > reward_values[1],
        "Low volatility HOLD should have better reward than high volatility HOLD, got: {:?}",
        reward_values
    );
}