foxhunt/crates/ml/tests/risk_action_masking_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! TDD Tests for Risk-Based Action Masking in DQN
//!
//! Tests verify that action masking correctly filters invalid actions based on:
//! 1. Position limit violations (±2.0 max position)
//! 2. Drawdown limit violations (15% max drawdown)
//! 3. Value-at-Risk (VaR) violations
//! 4. Cash reserve violations (20% minimum cash)
//! 5. Position-reducing actions are never masked
//! 6. Action diversity is preserved with masking
//!
//! Expected: 30-50% of invalid actions filtered before Q-value computation
//! Impact: Faster training, better convergence, safer trading

use ml::dqn::action_space::{FactoredAction, OrderType, get_valid_action_mask};
use std::time::Instant;
use tracing::info;

/// Helper struct to simulate portfolio state for risk checking
#[derive(Debug, Clone)]
struct PortfolioState {
    /// Current portfolio position (-2.0 to +2.0)
    pub position: f64,
    /// Current portfolio value in dollars
    pub portfolio_value: f64,
    /// Cash reserve in dollars
    pub cash_reserve: f64,
    /// Current peak portfolio value (for drawdown calculation)
    pub peak_value: f64,
    /// VaR limit as percentage of portfolio
    pub var_limit_pct: f64,
}

impl PortfolioState {
    /// Create a new portfolio state
    fn new(position: f64, portfolio_value: f64, cash_reserve: f64) -> Self {
        Self {
            position,
            portfolio_value,
            cash_reserve,
            peak_value: portfolio_value,
            var_limit_pct: 5.0, // 5% VaR limit
        }
    }

    /// Check if action would violate position limit
    fn violates_position_limit(&self, action: &FactoredAction, max_position: f64) -> bool {
        let target_exposure = action.target_exposure();

        // Fix: At position limits, mask actions in the same direction
        // At max long position (≥2.0), mask all LONG actions (Long50, Long100)
        // At max short position (≤-2.0), mask all SHORT actions (Short50, Short100)
        if self.position >= max_position {
            // At or above max long position, mask all long (buy) actions
            return action.is_buy();
        } else if self.position <= -max_position {
            // At or below max short position, mask all short (sell) actions
            return action.is_sell();
        }

        // Otherwise, just check if target would exceed absolute limit
        target_exposure.abs() > max_position
    }

    /// Check if action would violate drawdown limit (15% max)
    fn violates_drawdown_limit(&self, action: &FactoredAction, max_drawdown_pct: f64) -> bool {
        // Simulate action: assume 1% price movement in direction of action
        let price_movement = if action.is_buy() {
            0.01
        } else if action.is_sell() {
            -0.01
        } else {
            0.0 // HOLD
        };

        // Calculate portfolio value after action
        let position_change = action.target_exposure() - self.position;
        let portfolio_change = self.portfolio_value * price_movement * position_change.abs();
        let new_portfolio_value = self.portfolio_value + portfolio_change;

        // Check drawdown
        let drawdown_pct = ((self.peak_value - new_portfolio_value) / self.peak_value) * 100.0;
        drawdown_pct > max_drawdown_pct
    }

    /// Check if action would violate VaR limit
    fn violates_var_limit(&self, action: &FactoredAction) -> bool {
        // VaR = maximum potential loss at 95% confidence (2 sigma move)
        let var_dollar = self.portfolio_value * (self.var_limit_pct / 100.0);

        // Calculate exposure change (absolute values)
        let current_exposure = self.position.abs();
        let target_exposure = action.target_exposure().abs();
        let exposure_change = target_exposure - current_exposure;

        // Only check if exposure is increasing
        if exposure_change <= 0.0 {
            return false; // Reducing exposure never violates VaR
        }

        // Fix: Potential loss is exposure change times a volatility factor (e.g., 2% daily move)
        // This represents the potential loss from a 2-sigma price movement on the new exposure
        // VaR formula: potential_loss = (exposure_change * portfolio_value) * volatility_pct
        // Use 2% as a reasonable daily volatility estimate for futures
        let volatility_pct = 0.02; // 2% potential daily move
        let potential_loss = exposure_change * self.portfolio_value * volatility_pct;

        // Violates if potential loss exceeds VaR limit
        potential_loss > var_dollar
    }

