foxhunt/ml/tests/ppo_risk_integration_tests.rs

//! PPO Risk Management Integration Tests (TDD Red Phase)
//!
//! Tests for integrating PortfolioTracker, CircuitBreaker, and RewardNormalizer
//! into the PPO training pipeline for Wave 1 of the DQN-PPO alignment project.
//!
//! **Expected Status**: These tests should FAIL until implementation is complete.

use anyhow::Result;
use candle_core::Device;
use ml::dqn::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
use ml::dqn::portfolio_tracker::PortfolioTracker;
use ml::dqn::reward::RewardNormalizer;
use ml::dqn::TradingAction;
use ml::ppo::ppo::{PPOConfig, WorkingPPO};
use ml::ppo::trajectories::{Trajectory, TrajectoryBatch, TrajectoryStep};
use std::time::Duration;

/// Test 1: PPO PortfolioTracker Integration
///
/// **Goal**: Verify PortfolioTracker is properly integrated into PPO training loop
/// **Expected Behavior**:
/// - PPO struct contains PortfolioTracker field
/// - PortfolioTracker updates after each action
/// - Portfolio state (cash, position, P&L) is tracked correctly
#[test]
fn test_ppo_portfolio_tracker_integration() -> Result<()> {
    // Create PPO with PortfolioTracker
    let config = PPOConfig {
        state_dim: 64,
        num_actions: 3,
        policy_hidden_dims: vec![32, 16],
        value_hidden_dims: vec![32, 16],
        policy_learning_rate: 1e-4,
        value_learning_rate: 1e-3,
        ..Default::default()
    };

    let mut ppo = WorkingPPO::new(config)?;

    // Access PortfolioTracker (this will fail until implementation)
    let initial_cash = ppo.portfolio_tracker.cash_balance();
    assert_eq!(
        initial_cash, 10_000.0,
        "Initial cash should be 10,000"
    );

    // Simulate action execution
    let state = vec![0.5; 64];
    let (action, _value) = ppo.act(&state)?;

    // Verify PortfolioTracker updated
    let current_position = ppo.portfolio_tracker.current_position();
    match action {
        TradingAction::Buy => {
            assert!(
                current_position > 0.0,
                "Buy action should create long position"
            );
        },
        TradingAction::Sell => {
            assert!(
                current_position < 0.0,
                "Sell action should create short position"
            );
        },
        TradingAction::Hold => {
            assert_eq!(
                current_position, 0.0,
                "Hold action should keep neutral position"
            );
        },
    }

    Ok(())
}

/// Test 2: PPO CircuitBreaker Integration
///
/// **Goal**: Verify CircuitBreaker stops training on consecutive failures
/// **Expected Behavior**:
/// - CircuitBreaker opens after N consecutive failures (default: 5)
/// - Training loop stops when circuit is open
/// - CircuitBreaker state is logged
#[test]
fn test_ppo_circuit_breaker_integration() -> Result<()> {
    // Create CircuitBreaker with aggressive threshold for testing
    let circuit_config = CircuitBreakerConfig {
        failure_threshold: 3, // Open after 3 failures
        success_threshold: 2,
        timeout_duration: Duration::from_secs(10),
        half_open_max_calls: 1,
    };
    let circuit_breaker = CircuitBreaker::new(circuit_config);

    // Simulate 3 consecutive failures
    circuit_breaker.record_failure();
    circuit_breaker.record_failure();
    circuit_breaker.record_failure();

    // Circuit should be open
    let stats = circuit_breaker.stats();
    assert_eq!(
        stats.state,
        ml::dqn::circuit_breaker::CircuitState::Open,
        "Circuit should be open after 3 failures"
    );
    assert!(!circuit_breaker.allow_request(), "Circuit should block requests");

    // Create PPO config with circuit breaker
    let config = PPOConfig {
        state_dim: 64,
        num_actions: 3,
        ..Default::default()
    };

    let mut ppo = WorkingPPO::new(config)?;

    // Verify PPO has circuit breaker field (will fail until implementation)
    assert!(
        ppo.circuit_breaker.is_some(),
        "PPO should have CircuitBreaker field"
    );

    // Simulate training with open circuit
    let state = vec![0.5; 64];
    let mut trajectory = Trajectory::new();
    for _ in 0..10 {
        let (action, value) = ppo.act(&state)?;
        let step = TrajectoryStep::new(state.clone(), action, -1.0, value, -0.1, false);
        trajectory.add_step(step);
    }

    let mut batch = TrajectoryBatch::from_trajectories(
        vec![trajectory],
        vec![0.1; 10],
        vec![0.2; 10],
    );

    // Training update should respect circuit breaker state
    let result = ppo.update(&mut batch);

    // If circuit is open, update should either:
    // 1. Return an error indicating circuit is open
    // 2. Skip the update and return zero loss
    match result {
        Err(e) => {
            assert!(
                e.to_string().contains("circuit") || e.to_string().contains("breaker"),
                "Error should mention circuit breaker: {}",
                e
            );
        },
        Ok((policy_loss, value_loss)) => {
            // If update proceeds, losses should be minimal (no actual update)
            assert!(
                policy_loss < 0.01 && value_loss < 0.01,
                "Losses should be minimal when circuit is open"
            );
        },
    }

    Ok(())
}

