//! Integration tests for entropy regularization in DQN training
//!
//! These tests verify that entropy regularization is properly integrated
//! into the DQN loss calculation and maintains action diversity.

use candle_core::{DType, Device, Tensor};
use ml::dqn::{EntropyRegularizer, Experience, FactoredAction, WorkingDQN, WorkingDQNConfig};
use ml::MLError;

/// Test that entropy regularizer calculates correct penalties for 45-action space
#[test]
fn test_entropy_regularizer_45_actions() -> Result<(), MLError> {
    let regularizer = EntropyRegularizer::new();

    // Create Q-values for 45 actions (batch_size=1)
    let device = Device::cuda_if_available(0)?;

    // Case 1: Uniform distribution (high entropy, should get bonus)
    let uniform_q = vec![1.0f32; 45];
    let q_tensor = Tensor::from_vec(uniform_q, (1, 45), &device)?;
    let bonus = regularizer.calculate_entropy_bonus(&q_tensor)?;

    // Uniform distribution: entropy = log(45) ≈ 3.807, normalized = 1.0
    // Expected: 1.0 > 0.7, so bonus = 1.0 * 2.0 = 2.0
    assert!(
        bonus > 1.8,
        "Uniform distribution should get bonus > 1.8, got {}",
        bonus
    );

    // Case 2: Deterministic policy (low entropy, should get penalty)
    let mut deterministic_q = vec![0.0f32; 45];
    deterministic_q[0] = 10.0; // Only action 0 has high Q-value
    let q_tensor = Tensor::from_vec(deterministic_q, (1, 45), &device)?;
    let penalty = regularizer.calculate_entropy_bonus(&q_tensor)?;

    // Deterministic: entropy ≈ 0, normalized ≈ 0
    // Expected: 0 < 0.7, so penalty = -(0.7 - 0) * 3.0 = -2.1
    assert!(
        penalty < -2.0,
        "Deterministic policy should get penalty < -2.0, got {}",
        penalty
    );

    // Case 3: High diversity (all 45 actions roughly equal)
    let high_diversity_q: Vec<f32> = (0..45).map(|i| 5.0 - (i as f32 * 0.05)).collect();
    let q_tensor = Tensor::from_vec(high_diversity_q, (1, 45), &device)?;
    let result = regularizer.calculate_entropy_bonus(&q_tensor)?;

    // High diversity should get bonus (normalized entropy > 0.7)
    assert!(
        result > 0.0,
        "High diversity should get bonus > 0, got {}",
        result
    );

    Ok(())
}

/// Test that DQN calculate_entropy_penalty uses 45-action space
#[test]
fn test_dqn_entropy_penalty_45_actions() -> Result<(), MLError> {
    let config = WorkingDQNConfig {
        state_dim: 128,
        num_actions: 45,
        hidden_dims: vec![256, 128],
        learning_rate: 1e-4,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        replay_buffer_capacity: 10000,
        batch_size: 32,
        min_replay_size: 100,
        target_update_freq: 100,
        use_double_dqn: true,
        use_huber_loss: true,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
        tau: 0.001,
        use_soft_updates: false,
        warmup_steps: 0,
        entropy_coefficient: 0.01,
    };

    let mut dqn = WorkingDQN::new(config)?;

    // Simulate taking diverse actions (all 45 actions once)
    let dummy_state = vec![0.0f32; 128];
    for action_idx in 0..45 {
        let action = FactoredAction::from_index(action_idx)?;
        let reward = 0.1;
        let next_state = vec![0.0f32; 128];
        let done = false;

        // Track action for entropy calculation
        dqn.track_action(action);

        // Store experience in replay buffer
        let exp = Experience::new(
            dummy_state.clone(),
            action.to_index() as u8,
            reward,
            next_state,
            done,
        );
        dqn.store_experience(exp)?;
    }

    // Calculate entropy - should be high (all 45 actions used equally)
    let entropy = dqn.get_action_entropy();

    // Max entropy for 45 actions: log2(45) ≈ 5.492
    // With uniform distribution: entropy = 5.492
    assert!(
        entropy > 5.4,
        "Uniform 45-action distribution should have entropy > 5.4, got {}",
        entropy
    );

    Ok(())
}

/// Test that get_action_entropy correctly counts 45 actions
#[test]
fn test_get_action_entropy_45_actions() -> Result<(), MLError> {
    let config = WorkingDQNConfig {
        state_dim: 128,
        num_actions: 45,
        hidden_dims: vec![256, 128],
        learning_rate: 1e-4,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        replay_buffer_capacity: 10000,
        batch_size: 32,
        min_replay_size: 100,
        target_update_freq: 100,
        use_double_dqn: true,
        use_huber_loss: true,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
        tau: 0.001,
        use_soft_updates: false,
        warmup_steps: 0,
        entropy_coefficient: 0.01,
    };

    let mut dqn = WorkingDQN::new(config)?;

    // Test 1: All same action (deterministic)
    let dummy_state = vec![0.0f32; 128];
    for _ in 0..100 {
        let action = FactoredAction::from_index(0)?;
        dqn.track_action(action);
    }

    let entropy_deterministic = dqn.get_action_entropy();
    assert!(
        entropy_deterministic < 0.1,
        "Deterministic policy should have entropy ≈ 0, got {}",
        entropy_deterministic
    );

    // Test 2: Uniform distribution (all 45 actions)
    let mut dqn = WorkingDQN::new(WorkingDQNConfig {
        state_dim: 128,
        num_actions: 45,
        hidden_dims: vec![256, 128],
        learning_rate: 1e-4,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        replay_buffer_capacity: 10000,
        batch_size: 32,
        min_replay_size: 100,
        target_update_freq: 100,
        use_double_dqn: true,
        use_huber_loss: true,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
        tau: 0.001,
        use_soft_updates: false,
        warmup_steps: 0,
        entropy_coefficient: 0.01,
    })?;

    for action_idx in 0..45 {
        let action = FactoredAction::from_index(action_idx)?;
        dqn.track_action(action);
    }

    let entropy_uniform = dqn.get_action_entropy();

    // Max entropy for 45 actions: log2(45) ≈ 5.492
    assert!(
        entropy_uniform > 5.4,
        "Uniform distribution should have entropy ≈ 5.492, got {}",
        entropy_uniform
    );

    Ok(())
}