    /// Check if action would violate cash reserve requirement (20% minimum)
    fn violates_cash_reserve(&self, action: &FactoredAction) -> bool {
        // Fix: HOLD actions have zero transaction cost, so they never violate cash reserve
        if action.is_hold() {
            return false;
        }

        let min_cash = self.portfolio_value * 0.20; // 20% minimum

        // Fix: Calculate transaction cost based on actual exposure change, not arbitrary 50%
        // Trade value = |target_exposure - current_exposure| * portfolio_value
        let exposure_change = (action.target_exposure() - self.position).abs();
        let trade_value = exposure_change * self.portfolio_value;
        let transaction_cost = action.calculate_transaction_cost(trade_value);

        (self.cash_reserve - transaction_cost) < min_cash
    }

    /// Check if action reduces position (never mask)
    fn reduces_position(&self, action: &FactoredAction) -> bool {
        // Fix: Position-reducing means moving TOWARDS zero, not just reducing absolute value
        // At position +2.0, Short actions reduce position (move towards 0)
        // At position -2.0, Long actions reduce position (move towards 0)

        let current_pos = self.position;
        let target_pos = action.target_exposure();

        // Check if moving towards zero
        if current_pos > 0.0 {
            // Long position: reducing means target is less than current (moving left towards 0)
            target_pos < current_pos
        } else if current_pos < 0.0 {
            // Short position: reducing means target is greater than current (moving right towards 0)
            target_pos > current_pos
        } else {
            // At zero position, no action reduces position
            false
        }
    }

    /// Check if all risk constraints are satisfied
    fn is_action_valid(&self, action: &FactoredAction, max_position: f64) -> bool {
        // Fix: HOLD actions are ALWAYS valid (safety valve) - they have zero cost and maintain position
        if action.is_hold() {
            return true;
        }

        // Fix: Check position limits FIRST (before position-reducing bypass)
        // At position limit, we want to mask directional actions (BUY at max long, SELL at max short)
        if self.violates_position_limit(action, max_position) {
            return false;
        }

        // Position-reducing actions are valid after passing position limit check
        if self.reduces_position(action) {
            return self.portfolio_value > 0.0;
        }

        // Check remaining risk limits
        !self.violates_drawdown_limit(action, 15.0)
            && !self.violates_var_limit(action)
            && !self.violates_cash_reserve(action)
    }

    /// Get valid action indices based on risk constraints
    fn get_valid_actions(&self, max_position: f64) -> Vec<usize> {
        (0..45)
            .filter(|&idx| {
                if let Ok(action) = FactoredAction::from_index(idx) {
                    self.is_action_valid(&action, max_position)
                } else {
                    false
                }
            })
            .collect()
    }

    /// Count masked actions
    fn count_masked_actions(&self, max_position: f64) -> usize {
        45 - self.get_valid_actions(max_position).len()
    }
}

// ============================================================================
// TEST 1: Mask actions exceeding position limit
// ============================================================================
#[test]
fn test_mask_actions_exceeding_position_limit() {
    let state = PortfolioState::new(0.0, 100_000.0, 20_000.0);
    let max_position = 1.0; // Restrictive limit

    let valid_actions = state.get_valid_actions(max_position);

    // At max_position=1.0, Long100 and Short100 should be masked
    // because they set exposure to ±1.0 which exceeds the limit in absolute terms
    for idx in &valid_actions {
        let action = FactoredAction::from_index(*idx).unwrap();
        assert!(
            !state.violates_position_limit(&action, max_position),
            "Action {} should not violate position limit",
            idx
        );
    }

    // Verify that some actions are actually masked
    assert!(
        valid_actions.len() < 45,
        "Expected some actions to be masked with restrictive limit"
    );
}

// ============================================================================
// TEST 2: Mask actions violating drawdown limit
// ============================================================================
#[test]
fn test_mask_actions_violating_drawdown() {
    // Portfolio near peak with little room for drawdown
    let mut state = PortfolioState::new(1.0, 100_000.0, 20_000.0);
    state.peak_value = 105_000.0; // Already down from peak

    // Aggressive BUY actions should be masked (would increase drawdown)
    let valid_actions = state.get_valid_actions(2.0);

    for idx in &valid_actions {
        let action = FactoredAction::from_index(*idx).unwrap();
        let violates_dd = state.violates_drawdown_limit(&action, 15.0);
        assert!(
            !violates_dd,
            "Action {} should not violate drawdown limit",
            idx
        );
    }

    // Verify that aggressive actions might be filtered
    let valid_buys = valid_actions
        .iter()
        .filter(|&&idx| {
            if let Ok(action) = FactoredAction::from_index(idx) {
                action.is_buy()
            } else {
                false
            }
        })
        .count();

    // With significant drawdown history, aggressive actions should be limited
    assert!(
        valid_buys <= 18,
        "Expected limited BUY actions when portfolio is down"
    );
}

