//! Unit tests for softmax sampling boundary bias fix (Wave 2 Agent 2E)
//!
//! Tests verify:
//! 1. Uniform distribution sampling (equal probabilities)
//! 2. Boundary cases (prob=0.0, prob=1.0)
//! 3. Statistical distribution over 10,000 samples
//! 4. No action index bias (chi-square test)

use ml::dqn::{WorkingDQN, WorkingDQNConfig};

/// Test that uniform probabilities produce uniform action distribution
#[test]
fn test_uniform_probability_sampling() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.epsilon_start = 0.0; // Disable epsilon-greedy for pure softmax testing
    config.epsilon_end = 0.0;
    config.temperature_start = 1.0; // Balanced temperature

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

    // Create state that produces uniform Q-values (should lead to uniform probabilities)
    let state = vec![0.0; 52];

    // Sample 10,000 actions
    let mut action_counts = [0, 0, 0];
    for _ in 0..10000 {
        let action = dqn.select_action(&state)?;
        action_counts[action as usize] += 1;
    }

    // Expected: ~3333 per action (uniform distribution)
    let expected = 10000.0 / 3.0;

    // Verify each action is within 5% of expected (chi-square tolerance)
    for (i, &count) in action_counts.iter().enumerate() {
        let ratio = count as f64 / expected;
        assert!(
            ratio >= 0.90 && ratio <= 1.10,
            "Action {} count ({}) deviates >10% from expected ({:.0}), ratio={:.3}",
            i,
            count,
            expected,
            ratio
        );
    }

    // Calculate chi-square statistic
    let chi_square: f64 = action_counts
        .iter()
        .map(|&count| {
            let diff = count as f64 - expected;
            (diff * diff) / expected
        })
        .sum();

    // Chi-square critical value at 95% confidence, 2 degrees of freedom: 5.991
    assert!(
        chi_square < 5.991,
        "Chi-square test failed: {:.3} > 5.991 (not uniform distribution)",
        chi_square
    );

    println!(
        "✓ Uniform sampling test passed: BUY={}, SELL={}, HOLD={}, χ²={:.3}",
        action_counts[0], action_counts[1], action_counts[2], chi_square
    );

    Ok(())
}

/// Test boundary case: probability = 0.0 (action should never be selected)
#[test]
fn test_zero_probability_boundary() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.epsilon_start = 0.0; // Disable epsilon-greedy
    config.epsilon_end = 0.0;
    config.temperature_start = 0.1; // Low temperature for peaked distribution

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

    // Create state that strongly favors HOLD (index 2)
    // With low temperature, this should make BUY/SELL probabilities very small
    let mut state = vec![0.0; 52];
    state[0] = -10.0; // Strong negative signal for BUY
    state[1] = -10.0; // Strong negative signal for SELL

    // Sample 1,000 actions
    let mut action_counts = [0, 0, 0];
    for _ in 0..1000 {
        let action = dqn.select_action(&state)?;
        action_counts[action as usize] += 1;
    }

    // HOLD should dominate (>95% of samples)
    let hold_ratio = action_counts[2] as f64 / 1000.0;
    assert!(
        hold_ratio > 0.95,
        "Expected HOLD to dominate with low temperature, got ratio={:.3}",
        hold_ratio
    );

    println!(
        "✓ Zero probability boundary test passed: BUY={}, SELL={}, HOLD={} ({:.1}%)",
        action_counts[0],
        action_counts[1],
        action_counts[2],
        hold_ratio * 100.0
    );

    Ok(())
}

/// Test boundary case: probability = 1.0 (action should always be selected)
#[test]
fn test_one_probability_boundary() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.epsilon_start = 0.0; // Disable epsilon-greedy
    config.epsilon_end = 0.0;
    config.temperature_start = 0.01; // Very low temperature for extremely peaked distribution

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

    // Create state that strongly favors BUY (index 0)
    let mut state = vec![0.0; 52];
    state[0] = 100.0; // Extremely strong signal for BUY
    state[1] = -100.0; // Strong negative signal for SELL
    state[2] = -100.0; // Strong negative signal for HOLD

    // Sample 1,000 actions
    let mut action_counts = [0, 0, 0];
    for _ in 0..1000 {
        let action = dqn.select_action(&state)?;
        action_counts[action as usize] += 1;
    }

    // BUY should dominate (>99% of samples with very low temperature)
    let buy_ratio = action_counts[0] as f64 / 1000.0;
    assert!(
        buy_ratio > 0.99,
        "Expected BUY to dominate with very low temperature, got ratio={:.3}",
        buy_ratio
    );

    println!(
        "✓ One probability boundary test passed: BUY={} ({:.1}%), SELL={}, HOLD={}",
        action_counts[0],
        buy_ratio * 100.0,
        action_counts[1],
        action_counts[2]
    );

    Ok(())
}

