//! Softmax exploration utilities for DQN action selection
//!
//! This module provides temperature-based softmax sampling to replace
//! argmax tie-breaking bias and fix action diversity collapse.
//!
//! Key features:
//! - Log-sum-exp trick for numerical stability
//! - Temperature scaling for exploration control
//! - Entropy calculation for monitoring
//! - Batched and single-state support

use ml_core::MLError;
use rand::{thread_rng, Rng};

/// Compute softmax probabilities with temperature scaling (host-side, cold path).
///
/// Converts Q-values to probability distribution where higher values
/// get higher (but not exclusive) probability. Temperature controls
/// the exploration-exploitation tradeoff.
///
/// # Arguments
///
/// * `q_values` - Slice of Q-values
/// * `temperature` - Temperature parameter (0.1 = greedy, 10.0 = uniform)
///
/// # Returns
///
/// Probability distribution where values sum to 1.0
///
/// # Numerical Stability
///
/// Uses log-sum-exp trick to prevent overflow/underflow:
/// ```text
/// softmax(x) = exp(x - max(x)) / sum(exp(x - max(x)))
/// ```
pub fn softmax_with_temperature(q_values: &[f32], temperature: f64) -> Result<Vec<f32>, MLError> {
    if q_values.is_empty() {
        return Err(MLError::InvalidInput("Empty Q-values".to_owned()));
    }

    // Clamp temperature to prevent division by zero
    let temp = temperature.max(1e-6) as f32;

    // Scale Q-values by temperature
    let scaled: Vec<f32> = q_values.iter().map(|&q| q / temp).collect();

    // Log-sum-exp trick: find max for numerical stability
    let max_val = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);

    // Compute exp(scaled - max)
    let exp_vals: Vec<f32> = scaled.iter().map(|&s| (s - max_val).exp()).collect();

    // Sum of exp values
    let sum_exp: f32 = exp_vals.iter().sum();

    if sum_exp <= 0.0 || !sum_exp.is_finite() {
        return Err(MLError::ModelError(format!(
            "Softmax sum is invalid: {sum_exp}"
        )));
    }

    // Normalize to get probabilities
    let probs: Vec<f32> = exp_vals.iter().map(|&e| e / sum_exp).collect();

    Ok(probs)
}

/// Compute softmax probabilities for a batch of Q-value rows.
///
/// Each row is independently normalized.
///
/// # Arguments
///
/// * `q_values_batch` - Batch of Q-value slices (each row = one state)
/// * `temperature` - Temperature parameter
///
/// # Returns
///
/// Batch of probability distributions
pub fn softmax_with_temperature_batch(
    q_values_batch: &[Vec<f32>],
    temperature: f64,
) -> Result<Vec<Vec<f32>>, MLError> {
    q_values_batch
        .iter()
        .map(|row| softmax_with_temperature(row, temperature))
        .collect()
}

/// Sample action from softmax distribution
///
/// Stochastically selects an action according to the softmax probability
/// distribution over Q-values. Higher Q-values have higher probability
/// but all actions have non-zero chance (unlike argmax).
///
/// # Arguments
///
/// * `q_values` - Slice of Q-values
/// * `temperature` - Temperature parameter (0.1 = greedy, 10.0 = uniform)
///
/// # Returns
///
/// Action index (0 to num_actions-1)
pub fn sample_from_softmax(q_values: &[f32], temperature: f64) -> Result<u32, MLError> {
    let probs = softmax_with_temperature(q_values, temperature)?;

    // Sample from categorical distribution
    let mut rng = thread_rng();
    let rand_val: f32 = rng.gen(); // Uniform [0, 1)

    // Cumulative sum to find action
    let mut cumsum = 0.0;
    for (i, &prob) in probs.iter().enumerate() {
        cumsum += prob;
        if rand_val < cumsum {
            return Ok(i as u32);
        }
    }

    // Fallback: return last action (handles floating-point edge cases)
    Ok((probs.len().saturating_sub(1)) as u32)
}

/// Compute entropy of softmax distribution
///
/// Calculates Shannon entropy to measure exploration level:
/// - High entropy (~log2(N)): Uniform exploration
/// - Low entropy (~0): Greedy exploitation
///
/// # Arguments
///
/// * `q_values` - Slice of Q-values
/// * `temperature` - Temperature parameter
///
/// # Returns
///
/// Entropy in bits (base-2 logarithm)
pub fn softmax_entropy(q_values: &[f32], temperature: f64) -> Result<f64, MLError> {
    let probs = softmax_with_temperature(q_values, temperature)?;

    // Calculate Shannon entropy: H = -sum(p_i * log2(p_i))
    let mut entropy = 0.0_f64;
    for &prob in &probs {
        if prob > 1e-10 {
            // Skip near-zero probabilities to avoid log(0)
            entropy -= (prob as f64) * (prob as f64).log2();
        }
    }

    Ok(entropy)
}

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

    #[test]
    fn test_softmax_basic() -> Result<(), MLError> {
        let q_values = [1.0_f32, 2.0, 3.0];

        let probs = softmax_with_temperature(&q_values, 1.0)?;

        // Probabilities should sum to 1.0
        let sum: f32 = probs.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5, "Sum should be 1.0, got {}", sum);

        // Highest Q-value should have highest probability
        assert!(probs[2] > probs[1] && probs[1] > probs[0]);

        Ok(())
    }

    #[test]
    fn test_sampling_basic() -> Result<(), MLError> {
        let q_values = [1.0_f32, 2.0, 3.0];

        // Sample 100 times to verify it works
        for _ in 0..100 {
            let action = sample_from_softmax(&q_values, 1.0)?;
            assert!(action < 3, "Action should be in range [0, 3), got {}", action);
        }

        Ok(())
    }

    #[test]
    fn test_entropy_basic() -> Result<(), MLError> {
        // Uniform distribution (max entropy)
        let q_uniform = [0.0_f32; 3];
        let entropy_uniform = softmax_entropy(&q_uniform, 1.0)?;

        // Max entropy for 3 actions: log2(3) ~ 1.585
        assert!(
            (entropy_uniform - 1.585).abs() < 0.01,
            "Uniform entropy should be ~1.585, got {}",
            entropy_uniform
        );

        // Deterministic distribution (low entropy)
        let q_det = [-1000.0_f32, 0.0, 1000.0];
        let entropy_det = softmax_entropy(&q_det, 0.1)?;

        assert!(
            entropy_det < 0.1,
            "Deterministic entropy should be low, got {}",
            entropy_det
        );

        Ok(())
    }
}