// ============================================================================
// TEST 3: Mask actions violating VaR limit
// ============================================================================
#[test]
fn test_mask_actions_violating_var() {
    let state = PortfolioState::new(2.0, 50_000.0, 10_000.0); // High exposure + low cash

    let valid_actions = state.get_valid_actions(2.0);

    // No valid action should violate VaR
    for idx in valid_actions {
        let action = FactoredAction::from_index(idx).unwrap();
        let violates_var = state.violates_var_limit(&action);
        assert!(
            !violates_var,
            "Action {} should not violate VaR limit",
            idx
        );
    }
}

// ============================================================================
// TEST 4: Mask actions violating cash reserve requirement
// ============================================================================
#[test]
fn test_mask_actions_violating_cash_reserve() {
    // Portfolio with insufficient cash reserve
    let state = PortfolioState::new(0.0, 100_000.0, 5_000.0); // Only 5% cash

    let valid_actions = state.get_valid_actions(2.0);

    for idx in &valid_actions {
        let action = FactoredAction::from_index(*idx).unwrap();
        let violates_cash = state.violates_cash_reserve(&action);
        assert!(
            !violates_cash,
            "Action {} should not violate cash reserve requirement",
            idx
        );
    }

    // Expensive market orders should be masked
    let market_actions: Vec<_> = (0..45)
        .filter(|&idx| {
            if let Ok(action) = FactoredAction::from_index(idx) {
                action.order == OrderType::Market
            } else {
                false
            }
        })
        .collect();

    let valid_market: Vec<_> = valid_actions
        .iter()
        .filter(|&&idx| {
            if let Ok(action) = FactoredAction::from_index(idx) {
                action.order == OrderType::Market
            } else {
                false
            }
        })
        .collect();

    assert!(
        valid_market.len() < market_actions.len(),
        "Expected market orders to be filtered when cash is low"
    );
}

// ============================================================================
// TEST 5: Allow position-reducing actions
// ============================================================================
#[test]
fn test_allow_position_reducing_actions() {
    // Portfolio with large long position
    let state = PortfolioState::new(1.5, 100_000.0, 20_000.0);

    let valid_actions = state.get_valid_actions(1.0);

    // Position-reducing actions (SHORT50, SHORT100, FLAT) should always be valid
    // Check that at least one SELL action is valid
    let valid_sells: Vec<_> = valid_actions
        .iter()
        .filter(|&&idx| {
            if let Ok(action) = FactoredAction::from_index(idx) {
                action.is_sell() || action.is_hold()
            } else {
                false
            }
        })
        .collect();

    assert!(
        !valid_sells.is_empty(),
        "Position-reducing actions should always be valid"
    );
}

// ============================================================================
// TEST 6: Mask all long actions at max position
// ============================================================================
#[test]
fn test_mask_all_long_actions_at_max_long() {
    // Portfolio at maximum long position
    let state = PortfolioState::new(2.0, 100_000.0, 20_000.0);
    let max_position = 2.0;

    let valid_actions = state.get_valid_actions(max_position);

    // BUY actions (Long50, Long100) should all be masked
    let valid_buys: Vec<_> = valid_actions
        .iter()
        .filter(|&&idx| {
            if let Ok(action) = FactoredAction::from_index(idx) {
                action.is_buy()
            } else {
                false
            }
        })
        .collect();

    assert!(
        valid_buys.is_empty(),
        "BUY actions should all be masked at max position"
    );

    // HOLD and SELL should be available
    let valid_holds_sells: Vec<_> = valid_actions
        .iter()
        .filter(|&&idx| {
            if let Ok(action) = FactoredAction::from_index(idx) {
                action.is_hold() || action.is_sell()
            } else {
                false
            }
        })
        .collect();

    assert!(
        !valid_holds_sells.is_empty(),
        "HOLD and SELL actions should be available at max position"
    );
}

