foxhunt/crates/ml/tests/ppo_entropy_regularization_tests.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! 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::new_cuda(0).expect("CUDA required"))?;
    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::new_cuda(0).expect("CUDA required")
    )?.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(())
}