foxhunt/ml/tests/ppo_entropy_regularization_tests.rs

//! TDD Tests for PPO Entropy Regularization
//!
//! These tests were written BEFORE implementation following TDD methodology.
//! Tests verify:
//! 1. Shannon entropy calculation from action probabilities
//! 2. Entropy bonus for high diversity (> 0.7 normalized entropy)
//! 3. Entropy penalty for low diversity (< 0.7 normalized entropy)
//! 4. Numerical stability with extreme probabilities

use candle_core::{Device, Tensor};
use ml::ppo::entropy_regularization::EntropyRegularizer;
use ml::MLError;

/// Helper function to create probability tensor from raw values
/// Assumes probabilities sum to 1.0 (normalized)
fn create_prob_tensor(probs: &[f32]) -> Result<Tensor, MLError> {
    let tensor = Tensor::new(probs, &Device::Cpu)?;
    Ok(tensor.reshape(&[1, probs.len()])?)
}

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

    // Test 1: Uniform distribution - maximum entropy
    // For 3 actions: H = -3 * (1/3 * log(1/3)) = log(3) ≈ 1.0986
    // Normalized: 1.0986 / 1.0986 = 1.0
    let uniform_probs = create_prob_tensor(&[1.0/3.0, 1.0/3.0, 1.0/3.0])?;
    let entropy = regularizer.calculate_entropy(&uniform_probs)?;

    // Expected: log(3) ≈ 1.0986
    assert!(
        (entropy - 1.0986).abs() < 0.01,
        "Expected uniform entropy ~1.0986, got {}",
        entropy
    );

    // Test 2: Deterministic distribution - minimum entropy
    // For [1.0, 0.0, 0.0]: H = -(1 * log(1) + 0 * log(0) + 0 * log(0)) = 0
    let deterministic_probs = create_prob_tensor(&[1.0, 0.0, 0.0])?;
    let entropy = regularizer.calculate_entropy(&deterministic_probs)?;

    // Expected: 0 (with small epsilon tolerance for numerical stability)
    assert!(
        entropy < 0.01,
        "Expected deterministic entropy ~0, got {}",
        entropy
    );

    // Test 3: Moderate distribution
    // For [0.5, 0.3, 0.2]: H = -(0.5*log(0.5) + 0.3*log(0.3) + 0.2*log(0.2))
    // Expected: ~1.03
    let moderate_probs = create_prob_tensor(&[0.5, 0.3, 0.2])?;
    let entropy = regularizer.calculate_entropy(&moderate_probs)?;

    // Expected: ~1.03
    assert!(
        (entropy - 1.03).abs() < 0.05,
        "Expected moderate entropy ~1.03, got {}",
        entropy
    );

    Ok(())
}

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

    // Create probabilities with high diversity (normalized entropy > 0.7)
    // Uniform distribution: normalized entropy = 1.0
    let high_diversity_probs = create_prob_tensor(&[1.0/3.0, 1.0/3.0, 1.0/3.0])?;
    let bonus = regularizer.calculate_entropy_bonus(&high_diversity_probs)?;

    // Expected: normalized_entropy * 2.0 = 1.0 * 2.0 = 2.0
    assert!(
        bonus > 0.0,
        "High diversity should give positive bonus, got {}",
        bonus
    );
    assert!(
        (bonus - 2.0).abs() < 0.1,
        "Expected bonus ~2.0 (1.0 * 2.0), got {}",
        bonus
    );

    // Test with slightly less uniform (but still > 0.7)
    let moderate_high_probs = create_prob_tensor(&[0.4, 0.35, 0.25])?;
    let bonus2 = regularizer.calculate_entropy_bonus(&moderate_high_probs)?;

    // Should still be positive (normalized entropy ~0.95)
    assert!(
        bonus2 > 0.0,
        "Moderate high diversity should give positive bonus, got {}",
        bonus2
    );
    assert!(
        bonus2 > 1.4, // At least 0.7 * 2.0
        "Expected bonus > 1.4, got {}",
        bonus2
    );

    Ok(())
}

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

    // Create probabilities with low diversity (normalized entropy < 0.7)
    // Near-deterministic: [0.9, 0.05, 0.05] → low entropy
    let low_diversity_probs = create_prob_tensor(&[0.9, 0.05, 0.05])?;
    let penalty = regularizer.calculate_entropy_bonus(&low_diversity_probs)?;

    // Expected: -(0.7 - normalized_entropy) * 3.0
    // For [0.9, 0.05, 0.05]: H ≈ 0.57, normalized ≈ 0.52
    // Penalty: -(0.7 - 0.52) * 3.0 ≈ -0.54
    assert!(
        penalty < 0.0,
        "Low diversity should give negative penalty, got {}",
        penalty
    );
    assert!(
        penalty < -0.4,
        "Expected penalty < -0.4, got {}",
        penalty
    );

    // Test with very low diversity (almost deterministic)
    let very_low_probs = create_prob_tensor(&[0.98, 0.01, 0.01])?;
    let penalty2 = regularizer.calculate_entropy_bonus(&very_low_probs)?;

    // Should be more negative than moderate low diversity
    assert!(
        penalty2 < penalty,
        "Very low diversity penalty ({}) should be more negative than moderate ({})",
        penalty2,
        penalty
    );
    assert!(
        penalty2 < -1.5,
        "Expected very low diversity penalty < -1.5, got {}",
        penalty2
    );

    Ok(())
}

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

    // Test with extreme probabilities that could cause numerical issues

    // Test 1: Very small probabilities (close to 0)
    let small_probs = create_prob_tensor(&[0.9999, 0.0001, 0.0])?;
    let result1 = regularizer.calculate_entropy_bonus(&small_probs);
    assert!(result1.is_ok(), "Should handle small probabilities without error");
    let bonus1 = result1?;
    assert!(bonus1.is_finite(), "Result should be finite, got {}", bonus1);

    // Test 2: Very close to uniform (but with floating point rounding)
    let near_uniform = create_prob_tensor(&[0.3333, 0.3334, 0.3333])?;
    let result2 = regularizer.calculate_entropy_bonus(&near_uniform);
    assert!(result2.is_ok(), "Should handle near-uniform probabilities");
    let bonus2 = result2?;
    assert!(bonus2.is_finite(), "Result should be finite, got {}", bonus2);

    // Test 3: Check that epsilon protection prevents log(0) = -∞
    // This is handled internally by adding 1e-8 epsilon
    let zero_prob = create_prob_tensor(&[1.0, 0.0, 0.0])?;
    let result3 = regularizer.calculate_entropy(&zero_prob);
    assert!(result3.is_ok(), "Should handle zero probabilities with epsilon");
    let entropy = result3?;
    assert!(entropy.is_finite(), "Entropy should be finite with zero probs, got {}", entropy);
    assert!(entropy >= 0.0, "Entropy should be non-negative, got {}", entropy);

    // Test 4: Batch processing with mixed distributions
    // Shape: [4, 3] - batch of 4 probability distributions
    let batch_probs = Tensor::new(
        &[
            0.33, 0.33, 0.34,  // Uniform
            0.9, 0.05, 0.05,   // Low diversity
            0.6, 0.3, 0.1,     // Medium diversity
            0.98, 0.01, 0.01,  // Very low diversity
        ],
        &Device::Cpu
    )?.reshape(&[4, 3])?;

    let result4 = regularizer.calculate_entropy_bonus(&batch_probs);
    assert!(result4.is_ok(), "Should handle batch processing");
    let bonus4 = result4?;
    assert!(bonus4.is_finite(), "Batch result should be finite, got {}", bonus4);

    Ok(())
}