// ============================================================================
// TEST 7: Mask all short actions at min position
// ============================================================================
#[test]
fn test_mask_all_short_actions_at_max_short() {
    // Portfolio at maximum short position
    let state = PortfolioState::new(-2.0, 100_000.0, 20_000.0);
    let max_position = 2.0;

    let valid_actions = state.get_valid_actions(max_position);

    // SELL actions (Short50, Short100) should all be masked
    let valid_sells: Vec<_> = valid_actions
        .iter()
        .filter(|&&idx| {
            if let Ok(action) = FactoredAction::from_index(idx) {
                action.is_sell()
            } else {
                false
            }
        })
        .collect();

    assert!(
        valid_sells.is_empty(),
        "SELL actions should all be masked at min position"
    );

    // HOLD and BUY should be available
    let valid_holds_buys: Vec<_> = valid_actions
        .iter()
        .filter(|&&idx| {
            if let Ok(action) = FactoredAction::from_index(idx) {
                action.is_hold() || action.is_buy()
            } else {
                false
            }
        })
        .collect();

    assert!(
        !valid_holds_buys.is_empty(),
        "HOLD and BUY actions should be available at min position"
    );
}

// ============================================================================
// TEST 8: HOLD always valid (unless bankrupt)
// ============================================================================
#[test]
fn test_valid_actions_include_hold() {
    let states = vec![
        PortfolioState::new(-2.0, 100_000.0, 5_000.0),  // Min position, low cash
        PortfolioState::new(0.0, 100_000.0, 5_000.0),   // Flat, low cash
        PortfolioState::new(2.0, 100_000.0, 5_000.0),   // Max position, low cash
        PortfolioState::new(1.0, 20_000.0, 1_000.0),    // Small portfolio
    ];

    for state in states {
        let valid_actions = state.get_valid_actions(2.0);

        // HOLD actions (all FLAT with any order type and urgency)
        let valid_holds: Vec<_> = valid_actions
            .iter()
            .filter(|&&idx| {
                if let Ok(action) = FactoredAction::from_index(idx) {
                    action.is_hold()
                } else {
                    false
                }
            })
            .collect();

        assert!(
            !valid_holds.is_empty(),
            "HOLD action should always be valid for position {}",
            state.position
        );
    }
}

// ============================================================================
// TEST 9: Action masking performance (<1ms for 45 actions)
// ============================================================================
#[test]
fn test_action_mask_performance() {
    let state = PortfolioState::new(0.5, 100_000.0, 20_000.0);

    // Measure masking performance
    let start = Instant::now();
    for _ in 0..1000 {
        let _ = state.get_valid_actions(2.0);
    }
    let elapsed = start.elapsed();

    let avg_time_us = elapsed.as_micros() / 1000;
    info!(avg_time_us, "Average masking time per call (µs, 1000 iterations)");

    // Should complete 1000 calls in <1 second (1000 µs average per call)
    assert!(
        elapsed.as_millis() < 1000,
        "Masking should complete in <1ms per call (got {:.3}ms)",
        elapsed.as_secs_f64() * 1000.0 / 1000.0
    );
}

// ============================================================================
// TEST 10: Masked actions not used in Q-value computation
// ============================================================================
#[test]
fn test_masked_actions_not_in_qvalue_computation() {
    let state = PortfolioState::new(1.8, 100_000.0, 20_000.0);
    let max_position = 2.0;

    let valid_actions = state.get_valid_actions(max_position);
    let mask = get_valid_action_mask(state.position, max_position);

    // Verify consistency: valid actions from risk check should align with position masking
    let position_masked: Vec<usize> = (0..45)
        .filter(|&idx| !mask[idx])
        .collect();

    let risk_masked: Vec<usize> = (0..45)
        .filter(|&idx| !valid_actions.contains(&idx))
        .collect();

    // Both masks should agree on position-based exclusions
    for idx in &position_masked {
        let action = FactoredAction::from_index(*idx).unwrap();
        let abs_exposure = action.target_exposure().abs();
        assert!(
            abs_exposure > max_position,
            "Position mask should only exclude actions exceeding limit"
        );
    }

    // Risk mask is a superset of position mask (position-based + risk-based)
    assert!(
        position_masked.len() <= risk_masked.len(),
        "Risk-based mask should exclude at least as many actions as position mask"
    );
}

