foxhunt/ml/tests/ppo_45_action_validation.rs

//! Test to validate PPO already supports 45-action factored space
//!
//! This test confirms that PPO network architecture and training
//! pipeline correctly handle the 45-action factored space:
//! - 5 exposure levels × 3 order types × 3 urgency levels = 45 actions
//!
//! Wave 9-A4: Verification that PPO is ready for factored actions

use anyhow::Result;
use candle_core::Device;

use ml::ppo::ppo::{PPOConfig, PPO};
use ml::ppo::trajectories::{Trajectory, TrajectoryStep};

#[test]
fn test_ppo_45_action_default_config() -> Result<()> {
    // Verify default config uses 45 actions
    let config = PPOConfig::default();
    assert_eq!(
        config.num_actions, 45,
        "PPO default config should use 45 actions (factored space)"
    );
    assert_eq!(config.state_dim, 64, "Default state dim should be 64");

    println!("✅ PASS: PPOConfig::default() uses 45 actions");
    Ok(())
}

#[test]
fn test_ppo_network_45_actions() -> Result<()> {
    // Create PPO with 45 actions
    let config = PPOConfig {
        state_dim: 128,
        num_actions: 45,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![256, 128, 64],
        policy_learning_rate: 3e-5,
        value_learning_rate: 1e-4,
        clip_epsilon: 0.2,
        value_loss_coeff: 1.0,
        entropy_coeff: 0.05,
        batch_size: 2048,
        mini_batch_size: 512,
        num_epochs: 20,
        max_grad_norm: 0.5,
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = PPO::with_device(config.clone(), device)?;

    // Verify action selection returns valid action indices (0-44)
    let test_state = vec![0.5f32; 128];
    let (action_idx, _value) = ppo.act(&test_state)?;

    assert!(
        action_idx < 45,
        "Action index {} should be < 45",
        action_idx
    );

    println!("✅ PASS: PPO network supports 45-action output");
    println!("   Sample action: {}", action_idx);

    Ok(())
}

#[test]
fn test_ppo_action_diversity_45_actions() -> Result<()> {
    // Verify PPO can sample all 45 actions over multiple episodes
    let config = PPOConfig {
        state_dim: 64,
        num_actions: 45,
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = PPO::with_device(config, device)?;

    let mut action_counts = vec![0usize; 45];

    // Sample 500 actions with different states
    for i in 0..500 {
        let test_state: Vec<f32> = (0..64).map(|j| (i + j) as f32 * 0.01).collect();
        let (action_idx, _value) = ppo.act(&test_state)?;
        assert!(action_idx < 45, "Action {} out of range", action_idx);
        action_counts[action_idx] += 1;
    }

    // Count how many unique actions were sampled
    let unique_actions = action_counts.iter().filter(|&&count| count > 0).count();

    println!(
        "✅ PASS: PPO sampled {} unique actions out of 45",
        unique_actions
    );
    println!(
        "   Coverage: {:.1}%",
        (unique_actions as f64 / 45.0) * 100.0
    );

    // We expect at least 30% coverage (13+ actions) over 500 samples
    assert!(
        unique_actions >= 13,
        "Expected at least 13 unique actions, got {}",
        unique_actions
    );

    Ok(())
}

#[test]
fn test_ppo_trajectory_45_actions() -> Result<()> {
    // Verify trajectory batch handles 45-action indices correctly
    let config = PPOConfig {
        state_dim: 64,
        num_actions: 45,
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = PPO::with_device(config, device)?;

    // Create trajectory with all 45 actions
    let mut trajectory = Trajectory::new();

    for action_idx in 0..45 {
        let state = vec![action_idx as f32 * 0.01; 64];
        let step = TrajectoryStep::new(
            state,
            action_idx,       // Action index 0-44
            -1.0,             // Log prob
            0.5,              // Value estimate
            0.1,              // Reward
            action_idx == 44, // Done on last action
        );
        trajectory.add_step(step);
    }

    // Verify all actions are valid
    let actions = trajectory.get_actions();
    assert_eq!(actions.len(), 45, "Should have 45 trajectory steps");

    for (i, &action) in actions.iter().enumerate() {
        assert_eq!(action, i, "Action index {} should match step {}", action, i);
    }

    println!("✅ PASS: Trajectories correctly store 45-action indices");
    Ok(())
}

#[test]
fn test_ppo_action_probabilities_sum_to_one() -> Result<()> {
    // Verify softmax probabilities sum to 1.0 for 45 actions
    let config = PPOConfig {
        state_dim: 64,
        num_actions: 45,
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = PPO::with_device(config, device)?;

    let test_state = vec![0.5f32; 64];
    let probs = ppo.predict(&test_state)?;

    assert_eq!(probs.len(), 45, "Should return 45 probabilities");

    let sum: f32 = probs.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-5,
        "Probabilities should sum to 1.0, got {}",
        sum
    );

    println!("✅ PASS: Action probabilities sum to 1.0");
    println!("   Probability sum: {:.6}", sum);

    Ok(())
}

#[test]
fn test_ppo_factored_action_mapping_documentation() -> Result<()> {
    // Document the factored action mapping
    println!("📋 Factored Action Space Mapping (45 actions):");
    println!();
    println!("   Exposure Levels (5): -100%, -50%, 0%, +50%, +100%");
    println!("   Order Types (3): Market, Limit, Stop");
    println!("   Urgency Levels (3): Low, Medium, High");
    println!();
    println!("   Total Actions: 5 × 3 × 3 = 45");
    println!();
    println!("   Example Mapping:");
    println!("     Action 0:  Exposure=-100%, Market, Low urgency");
    println!("     Action 1:  Exposure=-100%, Market, Medium urgency");
    println!("     Action 2:  Exposure=-100%, Market, High urgency");
    println!("     Action 3:  Exposure=-100%, Limit, Low urgency");
    println!("     ...");
    println!("     Action 44: Exposure=+100%, Stop, High urgency");
    println!();
    println!("✅ PASS: Factored action space documented");

    Ok(())
}

#[test]
fn test_ppo_config_from_hyperparameters() -> Result<()> {
    // Verify PpoHyperparameters → PPOConfig conversion preserves 45 actions
    use ml::trainers::ppo::PpoHyperparameters;

    let hyperparams = PpoHyperparameters::conservative();
    let config: ml::ppo::ppo::PPOConfig = hyperparams.into();

    assert_eq!(
        config.num_actions, 45,
        "PpoHyperparameters should convert to 45 actions"
    );
    assert_eq!(
        config.state_dim, 54,
        "State dimension should be 54 (Wave 3 features, 54→54)"
    );

    println!("✅ PASS: PpoHyperparameters converts to 45-action PPOConfig");
    println!("   State dim: {}", config.state_dim);
    println!("   Actions: {}", config.num_actions);

    Ok(())
}