foxhunt/ml/tests/ppo_action_space_tests.rs

//! Tests for PPO Action Space Abstraction
//!
//! This test suite validates the unified action space interface for both
//! discrete and continuous PPO implementations.

use candle_core::Device;
use ml::ppo::{
    ActionSpace, ActionType, ContinuousPPOConfig, ExposureLevel, FactoredAction, OrderType,
    PPOConfig, Urgency, ContinuousPolicyConfig, UnifiedPPO, UnifiedPPOConfig,
};
use ml::MLError;

#[test]
fn test_action_space_discrete_creation() -> Result<(), MLError> {
    let action = FactoredAction::new(ExposureLevel::Long100, OrderType::Market, Urgency::Normal);
    let action_space = ActionSpace::discrete(action);

    assert_eq!(action_space.action_type(), ActionType::Discrete);
    assert!(action_space.as_discrete().is_some());
    assert!(action_space.as_continuous().is_none());

    Ok(())
}

#[test]
fn test_action_space_continuous_creation() -> Result<(), MLError> {
    use ml::ppo::ContinuousAction;

    let action = ContinuousAction::new(0.75);
    let action_space = ActionSpace::continuous(action);

    assert_eq!(action_space.action_type(), ActionType::Continuous);
    assert!(action_space.as_continuous().is_some());
    assert!(action_space.as_discrete().is_none());

    Ok(())
}

#[test]
fn test_action_space_tensor_conversion_discrete() -> Result<(), MLError> {
    let device = Device::Cpu;
    let action = FactoredAction::new(ExposureLevel::Long50, OrderType::LimitMaker, Urgency::Aggressive);
    let action_space = ActionSpace::discrete(action);

    // Convert to tensor
    let tensor = action_space.to_tensor(&device)?;

    // Convert back
    let recovered = ActionSpace::from_tensor(&tensor, ActionType::Discrete)?;

    assert_eq!(action_space, recovered);

    Ok(())
}

#[test]
fn test_action_space_tensor_conversion_continuous() -> Result<(), MLError> {
    use ml::ppo::ContinuousAction;

    let device = Device::Cpu;
    let action = ContinuousAction::new(0.65);
    let action_space = ActionSpace::continuous(action);

    // Convert to tensor
    let tensor = action_space.to_tensor(&device)?;

    // Convert back
    let recovered = ActionSpace::from_tensor(&tensor, ActionType::Continuous)?;

    // Check values match (allowing for floating point error)
    let original_pos = action_space.as_continuous().unwrap().position_size();
    let recovered_pos = recovered.as_continuous().unwrap().position_size();
    assert!((original_pos - recovered_pos).abs() < 1e-6);

    Ok(())
}

#[test]
fn test_unified_ppo_discrete_creation() -> Result<(), MLError> {
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 45,
        policy_hidden_dims: vec![32],
        value_hidden_dims: vec![32],
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = UnifiedPPO::with_device(UnifiedPPOConfig::Discrete(config), device)?;

    assert_eq!(ppo.action_type(), ActionType::Discrete);

    Ok(())
}

#[test]
fn test_unified_ppo_continuous_creation() -> Result<(), MLError> {
    let config = ContinuousPPOConfig {
        state_dim: 16,
        policy_config: ContinuousPolicyConfig {
            state_dim: 16,
            hidden_dims: vec![32],
            ..Default::default()
        },
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = UnifiedPPO::with_device(UnifiedPPOConfig::Continuous(config), device)?;

    assert_eq!(ppo.action_type(), ActionType::Continuous);

    Ok(())
}

#[test]
fn test_unified_ppo_discrete_action_sampling() -> Result<(), MLError> {
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 45,
        ..Default::default()
    };

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

    // Create dummy state
    let state = candle_core::Tensor::zeros((1, 16), candle_core::DType::F32, &device)?;

    // Sample action
    let action = ppo.act(&state)?;

    assert_eq!(action.action_type(), ActionType::Discrete);
    assert!(action.as_discrete().is_some());
    assert!(action.is_valid());

    Ok(())
}

#[test]
fn test_unified_ppo_continuous_action_sampling() -> Result<(), MLError> {
    let config = ContinuousPPOConfig {
        state_dim: 16,
        policy_config: ContinuousPolicyConfig {
            state_dim: 16,
            ..Default::default()
        },
        ..Default::default()
    };

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

    // Create dummy state
    let state = candle_core::Tensor::zeros((1, 16), candle_core::DType::F32, &device)?;

    // Sample action
    let action = ppo.act(&state)?;

    assert_eq!(action.action_type(), ActionType::Continuous);
    assert!(action.as_continuous().is_some());
    assert!(action.is_valid());

    Ok(())
}

