//! Entropy regularization for DQN to prevent action collapse
//!
//! This module implements entropy-based reward shaping and stochastic action selection
//! to maintain policy diversity during training. Key features:
//! - Shannon entropy calculation with normalization
//! - Bonus/penalty system based on entropy threshold
//! - Temperature-controlled softmax action selection
//!
//! # Example
//! ```rust,no_run
//! use ml_dqn::entropy_regularization::EntropyRegularizer;
//!
//! let regularizer = EntropyRegularizer::new();
//! let q_values = vec![3.0_f32, 2.0, 1.5, 1.0, 0.5];
//! let bonus = regularizer.calculate_entropy_bonus(&q_values).unwrap();
//! let action = regularizer.softmax_action_selection(&q_values, 1.0).unwrap();
//! ```

use ml_core::MLError;

/// Entropy regularizer for preventing policy collapse
///
/// 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 for 5 actions: log(5) ~= 1.609
    max_entropy: f64,
    /// Normalized entropy threshold (0.7) for bonus/penalty
    entropy_threshold: f64,
}

impl EntropyRegularizer {
    /// Create a new entropy regularizer
    ///
    /// # Configuration
    /// - `max_entropy`: log(5) ~= 1.6094 for 5 exposure actions (Short100, Short50, Flat, Long50, Long100)
    /// - `entropy_threshold`: 0.7 normalized entropy
    ///   - Above 0.7: bonus = `normalized_entropy` (scaled by `entropy_coefficient` in DQN loss)
    ///   - Below 0.7: penalty = -(threshold - `normalized_entropy`) (scaled by `entropy_coefficient` in DQN loss)
    pub fn new() -> Self {
        Self {
            max_entropy: (5.0_f64).ln(), // log(5) for 5-action exposure space
            entropy_threshold: 0.7,
        }
    }

    /// Calculate entropy bonus/penalty from Q-values (CPU).
    ///
    /// Accepts flat Q-values `[num_actions]` or batched `[batch_size * num_actions]`
    /// (with `num_actions_hint` to split). For GPU-resident computation, use the
    /// fused CUDA kernel in the DQN trainer.
    ///
    /// # Arguments
    /// * `q_values` - Q-value slice, shape [`num_actions`] or [`batch_size * num_actions`]
    ///
    /// # 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, q_values: &[f32]) -> Result<f64, MLError> {
        if q_values.is_empty() {
            return Err(MLError::InvalidInput("empty q_values".into()));
        }

        // Determine batch layout
        // If length is divisible by 5 and > 5, treat as batched [batch, 5]
        let num_actions = if q_values.len() >= 5 && q_values.len() % 5 == 0 {
            5
        } else {
            q_values.len()
        };
        let batch_size = q_values.len() / num_actions;

        let mut total_entropy = 0.0_f64;

        for b in 0..batch_size {
            let base = b * num_actions;
            let slice = q_values.get(base..base + num_actions).ok_or_else(|| {
                MLError::ModelError("q_values slice out of bounds".into())
            })?;

            // Step 1: Softmax with LogSumExp trick
            let max_q = slice.iter().copied().fold(f32::NEG_INFINITY, f32::max);
            let mut exp_sum = 0.0_f32;
            let mut exps = Vec::with_capacity(num_actions);
            for &q in slice {
                let e = (q - max_q).exp();
                exps.push(e);
                exp_sum += e;
            }

            // Step 2: Shannon entropy H(pi) = -sum pi(a|s) * log(pi(a|s))
            let mut entropy = 0.0_f64;
            for &e in &exps {
                let p = (e / exp_sum) as f64;
                if p > 1e-10 {
                    entropy -= p * p.ln();
                }
            }
            total_entropy += entropy;
        }

        let avg_entropy = total_entropy / batch_size as f64;

        // Step 4: Normalize to [0, 1]
        let normalized_entropy = avg_entropy / self.max_entropy;

