//! Action Masking for PPO
//!
//! Prevents the agent from selecting actions that would violate position limits.
//! This is adapted from DQN's action masking but works with PPO's policy network.
//!
//! # Key Differences from DQN:
//! - **DQN**: Masks Q-values after forward pass, selects argmax from unmasked actions
//! - **PPO**: Masks logits BEFORE softmax, samples from masked probability distribution
//!
//! # 3-Action Space:
//! - Action 0: BUY (increase position by +1.0)
//! - Action 1: SELL (decrease position by -1.0)
//! - Action 2: HOLD (keep position unchanged)
//!
//! # 45-Action Factored Space (5 exposure x 3 order x 3 urgency):
//! Index = exposure * 9 + order * 3 + urgency
//! - `exposure_idx` = idx / 9
//!   - 0 = Short100 (indices 0..9)
//!   - 1 = Short50  (indices 9..18)
//!   - 2 = Flat     (indices 18..27)
//!   - 3 = Long50   (indices 27..36)
//!   - 4 = Long100  (indices 36..45)
//! - Masking: Long50/Long100 masked at max position, Short100/Short50 masked at min position

/// Creates an action mask based on current position and position limits.
///
/// # Arguments
/// * `current_position` - Current portfolio position
/// * `max_position` - Maximum allowed position magnitude (typically 2.0)
/// * `num_actions` - Number of actions in the action space (3 or 45)
///
/// # Returns
/// Boolean mask where `true` = action is valid, `false` = action is invalid
pub fn create_action_mask(current_position: f64, max_position: f64, num_actions: usize) -> Vec<bool> {
    let mut mask = vec![true; num_actions];

    if num_actions == 3 {
        if current_position + 1.0 > max_position {
            mask[0] = false;
        }
        if current_position - 1.0 < -max_position {
            mask[1] = false;
        }
    } else if num_actions == 45 {
        for action_idx in 0..num_actions {
            let exposure_idx = action_idx / 9;
            match exposure_idx {
                3 | 4 => {
                    if current_position >= max_position {
                        if let Some(m) = mask.get_mut(action_idx) {
                            *m = false;
                        }
                    }
                }
                0 | 1 => {
                    if current_position <= -max_position {
                        if let Some(m) = mask.get_mut(action_idx) {
                            *m = false;
                        }
                    }
                }
                _ => {}
            }
        }
    } else {
        // Unsupported action space size, all actions remain valid
    }

    mask
}

/// Applies action mask to logits by setting masked actions to -1e9.
///
/// Invalid actions (mask[i] = false) have their logits set to -1e9 (effectively -inf).
/// This ensures that after softmax, masked actions have probability ~0.
///
/// # Arguments
/// * `logits` - Raw logits from policy network (shape: [`num_actions`])
/// * `mask` - Boolean mask where false = invalid action
///
/// # Returns
/// Masked logits where invalid actions have logit = -1e9
pub fn apply_mask_to_logits_vec(logits: &[f32], mask: &[bool]) -> Vec<f32> {
    logits
        .iter()
        .zip(mask.iter())
        .map(|(&logit, &valid)| if valid { logit } else { -1e9 })
        .collect()
}

#[cfg(test)]
#[allow(
    clippy::bool_assert_comparison,
    clippy::wildcard_enum_match_arm
)]
mod tests {
    use super::*;

    #[test]
    fn test_create_action_mask_flat_position() {
        let mask = create_action_mask(0.0, 2.0, 3);
        assert_eq!(mask, vec![true, true, true]);
    }

    #[test]
    fn test_create_action_mask_at_max_position() {
        let mask = create_action_mask(2.0, 2.0, 3);
        assert_eq!(mask[0], false);
        assert_eq!(mask[1], true);
        assert_eq!(mask[2], true);
    }

    #[test]
    fn test_create_action_mask_at_min_position() {
        let mask = create_action_mask(-2.0, 2.0, 3);
        assert_eq!(mask[0], true);
        assert_eq!(mask[1], false);
        assert_eq!(mask[2], true);
    }

