//! Integration tests for softmax-based action selection in DQN
//!
//! Validates that the softmax action selection produces balanced probability
//! distributions instead of always choosing the argmax (greedy) action.

use ml::dqn::dqn::{WorkingDQN, WorkingDQNConfig};
use ml::MLError;
use std::collections::HashMap;

/// Test that softmax produces balanced probabilities for similar Q-values
#[test]
fn test_softmax_produces_balanced_probabilities() -> Result<(), MLError> {
    // Create a DQN agent with temperature = 1.0 (balanced exploration)
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.hidden_dims = vec![64, 32];
    config.num_actions = 3;
    config.temperature_start = 1.0;
    config.temperature_min = 1.0; // Fixed temperature for testing
    config.temperature_decay = 1.0; // No decay
    config.epsilon_start = 0.0; // Disable epsilon-greedy for this test
    config.epsilon_end = 0.0;
    config.warmup_steps = 0; // No warmup

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

    // Create a dummy state vector
    let state = vec![0.0f32; 128];

    // Run 10,000 samples through softmax action selection
    let num_samples = 10_000;
    let mut action_counts: HashMap<u8, usize> = HashMap::new();

    for _ in 0..num_samples {
        let action = agent.select_action(&state)?;
        let action_idx = action.to_int();
        *action_counts.entry(action_idx).or_insert(0) += 1;
    }

    // Calculate empirical probabilities
    let buy_prob = *action_counts.get(&0).unwrap_or(&0) as f64 / num_samples as f64;
    let sell_prob = *action_counts.get(&1).unwrap_or(&0) as f64 / num_samples as f64;
    let hold_prob = *action_counts.get(&2).unwrap_or(&0) as f64 / num_samples as f64;

    println!("Action distribution over {} samples:", num_samples);
    println!(
        "  BUY  (0): {:.2}% ({} samples)",
        buy_prob * 100.0,
        action_counts.get(&0).unwrap_or(&0)
    );
    println!(
        "  SELL (1): {:.2}% ({} samples)",
        sell_prob * 100.0,
        action_counts.get(&1).unwrap_or(&0)
    );
    println!(
        "  HOLD (2): {:.2}% ({} samples)",
        hold_prob * 100.0,
        action_counts.get(&2).unwrap_or(&0)
    );

    // Verify distribution is NOT 100% argmax (action 0)
    // With softmax and temperature=1.0, no single action should dominate >95%
    assert!(
        buy_prob < 0.95,
        "BUY action dominates with {:.2}% probability (expected <95%)",
        buy_prob * 100.0
    );
    assert!(
        sell_prob < 0.95,
        "SELL action dominates with {:.2}% probability (expected <95%)",
        sell_prob * 100.0
    );
    assert!(
        hold_prob < 0.95,
        "HOLD action dominates with {:.2}% probability (expected <95%)",
        hold_prob * 100.0
    );

    // Verify each action gets at least 5% of samples (balanced distribution)
    assert!(
        buy_prob > 0.05,
        "BUY action too rare: {:.2}% (expected >5%)",
        buy_prob * 100.0
    );
    assert!(
        sell_prob > 0.05,
        "SELL action too rare: {:.2}% (expected >5%)",
        sell_prob * 100.0
    );
    assert!(
        hold_prob > 0.05,
        "HOLD action too rare: {:.2}% (expected >5%)",
        hold_prob * 100.0
    );

    // Verify probabilities sum to ~1.0 (within rounding error)
    let total_prob = buy_prob + sell_prob + hold_prob;
    assert!(
        (total_prob - 1.0).abs() < 0.01,
        "Probabilities don't sum to 1.0: {:.4}",
        total_prob
    );

    Ok(())
}

/// Test that low temperature produces more greedy behavior
#[test]
fn test_low_temperature_increases_greediness() -> Result<(), MLError> {
    // Create a DQN agent with low temperature (more greedy)
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.hidden_dims = vec![64, 32];
    config.num_actions = 3;
    config.temperature_start = 0.3; // Low temperature
    config.temperature_min = 0.3;
    config.temperature_decay = 1.0; // No decay
    config.epsilon_start = 0.0; // Disable epsilon-greedy
    config.epsilon_end = 0.0;
    config.warmup_steps = 0;

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

    // Create a dummy state
    let state = vec![0.0f32; 128];

    // Run 1,000 samples
    let num_samples = 1_000;
    let mut action_counts: HashMap<u8, usize> = HashMap::new();

    for _ in 0..num_samples {
        let action = agent.select_action(&state)?;
        *action_counts.entry(action.to_int()).or_insert(0) += 1;
    }

    // Calculate probabilities
    let probs: Vec<f64> = (0..3)
        .map(|i| *action_counts.get(&i).unwrap_or(&0) as f64 / num_samples as f64)
        .collect();

    println!("Low temperature (0.3) action distribution:");
    println!("  BUY  (0): {:.2}%", probs[0] * 100.0);
    println!("  SELL (1): {:.2}%", probs[1] * 100.0);
    println!("  HOLD (2): {:.2}%", probs[2] * 100.0);

    // With low temperature (0.3), the distribution should be LESS uniform than high temp (2.0)
    // but MAY not be fully greedy because Q-values are close at initialization
    // We just verify it's not completely uniform (each action getting ~33%)
    let max_prob = probs.iter().cloned().fold(0.0f64, f64::max);
    let min_prob = probs.iter().cloned().fold(1.0f64, f64::min);
    let prob_range = max_prob - min_prob;

    // Verify there's at least 10% spread in the distribution
    // (not completely uniform like 33.3%, 33.3%, 33.4%)
    assert!(
        prob_range > 0.10,
        "Low temperature should produce non-uniform distribution (range={:.2}%)",
        prob_range * 100.0
    );

    Ok(())
}