#[test]
fn test_unified_ppo_config_switching() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create discrete PPO
    let discrete_config = PPOConfig {
        state_dim: 16,
        num_actions: 45,
        ..Default::default()
    };
    let discrete_ppo = UnifiedPPO::with_device(
        UnifiedPPOConfig::Discrete(discrete_config),
        device.clone(),
    )?;
    assert_eq!(discrete_ppo.action_type(), ActionType::Discrete);

    // Create continuous PPO
    let continuous_config = ContinuousPPOConfig {
        state_dim: 16,
        policy_config: ContinuousPolicyConfig {
            state_dim: 16,
            ..Default::default()
        },
        ..Default::default()
    };
    let continuous_ppo = UnifiedPPO::with_device(
        UnifiedPPOConfig::Continuous(continuous_config),
        device,
    )?;
    assert_eq!(continuous_ppo.action_type(), ActionType::Continuous);

    Ok(())
}

#[test]
fn test_action_space_validation() {
    // Valid discrete action
    let valid_discrete = ActionSpace::discrete(FactoredAction::from_index(0).unwrap());
    assert!(valid_discrete.is_valid());

    // All 45 discrete actions should be valid
    for idx in 0..45 {
        let action = ActionSpace::discrete(FactoredAction::from_index(idx).unwrap());
        assert!(action.is_valid());
    }

    // Valid continuous action
    use ml::ppo::ContinuousAction;
    let valid_continuous = ActionSpace::continuous(ContinuousAction::new(0.5));
    assert!(valid_continuous.is_valid());

    // Clamped continuous action (still valid)
    let clamped_continuous = ActionSpace::continuous(ContinuousAction::new(1.5));
    assert!(clamped_continuous.is_valid());
}

#[test]
fn test_action_space_description() {
    let discrete = ActionSpace::discrete(FactoredAction::new(
        ExposureLevel::Long100,
        OrderType::Market,
        Urgency::Aggressive,
    ));
    let desc = discrete.description();
    assert!(desc.contains("Long100"));

    use ml::ppo::ContinuousAction;
    let continuous = ActionSpace::continuous(ContinuousAction::new(0.75));
    let desc = continuous.description();
    assert!(desc.contains("75"));
}

#[test]
fn test_unified_ppo_save_discrete() -> Result<(), MLError> {
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 45,
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = UnifiedPPO::with_device(UnifiedPPOConfig::Discrete(config), device)?;

    // Save checkpoint
    let checkpoint_path = "/tmp/test_unified_ppo_discrete";
    ppo.save(checkpoint_path)?;

    // Verify files exist
    assert!(std::path::Path::new(&format!("{}_actor.safetensors", checkpoint_path)).exists());
    assert!(std::path::Path::new(&format!("{}_critic.safetensors", checkpoint_path)).exists());
    assert!(std::path::Path::new(&format!("{}_metadata.json", checkpoint_path)).exists());

    // Clean up
    std::fs::remove_file(format!("{}_actor.safetensors", checkpoint_path)).ok();
    std::fs::remove_file(format!("{}_critic.safetensors", checkpoint_path)).ok();
    std::fs::remove_file(format!("{}_metadata.json", checkpoint_path)).ok();

    Ok(())
}

#[test]
fn test_unified_ppo_save_continuous() -> Result<(), MLError> {
    let config = ContinuousPPOConfig {
        state_dim: 16,
        policy_config: ContinuousPolicyConfig {
            state_dim: 16,
            ..Default::default()
        },
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = UnifiedPPO::with_device(UnifiedPPOConfig::Continuous(config), device)?;

    // Save checkpoint
    let checkpoint_path = "/tmp/test_unified_ppo_continuous";
    ppo.save(checkpoint_path)?;

    // Verify files exist
    assert!(std::path::Path::new(&format!("{}_policy.safetensors", checkpoint_path)).exists());
    assert!(std::path::Path::new(&format!("{}_value.safetensors", checkpoint_path)).exists());
    assert!(std::path::Path::new(&format!("{}_metadata.json", checkpoint_path)).exists());

    // Clean up
    std::fs::remove_file(format!("{}_policy.safetensors", checkpoint_path)).ok();
    std::fs::remove_file(format!("{}_value.safetensors", checkpoint_path)).ok();
    std::fs::remove_file(format!("{}_metadata.json", checkpoint_path)).ok();

    Ok(())
}

#[test]
fn test_unified_ppo_load_discrete() -> Result<(), MLError> {
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 45,
        ..Default::default()
    };

    let device = Device::Cpu;
    let ppo = UnifiedPPO::with_device(UnifiedPPOConfig::Discrete(config), device)?;

    // Save checkpoint
    let checkpoint_path = "/tmp/test_unified_ppo_load_discrete";
    ppo.save(checkpoint_path)?;

    // Load checkpoint
    let loaded_ppo = UnifiedPPO::load(checkpoint_path)?;
    assert_eq!(loaded_ppo.action_type(), ActionType::Discrete);

    // Clean up
    std::fs::remove_file(format!("{}_actor.safetensors", checkpoint_path)).ok();
    std::fs::remove_file(format!("{}_critic.safetensors", checkpoint_path)).ok();
    std::fs::remove_file(format!("{}_metadata.json", checkpoint_path)).ok();

    Ok(())
}