//! Continuous Action Masking for PPO
//!
//! Provides state-dependent action bounds and constraint enforcement for continuous
//! position sizing. Unlike discrete action masking (which sets masked logits to -inf),
//! continuous masking adjusts the Gaussian distribution parameters to respect constraints.

use serde::{Deserialize, Serialize};

use ml_core::MLError;

/// Continuous action constraints based on current state
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContinuousActionConstraints {
    /// Minimum allowed position (e.g., -2.0 for max short)
    pub min_position: f32,
    /// Maximum allowed position (e.g., +2.0 for max long)
    pub max_position: f32,
    /// Soft constraint penalty coefficient
    pub penalty_coeff: f32,
    /// Soft threshold as fraction of `max_position`
    pub soft_threshold_fraction: f32,
}

impl ContinuousActionConstraints {
    /// Create constraints from current position and limits.
    ///
    /// # Arguments
    /// * `current_position` - Current portfolio position
    /// * `max_position_abs` - Maximum position magnitude (e.g., 2.0)
    pub fn from_position(current_position: f32, max_position_abs: f32) -> Self {
        let max_long = max_position_abs - current_position.max(0.0);
        let max_short = max_position_abs + current_position.min(0.0);

        Self {
            min_position: -max_short,
            max_position: max_long,
            penalty_coeff: 1.0,
            soft_threshold_fraction: 0.8,
        }
    }

    /// Create with custom parameters
    pub const fn new(
        min_position: f32,
        max_position: f32,
        penalty_coeff: f32,
        soft_threshold_fraction: f32,
    ) -> Self {
        Self {
            min_position,
            max_position,
            penalty_coeff,
            soft_threshold_fraction,
        }
    }

    /// Apply hard constraint (clip action to valid range)
    pub const fn clip_action(&self, action: f32) -> f32 {
        action.clamp(self.min_position, self.max_position)
    }

    /// Compute soft penalty for constraint violation
    pub fn compute_penalty(&self, action: f32) -> f32 {
        let soft_limit = self.max_position * self.soft_threshold_fraction;
        let violation = (action.abs() - soft_limit).max(0.0);
        self.penalty_coeff * violation * violation
    }

    /// Adjust Gaussian distribution to respect constraints (host-side).
    ///
    /// Returns (`adjusted_mean`, `adjusted_log_std`)
    pub fn adjust_distribution_host(
        &self,
        mean: f32,
        log_std: f32,
    ) -> (f32, f32) {
        let sigma = log_std.exp();
        let two_sigma = 2.0 * sigma;

        // Shift mean if distribution tails exceed bounds
        let mut adjusted_mean = mean;
        if mean + two_sigma > self.max_position {
            adjusted_mean = self.max_position - two_sigma;
        }
        if adjusted_mean - two_sigma < self.min_position {
            adjusted_mean = self.min_position + two_sigma;
        }

        // Limit sigma to quarter of allowed range
        let max_sigma = (self.max_position - self.min_position) / 4.0;
        let adjusted_sigma = sigma.min(max_sigma);
        let adjusted_log_std = adjusted_sigma.ln();

        (adjusted_mean, adjusted_log_std)
    }

    /// Get effective action range (for normalization)
    pub fn action_range(&self) -> f32 {
        self.max_position - self.min_position
    }

    /// Check if action is within hard bounds
    pub fn is_valid(&self, action: f32) -> bool {
        action >= self.min_position && action <= self.max_position
    }

    /// Check if action is within soft bounds
    pub fn is_within_soft_bounds(&self, action: f32) -> bool {
        let soft_limit = self.max_position * self.soft_threshold_fraction;
        action.abs() <= soft_limit
    }
}

/// Apply continuous action masking to policy network output (host-side).
///
/// Returns (`masked_mean`, `masked_log_std`) that respect constraints.
pub fn mask_continuous_actions_host(
    mean: f32,
    log_std: f32,
    constraints: &ContinuousActionConstraints,
) -> Result<(f32, f32), MLError> {
    let (adj_mean, adj_log_std) = constraints.adjust_distribution_host(mean, log_std);
    Ok((adj_mean, adj_log_std))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_from_position_flat() {
        let constraints = ContinuousActionConstraints::from_position(0.0, 2.0);
        assert!((constraints.max_position - 2.0).abs() < 1e-6);
        assert!((constraints.min_position - (-2.0)).abs() < 1e-6);
    }

    #[test]
    fn test_from_position_long() {
        let constraints = ContinuousActionConstraints::from_position(1.5, 2.0);
        assert!((constraints.max_position - 0.5).abs() < 1e-6);
        assert!((constraints.min_position - (-2.0)).abs() < 1e-6);
    }

    #[test]
    fn test_from_position_short() {
        let constraints = ContinuousActionConstraints::from_position(-1.0, 2.0);
        assert!((constraints.max_position - 2.0).abs() < 1e-6);
        assert!((constraints.min_position - (-1.0)).abs() < 1e-6);
    }

    #[test]
    fn test_clip_action_basic() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        assert_eq!(constraints.clip_action(2.5), 2.0);
        assert_eq!(constraints.clip_action(-3.0), -2.0);
        assert_eq!(constraints.clip_action(1.0), 1.0);
    }

    #[test]
    fn test_compute_penalty_within_soft_bounds() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        assert_eq!(constraints.compute_penalty(1.0), 0.0);
        assert_eq!(constraints.compute_penalty(-1.0), 0.0);
        assert_eq!(constraints.compute_penalty(1.6), 0.0);
    }

    #[test]
    fn test_compute_penalty_beyond_soft_bounds() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        let penalty = constraints.compute_penalty(1.8);
        let expected = 1.0 * 0.2 * 0.2;
        assert!((penalty - expected).abs() < 1e-6);
    }

    #[test]
    fn test_action_range() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        assert_eq!(constraints.action_range(), 4.0);
    }

    #[test]
    fn test_is_valid() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        assert!(constraints.is_valid(1.0));
        assert!(constraints.is_valid(2.0));
        assert!(!constraints.is_valid(2.1));
    }

    #[test]
    fn test_is_within_soft_bounds() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        assert!(constraints.is_within_soft_bounds(1.0));
        assert!(constraints.is_within_soft_bounds(1.6));
        assert!(!constraints.is_within_soft_bounds(1.7));
    }

    #[test]
    fn test_adjust_distribution_no_adjustment() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        let (adj_mean, adj_log_std) = constraints.adjust_distribution_host(0.0, -1.0);
        assert!((adj_mean - 0.0).abs() < 0.01);
        assert!((adj_log_std - (-1.0)).abs() < 0.01);
    }

    #[test]
    fn test_adjust_distribution_mean_too_high() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        let log_std = -0.693_f32; // std ~ 0.5
        let (adj_mean, _) = constraints.adjust_distribution_host(2.0, log_std);
        // Adjusted mean should be 2.0 - 2*0.5 = 1.0
        assert!((adj_mean - 1.0).abs() < 0.1);
    }

    #[test]
    fn test_mask_continuous_actions_integration() {
        let constraints = ContinuousActionConstraints::new(-2.0, 2.0, 1.0, 0.8);
        let (masked_mean, _) = mask_continuous_actions_host(2.0, -0.693, &constraints).unwrap();
        assert!(masked_mean < 2.0);
    }
}