// ============================================================================
// TEST 11: Action diversity with masking
// ============================================================================
#[test]
fn test_action_diversity_with_masking() {
    // Even with masking, should preserve action diversity
    let state = PortfolioState::new(0.5, 100_000.0, 20_000.0);

    let valid_actions = state.get_valid_actions(2.0);

    // Count exposure type distribution among valid actions
    let mut exposure_types = std::collections::HashSet::new();
    let mut order_types = std::collections::HashSet::new();

    for &idx in &valid_actions {
        if let Ok(action) = FactoredAction::from_index(idx) {
            exposure_types.insert(format!("{}", action.exposure));
            order_types.insert(format!("{}", action.order));
        }
    }

    // Should have multiple exposure types
    assert!(
        exposure_types.len() > 1,
        "Masking should preserve exposure type diversity"
    );

    // Should have multiple order types
    assert!(
        order_types.len() > 1,
        "Masking should preserve order type diversity"
    );

    // At least 50% of actions should remain valid (flexible masking)
    assert!(
        valid_actions.len() >= 22,
        "Masking should leave at least 50% of actions valid"
    );
}

// ============================================================================
// TEST 12: Mask logging (information for monitoring)
// ============================================================================
#[test]
fn test_mask_logging() {
    let state = PortfolioState::new(0.0, 100_000.0, 5_000.0);

    let valid_count = state.get_valid_actions(2.0).len();
    let masked_count = state.count_masked_actions(2.0);

    info!(
        valid_count,
        masked_count,
        masking_rate_pct = (masked_count as f64 / 45.0) * 100.0,
        "Mask statistics"
    );

    // Log should provide useful information
    assert!(valid_count + masked_count == 45, "Counts should sum to 45");

    // Verify masking rate is reasonable (10-50%)
    let masking_rate = (masked_count as f64 / 45.0) * 100.0;
    assert!(
        masking_rate >= 0.0 && masking_rate <= 100.0,
        "Masking rate should be 0-100%"
    );

    // Print detailed breakdown
    info!(
        masking_rate_pct = masking_rate,
        valid_actions_pct = (valid_count as f64 / 45.0) * 100.0,
        "Expected impact: reduce Q-value computation by masking rate, focusing on valid actions"
    );
}

// ============================================================================
// INTEGRATION TEST: Multi-scenario masking consistency
// ============================================================================
#[test]
fn test_risk_masking_consistency_across_scenarios() {
    let scenarios = vec![
        ("Flat portfolio", PortfolioState::new(0.0, 100_000.0, 20_000.0)),
        ("Long position", PortfolioState::new(1.5, 100_000.0, 20_000.0)),
        ("Short position", PortfolioState::new(-1.5, 100_000.0, 20_000.0)),
        ("Low cash", PortfolioState::new(0.5, 100_000.0, 3_000.0)),
        ("Large portfolio", PortfolioState::new(1.0, 500_000.0, 100_000.0)),
    ];

    for (name, state) in scenarios {
        let valid_actions = state.get_valid_actions(2.0);

        // Consistency check: valid actions should never violate constraints
        for &idx in &valid_actions {
            if let Ok(action) = FactoredAction::from_index(idx) {
                assert!(
                    state.is_action_valid(&action, 2.0),
                    "Scenario '{}': Invalid action {} in valid list",
                    name,
                    idx
                );
            }
        }

        info!(
            scenario = name,
            valid_actions = valid_actions.len(),
            valid_pct = (valid_actions.len() as f64 / 45.0) * 100.0,
            "Scenario actions valid"
        );
    }
}

// ============================================================================
// EDGE CASE TEST: Bankrupt portfolio
// ============================================================================
#[test]
fn test_masking_with_bankrupt_portfolio() {
    // Portfolio with portfolio_value <= 0
    let state = PortfolioState::new(0.0, 100.0, 50.0); // Very small portfolio

    let valid_actions = state.get_valid_actions(2.0);

    // Even with small portfolio, should have some valid actions (HOLD)
    assert!(
        !valid_actions.is_empty(),
        "Even bankrupt portfolio should have valid actions (HOLD)"
    );
}

// ============================================================================
// EDGE CASE TEST: Extreme position limits
// ============================================================================
#[test]
fn test_masking_with_extreme_limits() {
    let state = PortfolioState::new(0.0, 100_000.0, 20_000.0);

    // Very restrictive limit
    let valid_restrictive = state.get_valid_actions(0.3);

    // Very permissive limit
    let valid_permissive = state.get_valid_actions(3.0);

    // Permissive should have more or equal valid actions
    assert!(
        valid_permissive.len() >= valid_restrictive.len(),
        "Permissive limits should allow more or equal actions"
    );

    // Restrictive should mask some actions
    assert!(
        valid_restrictive.len() < 45,
        "Restrictive limits should mask some actions"
    );
}