/// Test that high temperature produces more exploration
#[test]
fn test_high_temperature_increases_exploration() -> Result<(), MLError> {
    // Create a DQN agent with high temperature (more exploration)
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.hidden_dims = vec![64, 32];
    config.num_actions = 3;
    config.temperature_start = 2.0; // High temperature
    config.temperature_min = 2.0;
    config.temperature_decay = 1.0; // No decay
    config.epsilon_start = 0.0; // Disable epsilon-greedy
    config.epsilon_end = 0.0;
    config.warmup_steps = 0;

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

    // Create a dummy state
    let state = vec![0.0f32; 128];

    // Run 1,000 samples
    let num_samples = 1_000;
    let mut action_counts: HashMap<u8, usize> = HashMap::new();

    for _ in 0..num_samples {
        let action = agent.select_action(&state)?;
        *action_counts.entry(action.to_int()).or_insert(0) += 1;
    }

    // Calculate probabilities
    let probs: Vec<f64> = (0..3)
        .map(|i| *action_counts.get(&i).unwrap_or(&0) as f64 / num_samples as f64)
        .collect();

    println!("High temperature (2.0) action distribution:");
    println!("  BUY  (0): {:.2}%", probs[0] * 100.0);
    println!("  SELL (1): {:.2}%", probs[1] * 100.0);
    println!("  HOLD (2): {:.2}%", probs[2] * 100.0);

    // With high temperature, distribution should be more uniform
    // No single action should dominate >60%
    let max_prob = probs.iter().cloned().fold(0.0f64, f64::max);
    assert!(
        max_prob < 0.60,
        "High temperature should produce more uniform distribution (max_prob={:.2}%)",
        max_prob * 100.0
    );

    // Each action should get at least 15% of samples
    for (i, &prob) in probs.iter().enumerate() {
        assert!(
            prob > 0.15,
            "Action {} too rare with high temperature: {:.2}%",
            i,
            prob * 100.0
        );
    }

    Ok(())
}

/// Test temperature getter/setter methods
#[test]
fn test_temperature_getter_setter() -> Result<(), MLError> {
    let config = WorkingDQNConfig::emergency_safe_defaults();
    let mut agent = WorkingDQN::new(config)?;

    // Check initial temperature
    let initial_temp = agent.get_temperature();
    assert_eq!(initial_temp, 1.0, "Initial temperature should be 1.0");

    // Set new temperature
    agent.set_temperature(0.5);
    assert_eq!(
        agent.get_temperature(),
        0.5,
        "Temperature should be updated to 0.5"
    );

    // Test that very low temperature is clamped to 0.01 (prevent division by zero)
    agent.set_temperature(0.001);
    assert!(
        agent.get_temperature() >= 0.01,
        "Temperature should be clamped to minimum 0.01"
    );

    Ok(())
}

/// Test temperature decay over epochs
#[test]
fn test_temperature_decay() -> Result<(), MLError> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.temperature_start = 1.0;
    config.temperature_min = 0.3;
    config.temperature_decay = 0.99; // 1% decay per epoch

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

    // Initial temperature
    assert_eq!(agent.get_temperature(), 1.0);

    // Update temperature 10 times
    for i in 1..=10 {
        agent.update_temperature();
        let expected = (1.0 * 0.99_f64.powi(i)).max(0.3);
        let actual = agent.get_temperature();
        assert!(
            (actual - expected).abs() < 0.001,
            "After {} updates: expected {:.4}, got {:.4}",
            i,
            expected,
            actual
        );
    }

    // Update many times - should converge to minimum
    for _ in 0..1000 {
        agent.update_temperature();
    }
    assert_eq!(
        agent.get_temperature(),
        0.3,
        "Temperature should converge to minimum"
    );

    Ok(())
}