    #[test]
    fn test_apply_mask_to_logits() {
        let logits = vec![0.5_f32, 0.3, 0.2];
        let mask = vec![false, true, true];

        let masked = apply_mask_to_logits_vec(&logits, &mask);

        assert!(masked[0] < -1e8);
        assert_eq!(masked[1], 0.3);
        assert_eq!(masked[2], 0.2);
    }

    #[test]
    fn test_apply_mask_all_valid() {
        let logits = vec![0.5_f32, 0.3, 0.2];
        let mask = vec![true, true, true];

        let masked = apply_mask_to_logits_vec(&logits, &mask);

        assert_eq!(masked[0], 0.5);
        assert_eq!(masked[1], 0.3);
        assert_eq!(masked[2], 0.2);
    }

    #[test]
    fn test_create_action_mask_45_flat_position() {
        let mask = create_action_mask(0.0, 2.0, 45);
        assert_eq!(mask.len(), 45);
        assert!(mask.iter().all(|&v| v));
    }

    #[test]
    fn test_create_action_mask_45_at_max_position() {
        let mask = create_action_mask(2.0, 2.0, 45);
        assert_eq!(mask.len(), 45);
        for i in 0..9 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
        for i in 9..18 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
        for i in 18..27 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
        for i in 27..36 {
            assert_eq!(mask.get(i).copied().unwrap_or(true), false);
        }
        for i in 36..45 {
            assert_eq!(mask.get(i).copied().unwrap_or(true), false);
        }
    }

    #[test]
    fn test_create_action_mask_45_at_min_position() {
        let mask = create_action_mask(-2.0, 2.0, 45);
        for i in 0..9 {
            assert_eq!(mask.get(i).copied().unwrap_or(true), false);
        }
        for i in 9..18 {
            assert_eq!(mask.get(i).copied().unwrap_or(true), false);
        }
        for i in 18..27 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
        for i in 27..36 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
        for i in 36..45 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
    }

    #[test]
    fn test_create_action_mask_45_partial_position() {
        let mask = create_action_mask(1.0, 2.0, 45);
        assert!(mask.iter().all(|&v| v));
    }

    #[test]
    fn test_apply_mask_all_invalid() {
        let logits = vec![0.5_f32, 0.3, 0.2];
        let mask = vec![false, false, false];

        let masked = apply_mask_to_logits_vec(&logits, &mask);

        assert!(masked[0] < -1e8);
        assert!(masked[1] < -1e8);
        assert!(masked[2] < -1e8);
    }

    #[test]
    fn test_45_action_masking_canonical_layout() {
        use ml_core::action_space::{ExposureLevel, FactoredAction};

        let max_position = 5.0;

        let mask = create_action_mask(max_position, max_position, 45);
        for idx in 0..45 {
            if let Ok(action) = FactoredAction::from_index(idx) {
                match action.exposure {
                    ExposureLevel::Long50 | ExposureLevel::Long100 => {
                        assert!(
                            !mask.get(idx).copied().unwrap_or(true),
                            "Long action {} should be masked at max position",
                            idx
                        );
                    }
                    _ => {
                        assert!(
                            mask.get(idx).copied().unwrap_or(false),
                            "Non-long action {} should NOT be masked at max position",
                            idx
                        );
                    }
                }
            }
        }

        let mask = create_action_mask(-max_position, max_position, 45);
        for idx in 0..45 {
            if let Ok(action) = FactoredAction::from_index(idx) {
                match action.exposure {
                    ExposureLevel::Short50 | ExposureLevel::Short100 => {
                        assert!(
                            !mask.get(idx).copied().unwrap_or(true),
                            "Short action {} should be masked at min position",
                            idx
                        );
                    }
                    _ => {
                        assert!(
                            mask.get(idx).copied().unwrap_or(false),
                            "Non-short action {} should NOT be masked at min position",
                            idx
                        );
                    }
                }
            }
        }
    }
}
