foxhunt/crates/ml-ppo/src/action_space.rs

//! PPO Action Space Abstraction
//!
//! This module provides a unified interface for both discrete and continuous action spaces.
//! It enables the same PPO framework to handle:
//! - Discrete actions: 63-action factored space (7x3x3 = exposure x order x urgency)
//! - Continuous actions: Gaussian policy for position sizing (0.0 to 1.0)

use ml_core::action_space::FactoredAction;
use crate::continuous_policy::ContinuousAction;
use ml_core::MLError;
use serde::{Deserialize, Serialize};

/// Action space type discriminator
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ActionType {
    /// Discrete action space (63-action factored space)
    Discrete,
    /// Continuous action space (Gaussian policy)
    Continuous,
}

/// Unified action space supporting both discrete and continuous actions
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ActionSpace {
    /// Discrete trading action (63-action factored space)
    Discrete(FactoredAction),
    /// Continuous trading action (position sizing)
    Continuous(ContinuousAction),
}

impl ActionSpace {
    /// Create discrete action from factored action
    pub const fn discrete(action: FactoredAction) -> Self {
        ActionSpace::Discrete(action)
    }

    /// Create continuous action from position size
    pub const fn continuous(action: ContinuousAction) -> Self {
        ActionSpace::Continuous(action)
    }

    /// Get action type
    pub const fn action_type(&self) -> ActionType {
        match self {
            ActionSpace::Discrete(_) => ActionType::Discrete,
            ActionSpace::Continuous(_) => ActionType::Continuous,
        }
    }

    /// Convert action to index (discrete only)
    pub fn to_index(&self) -> Result<usize, MLError> {
        match self {
            ActionSpace::Discrete(action) => Ok(action.to_index()),
            ActionSpace::Continuous(_) => Err(MLError::InvalidInput(
                "Cannot convert continuous action to index".to_owned(),
            )),
        }
    }

    /// Create action from index (discrete)
    pub fn from_index(idx: usize) -> Result<Self, MLError> {
        let action = FactoredAction::from_index(idx)?;
        Ok(ActionSpace::Discrete(action))
    }

    /// Get discrete action (if applicable)
    pub const fn as_discrete(&self) -> Option<&FactoredAction> {
        match self {
            ActionSpace::Discrete(action) => Some(action),
            ActionSpace::Continuous(_) => None,
        }
    }

    /// Get continuous action (if applicable)
    pub const fn as_continuous(&self) -> Option<&ContinuousAction> {
        match self {
            ActionSpace::Discrete(_) => None,
            ActionSpace::Continuous(action) => Some(action),
        }
    }

    /// Validate action is within bounds
    pub fn is_valid(&self) -> bool {
        match self {
            ActionSpace::Discrete(action) => action.to_index() < 63,
            ActionSpace::Continuous(action) => action.is_valid(),
        }
    }

    /// Get human-readable description
    pub fn description(&self) -> String {
        match self {
            ActionSpace::Discrete(action) => {
                format!("{}", action)
            }
            ActionSpace::Continuous(action) => {
                format!("Position: {:.2}%", action.position_size() * 100.0)
            }
        }
    }
}

impl std::fmt::Display for ActionSpace {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.description())
    }
}

impl Default for ActionSpace {
    fn default() -> Self {
        use ml_core::action_space::{ExposureLevel, OrderType, Urgency};
        ActionSpace::Discrete(FactoredAction::new(
            ExposureLevel::Flat,
            OrderType::Market,
            Urgency::Normal,
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ml_core::action_space::{ExposureLevel, OrderType, Urgency};

    #[test]
    fn test_action_type() {
        let discrete = ActionSpace::discrete(FactoredAction::new(
            ExposureLevel::LongFull,
            OrderType::Market,
            Urgency::Normal,
        ));
        assert_eq!(discrete.action_type(), ActionType::Discrete);

        let continuous = ActionSpace::continuous(ContinuousAction::new(0.5));
        assert_eq!(continuous.action_type(), ActionType::Continuous);
    }

    #[test]
    fn test_as_discrete() {
        let action = FactoredAction::new(
            ExposureLevel::ShortSmall,
            OrderType::IoC,
            Urgency::Patient,
        );
        let action_space = ActionSpace::discrete(action);

        assert!(action_space.as_discrete().is_some());
        assert!(action_space.as_continuous().is_none());
        assert_eq!(*action_space.as_discrete().unwrap(), action);
    }

    #[test]
    fn test_as_continuous() {
        let action = ContinuousAction::new(0.3);
        let action_space = ActionSpace::continuous(action);

        assert!(action_space.as_continuous().is_some());
        assert!(action_space.as_discrete().is_none());
        assert_eq!(*action_space.as_continuous().unwrap(), action);
    }

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

        let valid_continuous = ActionSpace::continuous(ContinuousAction::new(0.5));
        assert!(valid_continuous.is_valid());
    }

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

        let continuous = ActionSpace::continuous(ContinuousAction::new(0.65));
        let desc = continuous.description();
        assert!(desc.contains("65"));
    }

    #[test]
    fn test_all_discrete_actions() -> Result<(), MLError> {
        for idx in 0..63 {
            let action = FactoredAction::from_index(idx)?;
            let action_space = ActionSpace::discrete(action);
            assert!(action_space.is_valid());

            let recovered_idx = action_space.to_index()?;
            assert_eq!(recovered_idx, idx);
        }
        Ok(())
    }

    #[test]
    fn test_default() {
        let default_action = ActionSpace::default();
        assert_eq!(default_action.action_type(), ActionType::Discrete);

        let action = default_action.as_discrete().unwrap();
        assert!(action.is_hold());
    }
}