foxhunt/ml/tests/ppo_45_action_network_tests.rs

//! PPO 45-Action Network Tests (TDD Red Phase)
//!
//! Tests for PPO network architecture expansion from 3 → 45 actions.
//!
//! Test coverage:
//! 1. Policy network outputs 45 logits (not 3)
//! 2. Value network still outputs 1 value (unchanged)
//! 3. Softmax over 45 logits sums to 1.0
//! 4. Full forward pass with 45 actions works correctly
//!
//! Expected behavior (TDD Red → Green):
//! - RED: Tests fail because num_actions=3 (current implementation)
//! - GREEN: Tests pass after num_actions=45 changes

use anyhow::Result;
use candle_core::{Device, Tensor};

#[test]
fn test_policy_network_45_output() -> Result<()> {
    use ml::ppo::ppo::{PPOConfig, PPO};

    // Create PPO with 45 actions
    let config = PPOConfig {
        num_actions: 45, // 5×3×3 factored action space
        state_dim: 54,   // Wave 3 features (54→54)
        ..Default::default()
    };

    let ppo = PPO::new(config)?;

    // Create dummy state
    let state = Tensor::zeros(&[1, 54], candle_core::DType::F32, &Device::Cpu)?;

    // Forward pass through policy network
    let logits = ppo.actor.forward(&state)?;

    // ASSERT: Policy network outputs 45 logits (not 3)
    assert_eq!(
        logits.dims(),
        &[1, 45],
        "Policy network should output 45 logits, got: {:?}",
        logits.dims()
    );

    Ok(())
}

#[test]
fn test_value_network_single_output() -> Result<()> {
    use ml::ppo::ppo::{PPOConfig, PPO};

    // Create PPO with 45 actions
    let config = PPOConfig {
        num_actions: 45, // 5×3×3 factored action space
        state_dim: 54,   // Wave 3 features (54→54)
        ..Default::default()
    };

    let ppo = PPO::new(config)?;

    // Create dummy state
    let state = Tensor::zeros(&[1, 54], candle_core::DType::F32, &Device::Cpu)?;

    // Forward pass through value network
    let value = ppo.critic.forward(&state)?;

    // ASSERT: Value network still outputs 1 value (batch_size=1, scalar output)
    // Note: critic.forward() squeezes to [batch_size], so we expect [1] not [1, 1]
    assert_eq!(
        value.dims(),
        &[1],
        "Value network should output single value per batch item, got: {:?}",
        value.dims()
    );

    Ok(())
}

#[test]
fn test_policy_network_softmax() -> Result<()> {
    use ml::ppo::ppo::{PPOConfig, PPO};

    // Create PPO with 45 actions
    let config = PPOConfig {
        num_actions: 45,
        state_dim: 54,
        ..Default::default()
    };

    let ppo = PPO::new(config)?;

    // Create dummy state
    let state = Tensor::zeros(&[1, 54], candle_core::DType::F32, &Device::Cpu)?;

    // Get action probabilities (softmax of logits)
    let probs = ppo.actor.action_probabilities(&state)?;

    // ASSERT 1: Probabilities have 45 dimensions
    assert_eq!(
        probs.dims(),
        &[1, 45],
        "Action probabilities should have 45 dimensions, got: {:?}",
        probs.dims()
    );

    // ASSERT 2: Probabilities sum to 1.0 (softmax property)
    let probs_vec = probs.flatten_all()?.to_vec1::<f32>()?;
    let sum: f32 = probs_vec.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-5,
        "Softmax probabilities should sum to 1.0, got: {}",
        sum
    );

    // ASSERT 3: Each probability is in [0, 1]
    for (i, &prob) in probs_vec.iter().enumerate() {
        assert!(
            prob >= 0.0 && prob <= 1.0,
            "Probability {} is out of range [0, 1]: {}",
            i,
            prob
        );
    }

    Ok(())
}