/// Test that sampling doesn't favor lower-index actions (no bias)
#[test]
fn test_no_action_index_bias() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.epsilon_start = 0.0; // Disable epsilon-greedy
    config.epsilon_end = 0.0;
    config.temperature_start = 1.0; // Balanced temperature

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

    // Test multiple states to ensure no systematic bias
    let test_cases = vec![
        ("uniform", vec![0.0; 52]),
        ("slightly positive", vec![0.5; 52]),
        ("slightly negative", vec![-0.5; 52]),
    ];

    for (name, state) in test_cases {
        let mut action_counts = [0, 0, 0];

        // Sample 10,000 actions
        for _ in 0..10000 {
            let action = dqn.select_action(&state)?;
            action_counts[action as usize] += 1;
        }

        // Check for lower-index bias (BUY should not be significantly higher than others)
        let expected = 10000.0 / 3.0;
        let buy_ratio = action_counts[0] as f64 / expected;
        let sell_ratio = action_counts[1] as f64 / expected;
        let hold_ratio = action_counts[2] as f64 / expected;

        // None should deviate more than 10% from expected
        assert!(
            buy_ratio >= 0.90 && buy_ratio <= 1.10,
            "Test '{}': BUY ratio {:.3} deviates >10% from expected",
            name,
            buy_ratio
        );
        assert!(
            sell_ratio >= 0.90 && sell_ratio <= 1.10,
            "Test '{}': SELL ratio {:.3} deviates >10% from expected",
            name,
            sell_ratio
        );
        assert!(
            hold_ratio >= 0.90 && hold_ratio <= 1.10,
            "Test '{}': HOLD ratio {:.3} deviates >10% from expected",
            name,
            hold_ratio
        );

        println!(
            "✓ No bias test passed for '{}': BUY={} ({:.3}), SELL={} ({:.3}), HOLD={} ({:.3})",
            name,
            action_counts[0],
            buy_ratio,
            action_counts[1],
            sell_ratio,
            action_counts[2],
            hold_ratio
        );
    }

    Ok(())
}

/// Test edge case: sample exactly equals cumulative probability
#[test]
fn test_sample_equals_cumulative_edge_case() -> anyhow::Result<()> {
    // This test verifies the fix for the boundary condition bug
    // When sample == cumulative, the action should be selected
    // But with `<=` instead of `<`, it creates a bias towards lower indices

    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 52;
    config.epsilon_start = 0.0;
    config.epsilon_end = 0.0;
    config.temperature_start = 1.0;

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

    // Use uniform state to get equal probabilities
    let state = vec![0.0; 52];

    // Run large sample to catch edge cases
    let mut action_counts = [0, 0, 0];
    for _ in 0..100000 {
        let action = dqn.select_action(&state)?;
        action_counts[action as usize] += 1;
    }

    // With 100K samples, we should see very tight distribution
    let expected = 100000.0 / 3.0;

    // All actions should be within 2% of expected (stricter tolerance with more samples)
    for (i, &count) in action_counts.iter().enumerate() {
        let ratio = count as f64 / expected;
        assert!(
            ratio >= 0.98 && ratio <= 1.02,
            "Action {} count ({}) deviates >2% from expected ({:.0}), ratio={:.3}",
            i,
            count,
            expected,
            ratio
        );
    }

    // Chi-square test with large sample
    let chi_square: f64 = action_counts
        .iter()
        .map(|&count| {
            let diff = count as f64 - expected;
            (diff * diff) / expected
        })
        .sum();

    assert!(
        chi_square < 5.991,
        "Chi-square test failed with 100K samples: {:.3} > 5.991",
        chi_square
    );

    println!(
        "✓ Edge case test passed (100K samples): BUY={}, SELL={}, HOLD={}, χ²={:.3}",
        action_counts[0], action_counts[1], action_counts[2], chi_square
    );

    Ok(())
}

/// Test cumulative probability normalization (should sum to exactly 1.0)
#[test]
fn test_probability_normalization() -> anyhow::Result<()> {
    // This test verifies that softmax probabilities sum to 1.0
    // Important for the cumulative sampling to work correctly

    use candle_core::{Device, Tensor};
    use candle_nn;

    let device = Device::Cpu;

    // Test multiple Q-value scenarios
    let test_cases = vec![
        ("uniform", vec![0.5, 0.5, 0.5]),
        ("peaked", vec![1.0, 0.0, 0.0]),
        ("mixed", vec![0.7, 0.2, 0.1]),
    ];

    for (name, q_values) in test_cases {
        let q_tensor = Tensor::from_vec(q_values, (1, 3), &device)?;

        // Apply softmax (this is what DQN does internally)
        let probs = candle_nn::ops::softmax(&q_tensor, 1)?;
        let probs_vec = probs.flatten_all()?.to_vec1::<f32>()?;

        // Sum should be exactly 1.0 (within floating point tolerance)
        let sum: f32 = probs_vec.iter().sum();
        assert!(
            (sum - 1.0).abs() < 1e-6,
            "Test '{}': Probabilities don't sum to 1.0: sum={:.10}",
            name,
            sum
        );

        println!(
            "✓ Normalization test passed for '{}': probs={:?}, sum={:.10}",
            name, probs_vec, sum
        );
    }

    Ok(())
}