//! Entropy regularization for PPO to prevent policy collapse
//!
//! This module implements entropy-based reward shaping to maintain action diversity
//! during PPO training. Key features:
//! - Shannon entropy calculation from action probabilities
//! - Bonus/penalty system based on normalized entropy threshold
//! - Numerical stability with epsilon protection
//!
//! # Key Difference from DQN
//! Unlike DQN which computes softmax from Q-values, PPO already has action probabilities
//! from the policy network, making entropy calculation more direct.

use ml_core::MLError;

/// Entropy regularizer for preventing policy collapse in PPO
///
/// Implements Shannon entropy calculation and entropy-based reward shaping
/// to encourage exploration and maintain action diversity.
#[derive(Debug, Clone)]
pub struct EntropyRegularizer {
    /// Maximum possible entropy: `log(num_actions)`
    max_entropy: f64,
    /// Normalized entropy threshold (0.7) for bonus/penalty
    entropy_threshold: f64,
}

impl EntropyRegularizer {
    /// Number of factored actions (5 exposure x 3 order x 3 urgency)
    const NUM_ACTIONS: usize = 45;

    /// Create a new entropy regularizer for 45 factored actions.
    pub fn new() -> Self {
        Self {
            max_entropy: (Self::NUM_ACTIONS as f64).ln(),
            entropy_threshold: 0.7,
        }
    }

    /// Calculate Shannon entropy from action probabilities (host-side).
    ///
    /// # Arguments
    /// * `action_probs` - Action probability slice, shape [`num_actions`] or flattened [`batch_size * num_actions`]
    ///
    /// # Returns
    /// Raw Shannon entropy H(pi) = -sum pi(a|s) * log(pi(a|s))
    pub fn calculate_entropy(&self, action_probs: &[f32]) -> Result<f64, MLError> {
        if action_probs.is_empty() {
            return Err(MLError::InvalidInput(
                "action_probs must be non-empty".to_owned(),
            ));
        }

        let epsilon = 1e-8_f64;
        let entropy: f64 = action_probs
            .iter()
            .map(|&p| {
                let p = p as f64 + epsilon;
                -p * p.ln()
            })
            .sum();

        // Average if this looks like a batch (multiple of NUM_ACTIONS)
        if action_probs.len() > Self::NUM_ACTIONS && action_probs.len() % Self::NUM_ACTIONS == 0 {
            let batch_size = action_probs.len() / Self::NUM_ACTIONS;
            Ok(entropy / batch_size as f64)
        } else {
            Ok(entropy)
        }
    }

    /// Calculate entropy bonus/penalty from action probabilities
    ///
    /// # Returns
    /// - Positive value: Bonus for high entropy (> 0.7 normalized)
    /// - Negative value: Penalty for low entropy (< 0.7 normalized)
    pub fn calculate_entropy_bonus(&self, action_probs: &[f32]) -> Result<f64, MLError> {
        let avg_entropy = self.calculate_entropy(action_probs)?;
        let normalized_entropy = avg_entropy / self.max_entropy;

        if normalized_entropy > self.entropy_threshold {
            Ok(normalized_entropy * 2.0)
        } else {
            Ok(-(self.entropy_threshold - normalized_entropy) * 3.0)
        }
    }
}

impl Default for EntropyRegularizer {
    fn default() -> Self {
        Self::new()
    }
}

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

    /// Create uniform probability vector over 45 actions
    fn uniform_45() -> Vec<f32> {
        vec![1.0 / 45.0; 45]
    }

    /// Create near-deterministic probability vector over 45 actions
    fn deterministic_45(dominant_action: usize) -> Vec<f32> {
        let mut probs = vec![0.001 / 44.0; 45];
        if let Some(p) = probs.get_mut(dominant_action) {
            *p = 0.999;
        }
        probs
    }

    #[test]
    fn test_entropy_uniform_distribution() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        let probs = uniform_45();

        let entropy = regularizer.calculate_entropy(&probs)?;

        let expected = (45.0_f64).ln();
        assert!(
            (entropy - expected).abs() < 0.01,
            "Expected entropy ~{expected:.4}, got {entropy}",
        );
        Ok(())
    }

    #[test]
    fn test_entropy_deterministic() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        let probs = deterministic_45(0);

        let entropy = regularizer.calculate_entropy(&probs)?;

        assert!(entropy < 0.05, "Expected entropy ~0, got {entropy}");
        Ok(())
    }

    #[test]
    fn test_bonus_high_diversity() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        let probs = uniform_45();

        let bonus = regularizer.calculate_entropy_bonus(&probs)?;

        assert!(
            (bonus - 2.0).abs() < 0.05,
            "Expected bonus ~2.0, got {bonus}",
        );
        Ok(())
    }

    #[test]
    fn test_penalty_low_diversity() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        let probs = deterministic_45(38);

        let penalty = regularizer.calculate_entropy_bonus(&probs)?;

        assert!(penalty < 0.0, "Expected penalty < 0.0, got {penalty}");
        assert!(penalty < -1.0, "Expected penalty < -1.0, got {penalty}");
        Ok(())
    }
}