#[test]
fn test_network_forward_pass() -> Result<()> {
    use ml::ppo::ppo::{PPOConfig, PPO};

    // Create PPO with 45 actions
    let config = PPOConfig {
        num_actions: 45,
        state_dim: 54,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![256, 128, 64],
        ..Default::default()
    };

    let ppo = PPO::new(config)?;

    // Create batch of states (batch_size=8)
    let batch_size = 8;
    let state = Tensor::zeros(&[batch_size, 54], candle_core::DType::F32, &Device::Cpu)?;

    // ASSERT 1: Policy forward pass with batch
    let logits = ppo.actor.forward(&state)?;
    assert_eq!(
        logits.dims(),
        &[batch_size, 45],
        "Policy logits should have shape [batch_size, 45], got: {:?}",
        logits.dims()
    );

    // ASSERT 2: Value forward pass with batch
    let values = ppo.critic.forward(&state)?;
    assert_eq!(
        values.dims(),
        &[batch_size],
        "Value network should output [batch_size] values, got: {:?}",
        values.dims()
    );

    // ASSERT 3: Action probabilities with batch
    let probs = ppo.actor.action_probabilities(&state)?;
    assert_eq!(
        probs.dims(),
        &[batch_size, 45],
        "Action probabilities should have shape [batch_size, 45], got: {:?}",
        probs.dims()
    );

    // ASSERT 4: Each batch item's probabilities sum to 1.0
    let probs_vec = probs.to_vec2::<f32>()?;
    for (i, row) in probs_vec.iter().enumerate() {
        let sum: f32 = row.iter().sum();
        assert!(
            (sum - 1.0).abs() < 1e-5,
            "Batch item {} probabilities should sum to 1.0, got: {}",
            i,
            sum
        );
    }

    Ok(())
}

#[test]
fn test_backward_compatibility_3_actions() -> Result<()> {
    use ml::ppo::ppo::{PPOConfig, PPO};

    // Test that 3-action configuration still works (backward compatibility)
    let config = PPOConfig {
        num_actions: 3,
        state_dim: 54,
        ..Default::default()
    };

    let ppo = PPO::new(config)?;

    // Create dummy state
    let state = Tensor::zeros(&[1, 54], candle_core::DType::F32, &Device::Cpu)?;

    // Policy network should output 3 logits
    let logits = ppo.actor.forward(&state)?;
    assert_eq!(
        logits.dims(),
        &[1, 3],
        "Policy network with num_actions=3 should output 3 logits"
    );

    // Value network should still output 1 value
    let value = ppo.critic.forward(&state)?;
    assert_eq!(
        value.dims(),
        &[1],
        "Value network should output single value"
    );

    Ok(())
}

#[test]
fn test_ppo_config_default_45_actions() -> Result<()> {
    use ml::ppo::ppo::PPOConfig;

    // ASSERT: Default config uses 45 actions (not 3)
    let config = PPOConfig::default();
    assert_eq!(
        config.num_actions, 45,
        "Default PPOConfig should have num_actions=45, got: {}",
        config.num_actions
    );

    Ok(())
}

#[test]
fn test_ppo_trainer_default_45_actions() -> Result<()> {
    use ml::trainers::ppo::PpoHyperparameters;

    // ASSERT: Conservative params use 45 actions
    let params = PpoHyperparameters::conservative();
    let config: ml::ppo::ppo::PPOConfig = params.into();

    assert_eq!(
        config.num_actions, 45,
        "PpoHyperparameters should default to 45 actions, got: {}",
        config.num_actions
    );

    Ok(())
}

#[test]
fn test_hyperopt_adapter_default_45_actions() -> Result<()> {
    use ml::hyperopt::adapters::ppo::PPOParams;

    // ASSERT: PPOParams default uses 45 actions when converted to config
    let params = PPOParams::default();

    // Create config manually (since PPOParams → PPOConfig conversion is in PPOTrainer)
    // We verify the architecture matches 45-action space
    use ml::ppo::ppo::PPOConfig;
    use ml::ppo::gae::GAEConfig;

    let config = PPOConfig {
        state_dim: 54,
        num_actions: 45, // This is what hyperopt adapter should use
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![512, 384, 256, 128, 64],
        policy_learning_rate: params.policy_learning_rate,
        value_learning_rate: params.value_learning_rate,
        clip_epsilon: params.clip_epsilon as f32,
        value_loss_coeff: params.value_loss_coeff as f32,
        entropy_coeff: params.entropy_coeff as f32,
        gae_config: GAEConfig::default(),
        batch_size: 2048,
        mini_batch_size: 512,
        num_epochs: 20,
        max_grad_norm: 0.5,
        early_stopping_enabled: true,
        early_stopping_patience: 5,
        early_stopping_min_delta: 1e-4,
        early_stopping_min_epochs: 5,
        max_position_absolute: 2.0,
        transaction_cost_bps: 0.10,
        cash_reserve_pct: 20.0,
        circuit_breaker_threshold: 5,
        use_lstm: false,
        lstm_hidden_dim: 128,
        lstm_num_layers: 1,
        lstm_sequence_length: 32,
        accumulation_steps: 1,
        clip_epsilon_high: None,
    };

    assert_eq!(
        config.num_actions, 45,
        "Hyperopt adapter config should use 45 actions"
    );

    Ok(())
}