/// Test 3: PPO Reward Normalization Integration
///
/// **Goal**: Verify RewardNormalizer is applied to rewards during training
/// **Expected Behavior**:
/// - Rewards are normalized to ~N(0,1) distribution
/// - Normalization statistics (mean, std) are tracked
/// - Normalized rewards are used in advantage calculation
#[test]
fn test_ppo_reward_normalization_integration() -> Result<()> {
    // Create RewardNormalizer
    let mut normalizer = RewardNormalizer::new();

    // Feed rewards with non-zero mean and variance
    let raw_rewards = vec![10.0, 20.0, 30.0, 40.0, 50.0];
    for &reward in &raw_rewards {
        normalizer.update(reward);
    }

    // Verify normalization statistics
    let (mean, std) = normalizer.get_stats();
    assert!(
        (mean - 30.0).abs() < 0.1,
        "Mean should be ~30.0, got {}",
        mean
    );
    assert!(std > 0.0, "Std should be positive, got {}", std);

    // Verify normalized rewards have ~N(0,1) distribution
    let normalized_rewards: Vec<f64> = raw_rewards
        .iter()
        .map(|&r| normalizer.normalize(r))
        .collect();

    let norm_mean: f64 = normalized_rewards.iter().sum::<f64>() / normalized_rewards.len() as f64;
    let norm_var: f64 = normalized_rewards
        .iter()
        .map(|r| (r - norm_mean).powi(2))
        .sum::<f64>()
        / normalized_rewards.len() as f64;
    let norm_std = norm_var.sqrt();

    assert!(
        norm_mean.abs() < 0.5,
        "Normalized mean should be ~0, got {}",
        norm_mean
    );
    assert!(
        (norm_std - 1.0).abs() < 0.5,
        "Normalized std should be ~1.0, got {}",
        norm_std
    );

    // Create PPO config with reward normalization enabled
    let config = PPOConfig {
        state_dim: 64,
        num_actions: 3,
        ..Default::default()
    };

    let mut ppo = WorkingPPO::new(config)?;

    // Verify PPO has RewardNormalizer field (will fail until implementation)
    assert!(
        ppo.reward_normalizer.is_some(),
        "PPO should have RewardNormalizer field"
    );

    // Simulate training with reward normalization
    let state = vec![0.5; 64];
    let mut trajectory = Trajectory::new();
    for i in 0..10 {
        let (action, value) = ppo.act(&state)?;
        let reward = 10.0 + i as f32; // Non-normalized rewards (10.0 to 19.0)
        let step = TrajectoryStep::new(state.clone(), action, -1.0, value, reward, false);
        trajectory.add_step(step);
    }

    let mut batch = TrajectoryBatch::from_trajectories(
        vec![trajectory],
        vec![0.1; 10],
        vec![0.2; 10],
    );

    // Update should normalize rewards internally
    let (_policy_loss, _value_loss) = ppo.update(&mut batch)?;

    // Verify normalizer statistics were updated
    let (reward_mean, reward_std) = ppo.reward_normalizer.as_ref().unwrap().get_stats();
    assert!(
        reward_mean > 0.0,
        "Reward mean should be positive after training"
    );
    assert!(
        reward_std > 0.0,
        "Reward std should be positive after training"
    );

    Ok(())
}