        // Step 5: Apply bonus/penalty based on threshold
        if normalized_entropy > self.entropy_threshold {
            Ok(normalized_entropy)
        } else {
            Ok(-(self.entropy_threshold - normalized_entropy))
        }
    }

    /// Select action stochastically using temperature-controlled softmax (CPU).
    ///
    /// # Arguments
    /// * `q_values` - Q-value slice, shape [`num_actions`]
    /// * `temperature` - Temperature parameter controlling randomness
    ///   - Low (0.1): Near-deterministic (always picks highest Q-value)
    ///   - Medium (1.0): Balanced stochastic sampling [DEFAULT]
    ///   - High (10.0): Near-uniform random exploration
    ///
    /// # Returns
    /// Selected action index (0 = Short100, 1 = Short50, 2 = Flat, 3 = Long50, 4 = Long100)
    ///
    /// # Errors
    /// Returns error if temperature is zero or negative
    pub fn softmax_action_selection(&self, q_values: &[f32], temperature: f64) -> Result<i64, MLError> {
        if temperature <= 0.0 {
            return Err(MLError::InvalidInput(format!(
                "Temperature must be positive, got {}",
                temperature
            )));
        }
        if q_values.is_empty() {
            return Err(MLError::InvalidInput("empty q_values".into()));
        }

        let temp = temperature as f32;

        // Step 1: Temperature scaling
        let scaled: Vec<f32> = q_values.iter().map(|&q| q / temp).collect();

        // Step 2: Softmax with numerical stability
        let max_q = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        let mut exp_sum = 0.0_f32;
        let mut probs = Vec::with_capacity(q_values.len());
        for &s in &scaled {
            let e = (s - max_q).exp();
            probs.push(e);
            exp_sum += e;
        }
        for p in &mut probs {
            *p /= exp_sum;
        }

        // Step 3: Gumbel-max sampling (CPU-based)
        // Use simple cumulative probability sampling with system RNG
        use rand::Rng;
        let mut rng = rand::thread_rng();
        let u: f64 = rng.gen();
        let mut cumulative = 0.0_f64;
        for (i, &p) in probs.iter().enumerate() {
            cumulative += p as f64;
            if u <= cumulative {
                return Ok(i as i64);
            }
        }

        // Fallback: return last action
        Ok((q_values.len() - 1) as i64)
    }
}

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

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

    #[test]
    fn test_entropy_uniform_distribution() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        // 5 actions matching the real DQN exposure space
        let q_values = vec![1.0_f32, 1.0, 1.0, 1.0, 1.0]; // Uniform after softmax

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

        // Uniform over 5 actions: entropy = log(5), normalized = 1.0
        assert!(
            (bonus - 1.0).abs() < 0.01,
            "Expected bonus ~1.0, got {}",
            bonus
        );
        Ok(())
    }

    #[test]
    fn test_entropy_deterministic_policy() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        // 5 actions: one dominant, rest near zero
        let q_values = vec![1000.0_f32, 0.0, 0.0, 0.0, 0.0]; // Softmax -> [1, 0, 0, 0, 0]

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

        // Deterministic policy: entropy ~= 0, normalized ~= 0
        // penalty = -(0.7 - 0) = -0.7
        assert!(bonus < -0.6, "Expected penalty < -0.6, got {}", bonus);
        assert!(bonus > -0.8, "Expected penalty > -0.8, got {}", bonus);
        Ok(())
    }

    #[test]
    fn test_entropy_high_diversity() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        // 5 actions with high diversity (close Q-values -> near-uniform softmax)
        let q_values = vec![2.0_f32, 1.9, 1.8, 1.7, 1.6];

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

        // High diversity: normalized entropy > 0.7 -> bonus = normalized_entropy
        assert!(
            bonus > 0.9,
            "Expected bonus > 0.9 for near-uniform 5-action, got {}",
            bonus
        );
        Ok(())
    }

    #[test]
    fn test_entropy_low_diversity() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        // 5 actions with low diversity (one dominant Q-value)
        let q_values = vec![5.0_f32, 0.1, 0.2, 0.1, 0.1];

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

        // Low diversity: normalized entropy < 0.7
        // penalty = -(0.7 - normalized_entropy), should be negative
        assert!(bonus < 0.0, "Expected penalty < 0.0, got {}", bonus);
        Ok(())
    }

    #[test]
    fn test_softmax_action_selection() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        // 5 actions matching the real DQN exposure space
        let q_values = vec![3.0_f32, 2.0, 1.5, 1.0, 0.5];

        // Run 1000 samples to check probabilistic distribution
        let mut action_counts = [0; 5];
        for _ in 0..1000 {
            let action = regularizer.softmax_action_selection(&q_values, 1.0)?;
            assert!(action >= 0 && action < 5, "Invalid action: {}", action);
            action_counts[action as usize] += 1;
        }

        // With Q-values [3.0, 2.0, 1.5, 1.0, 0.5] and temp=1.0:
        // Action 0 should be selected most frequently
        assert!(
            action_counts[0] > 300,
            "Action 0 should be > 30%, got {}%",
            action_counts[0] / 10
        );

        // All actions should get some selection probability
        assert!(
            action_counts[1] > 50,
            "Action 1 should be > 5%, got {}%",
            action_counts[1] / 10
        );

        Ok(())
    }

    #[test]
    fn test_temperature_effect() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        // 5 actions matching the real DQN exposure space
        let q_values = vec![3.0_f32, 2.0, 1.5, 1.0, 0.5];

        // Low temperature (0.1): More deterministic
        let mut low_temp_counts = [0; 5];
        for _ in 0..1000 {
            let action = regularizer.softmax_action_selection(&q_values, 0.1)?;
            low_temp_counts[action as usize] += 1;
        }

        // High temperature (10.0): More uniform
        let mut high_temp_counts = [0; 5];
        for _ in 0..1000 {
            let action = regularizer.softmax_action_selection(&q_values, 10.0)?;
            high_temp_counts[action as usize] += 1;
        }

        // Low temp: Action 0 should dominate (> 90%)
        assert!(
            low_temp_counts[0] > 900,
            "Low temp: Action 0 should be > 90%, got {}%",
            low_temp_counts[0] / 10
        );

        // High temp: Actions should be more evenly distributed
        assert!(
            high_temp_counts[0] > 100 && high_temp_counts[4] > 100,
            "High temp should be more uniform: {:?}",
            high_temp_counts
        );

        Ok(())
    }

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

        // Test various Q-value distributions with 5 actions
        let test_cases = vec![
            vec![1.0_f32, 1.0, 1.0, 1.0, 1.0],        // Uniform
            vec![10.0, 0.0, 0.0, 0.0, 0.0],             // Deterministic
            vec![2.0, 1.8, 1.5, 1.2, 1.0],              // Moderate diversity
            vec![3.0, 2.5, 2.0, 1.5, 1.0],              // Higher diversity
        ];

        for q_vals in test_cases {
            // Compute softmax probabilities
            let max_q = q_vals.iter().copied().fold(f32::NEG_INFINITY, f32::max);
            let mut exp_sum = 0.0_f32;
            let mut exps = Vec::new();
            for &q in &q_vals {
                let e = (q - max_q).exp();
                exps.push(e);
                exp_sum += e;
            }

            let mut entropy = 0.0_f64;
            for &e in &exps {
                let p = (e / exp_sum) as f64;
                if p > 1e-10 {
                    entropy -= p * p.ln();
                }
            }
            let normalized_entropy = entropy / regularizer.max_entropy;

            // Verify normalization is in [0, 1]
            assert!(
                normalized_entropy >= 0.0 && normalized_entropy <= 1.0 + 1e-6,
                "Normalized entropy out of bounds: {} for Q-values {:?}",
                normalized_entropy,
                q_vals
            );
        }

        Ok(())
    }

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

        // Create batch of Q-values: [32 * 5] (5-action space)
        use rand::Rng;
        let mut rng = rand::thread_rng();
        let q_values: Vec<f32> = (0..32 * 5).map(|_| rng.gen::<f32>() * 2.0 - 1.0).collect();

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

        // Should return a single scalar value (not panic or return NaN)
        assert!(bonus.is_finite(), "Batch entropy should be finite, got {}", bonus);
        Ok(())
    }
}