/// Test 4: PPO Transaction Cost Integration
///
/// **Goal**: Verify transaction costs are deducted from rewards
/// **Expected Behavior**:
/// - Buy/Sell actions incur transaction costs (0.05-0.15%)
/// - Hold actions incur no transaction costs
/// - Transaction costs reduce net rewards
#[test]
fn test_ppo_transaction_cost_integration() -> Result<()> {
    // Create PPO config with transaction costs enabled
    let config = PPOConfig {
        state_dim: 64,
        num_actions: 3,
        policy_learning_rate: 1e-4,
        value_learning_rate: 1e-3,
        ..Default::default()
    };

    let mut ppo = WorkingPPO::new(config)?;

    // Access transaction cost configuration (will fail until implementation)
    assert!(
        ppo.transaction_cost_bps.is_some(),
        "PPO should have transaction_cost_bps field"
    );
    let tx_cost_bps = ppo.transaction_cost_bps.unwrap();
    assert!(
        tx_cost_bps >= 0.05 && tx_cost_bps <= 0.15,
        "Transaction cost should be 0.05-0.15%, got {}",
        tx_cost_bps
    );

    // Simulate Buy action
    let state = vec![0.5; 64];
    let mut trajectory = Trajectory::new();

    // Force Buy action and measure reward
    let (_, value) = ppo.act(&state)?;
    let buy_reward_before_cost = 1.0; // Hypothetical reward before costs
    let buy_step = TrajectoryStep::new(
        state.clone(),
        TradingAction::Buy,
        -1.0,
        value,
        buy_reward_before_cost,
        false,
    );
    trajectory.add_step(buy_step);

    // Force Hold action and measure reward
    let hold_reward_before_cost = 0.0; // Hold gets no P&L
    let hold_step = TrajectoryStep::new(
        state.clone(),
        TradingAction::Hold,
        -1.0,
        value,
        hold_reward_before_cost,
        false,
    );
    trajectory.add_step(hold_step);

    let mut batch = TrajectoryBatch::from_trajectories(
        vec![trajectory],
        vec![0.1, 0.0],
        vec![0.2, 0.0],
    );

    // Update applies transaction costs internally
    let (_policy_loss, _value_loss) = ppo.update(&mut batch)?;

    // Verify transaction costs were applied
    // Buy action should have reward reduced by transaction cost
    // Hold action should have no transaction cost
    let buy_reward_after_cost = batch.rewards[0];
    let hold_reward_after_cost = batch.rewards[1];

    assert!(
        buy_reward_after_cost < buy_reward_before_cost,
        "Buy action reward should be reduced by transaction costs"
    );
    assert_eq!(
        hold_reward_after_cost, hold_reward_before_cost,
        "Hold action reward should not have transaction costs"
    );

    Ok(())
}

/// Test 5: PPO Position Limit Integration
///
/// **Goal**: Verify position limits are enforced during action execution
/// **Expected Behavior**:
/// - Position size is clamped to [-max_position, +max_position]
/// - Actions that would exceed limits are rejected or reduced
/// - Position limit violations are logged
#[test]
fn test_ppo_position_limit_integration() -> Result<()> {
    // Create PPO config with strict position limits
    let config = PPOConfig {
        state_dim: 64,
        num_actions: 3,
        policy_learning_rate: 1e-4,
        value_learning_rate: 1e-3,
        ..Default::default()
    };

    let mut ppo = WorkingPPO::new(config)?;

    // Access position limit configuration (will fail until implementation)
    assert!(
        ppo.max_position_absolute.is_some(),
        "PPO should have max_position_absolute field"
    );
    let max_position = ppo.max_position_absolute.unwrap();
    assert!(
        max_position > 0.0,
        "Max position should be positive, got {}",
        max_position
    );

    // Simulate multiple Buy actions to test position limit
    let state = vec![0.5; 64];
    for _ in 0..10 {
        let (_action, _value) = ppo.act(&state)?;
        // Force Buy action
        let current_position = ppo.portfolio_tracker.current_position();

        // Verify position never exceeds max_position
        assert!(
            current_position.abs() <= max_position as f32,
            "Position {} should not exceed max_position {}",
            current_position,
            max_position
        );
    }

    Ok(())
}