foxhunt/ml/tests/replay_buffer_test.rs

//! Comprehensive test suite for DQN Experience Replay Buffer
//!
//! Tests cover:
//! - Experience storage and retrieval
//! - Capacity management and FIFO eviction
//! - Uniform random sampling
//! - Prioritized Experience Replay (PER)
//! - Priority updates and importance sampling weights
//! - Edge cases (empty buffer, single experience)
//! - Performance benchmarks

use ml::dqn::replay_buffer::{ReplayBuffer, ReplayBufferConfig};
use ml::dqn::{Experience, ExperienceBatch};
use std::collections::HashSet;
use std::time::Instant;

/// Helper to create a test experience with specific values
fn create_test_experience(state_id: u32, action: u8, reward: f32) -> Experience {
    let state_value = state_id as f32;
    Experience::new(
        vec![state_value, state_value + 0.1, state_value + 0.2],
        action,
        reward,
        vec![state_value + 1.0, state_value + 1.1, state_value + 1.2],
        false,
    )
}

/// Test 1: Buffer stores experiences correctly
#[test]
fn test_experience_storage() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_storage");
    let config = ReplayBufferConfig {
        capacity: 100,
        batch_size: 10,
        min_experiences: 5,
    };

    let buffer = ReplayBuffer::new(path, config)?;
    assert_eq!(buffer.size(), 0);
    assert!(!buffer.can_sample());

    // Add experiences
    for i in 0..10 {
        let exp = create_test_experience(i, (i % 3) as u8, (i * 10) as f32);
        buffer.push(exp)?;
    }

    let stats = buffer.stats();
    assert_eq!(stats.size, 10);
    assert_eq!(stats.experiences_added, 10);
    assert!(buffer.can_sample());

    Ok(())
}

/// Test 2: Buffer respects capacity (FIFO eviction)
#[test]
fn test_capacity_fifo_eviction() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_capacity");
    let config = ReplayBufferConfig {
        capacity: 5,
        batch_size: 2,
        min_experiences: 2,
    };

    let buffer = ReplayBuffer::new(path, config)?;

    // Add more experiences than capacity
    for i in 0..10 {
        let exp = create_test_experience(i, 0, i as f32);
        buffer.push(exp)?;
    }

    let stats = buffer.stats();
    assert_eq!(stats.size, 5, "Buffer size should be capped at capacity");
    assert_eq!(stats.capacity, 5);
    assert_eq!(stats.experiences_added, 10, "Should track all additions");

    // Sample and verify we get the LATEST experiences (5-9, not 0-4)
    let batch = buffer.sample(Some(5))?;
    assert_eq!(batch.batch_size, 5);

    // Extract rewards to verify FIFO eviction
    let rewards: Vec<f32> = batch.experiences.iter().map(|e| e.reward_f32()).collect();
    let mut sorted_rewards = rewards.clone();
    sorted_rewards.sort_by(|a, b| a.partial_cmp(b).unwrap());

    // Should contain rewards 5.0, 6.0, 7.0, 8.0, 9.0 (latest 5 experiences)
    let expected: Vec<f32> = vec![5.0, 6.0, 7.0, 8.0, 9.0];
    for (actual, expected) in sorted_rewards.iter().zip(expected.iter()) {
        assert!(
            (actual - expected).abs() < 0.001,
            "Expected {} but got {}",
            expected,
            actual
        );
    }

    Ok(())
}

/// Test 3: Uniform sampling returns random batch
#[test]
fn test_uniform_sampling() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_uniform");
    let config = ReplayBufferConfig {
        capacity: 100,
        batch_size: 10,
        min_experiences: 20,
    };

    let buffer = ReplayBuffer::new(path, config)?;

    // Add 50 experiences
    for i in 0..50 {
        let exp = create_test_experience(i, 0, i as f32);
        buffer.push(exp)?;
    }

    // Sample multiple times and verify randomness
    let batch1 = buffer.sample(Some(10))?;
    let batch2 = buffer.sample(Some(10))?;

    assert_eq!(batch1.batch_size, 10);
    assert_eq!(batch2.batch_size, 10);

    // Extract rewards
    let rewards1: Vec<f32> = batch1.experiences.iter().map(|e| e.reward_f32()).collect();
    let rewards2: Vec<f32> = batch2.experiences.iter().map(|e| e.reward_f32()).collect();

    // With 50 experiences and batch size 10, it's highly unlikely the same batch
    // is sampled twice (probability ~1e-10)
    assert_ne!(rewards1, rewards2, "Uniform sampling should be random");

    Ok(())
}

/// Test 4: Prioritized sampling favors high-TD-error transitions
#[test]
#[ignore] // Enable when PER is implemented
fn test_prioritized_sampling() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_prioritized");
    let config = ReplayBufferConfig {
        capacity: 100,
        batch_size: 10,
        min_experiences: 20,
    };

    let buffer = ReplayBuffer::new(path, config)?;

    // Add experiences with varying priorities
    // High priority experiences (TD error = 10.0)
    for i in 0..10 {
        let mut exp = create_test_experience(i, 0, i as f32);
        exp.priority = 10.0;
        buffer.push(exp)?;
    }

    // Low priority experiences (TD error = 0.1)
    for i in 10..50 {
        let mut exp = create_test_experience(i, 0, i as f32);
        exp.priority = 0.1;
        buffer.push(exp)?;
    }

    // Sample with prioritization (alpha=0.6)
    let (batch, weights, indices) = buffer.sample_prioritized(20, 0.6, 0.4)?;

    assert_eq!(batch.len(), 20);
    assert_eq!(weights.len(), 20);
    assert_eq!(indices.len(), 20);

    // Count high-priority experiences in batch
    let high_priority_count = batch.iter().filter(|exp| exp.priority > 5.0).count();

    // With alpha=0.6, we expect significantly more high-priority experiences
    // than random chance (random would be ~4 out of 20)
    assert!(
        high_priority_count >= 10,
        "Prioritized sampling should favor high-priority experiences (got {} / 20)",
        high_priority_count
    );

    Ok(())
}

/// Test 5: Priority updates work correctly
#[test]
#[ignore] // Enable when PER is implemented
fn test_priority_updates() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_priority_update");
    let config = ReplayBufferConfig {
        capacity: 100,
        batch_size: 10,
        min_experiences: 5,
    };

    let buffer = ReplayBuffer::new(path, config)?;

    // Add experiences with initial priorities
    for i in 0..20 {
        let mut exp = create_test_experience(i, 0, i as f32);
        exp.priority = 1.0;
        buffer.push(exp)?;
    }

    // Sample and get indices
    let (batch, _weights, indices) = buffer.sample_prioritized(5, 0.6, 0.4)?;

    // Update priorities based on TD errors (simulate learning)
    let new_priorities = vec![5.0, 3.0, 8.0, 1.0, 2.0];
    buffer.update_priorities(indices.clone(), new_priorities.clone())?;

    // Sample again - updated priorities should be reflected
    let (batch2, _weights2, _indices2) = buffer.sample_prioritized(10, 0.6, 0.4)?;

    // The experience with priority 8.0 should appear more frequently
    // This is probabilistic, so we check over multiple samples
    let mut high_priority_samples = 0;
    for _ in 0..100 {
        let (batch, _, _) = buffer.sample_prioritized(5, 0.6, 0.4)?;
        for exp in &batch {
            if exp.priority > 5.0 {
                high_priority_samples += 1;
            }
        }
    }

    // With 1 high-priority (8.0) out of 20 total, uniform would be ~25 / 500
    // Prioritized should be significantly higher
    assert!(
        high_priority_samples > 50,
        "Priority updates should affect sampling (got {} high-priority samples)",
        high_priority_samples
    );

    Ok(())
}

/// Test 6: Importance sampling weights computed correctly
#[test]
#[ignore] // Enable when PER is implemented
fn test_importance_sampling_weights() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_is_weights");
    let config = ReplayBufferConfig {
        capacity: 100,
        batch_size: 10,
        min_experiences: 10,
    };

    let buffer = ReplayBuffer::new(path, config)?;

    // Add experiences with varying priorities
    for i in 0..20 {
        let mut exp = create_test_experience(i, 0, i as f32);
        exp.priority = if i < 5 { 10.0 } else { 1.0 };
        buffer.push(exp)?;
    }

    // Sample with prioritization
    let (batch, weights, _indices) = buffer.sample_prioritized(10, 0.6, 0.4)?;

    assert_eq!(weights.len(), 10);

    // Weights should be normalized (max weight = 1.0)
    let max_weight = weights.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    assert!(
        (max_weight - 1.0).abs() < 1e-6,
        "Weights should be normalized"
    );

    // Low-priority experiences should have higher IS weights (compensate for bias)
    // High-priority experiences should have lower IS weights
    for (exp, weight) in batch.iter().zip(weights.iter()) {
        if exp.priority > 5.0 {
            assert!(
                *weight < 1.0,
                "High-priority experiences should have weight < 1.0"
            );
        }
    }

    Ok(())
}

/// Test 7: Buffer handles edge cases (empty, single experience)
#[test]
fn test_edge_cases() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_edge");
    let config = ReplayBufferConfig {
        capacity: 10,
        batch_size: 5,
        min_experiences: 5,
    };

    let buffer = ReplayBuffer::new(path, config)?;

    // Test 7a: Empty buffer cannot be sampled
    assert!(!buffer.can_sample());
    let result = buffer.sample(Some(5));
    assert!(result.is_err(), "Sampling empty buffer should fail");

    // Test 7b: Single experience
    let exp = create_test_experience(0, 0, 100.0);
    buffer.push(exp.clone())?;
    assert_eq!(buffer.size(), 1);
    assert!(!buffer.can_sample(), "Need min_experiences=5");

    // Test 7c: Exactly min_experiences
    for i in 1..5 {
        let exp = create_test_experience(i, 0, i as f32);
        buffer.push(exp)?;
    }
    assert_eq!(buffer.size(), 5);
    assert!(buffer.can_sample());

    let batch = buffer.sample(Some(3))?;
    assert_eq!(batch.batch_size, 3);

    // Test 7d: Batch size larger than buffer size should fail
    let result = buffer.sample(Some(10));
    assert!(result.is_err(), "Batch size > buffer size should fail");

    Ok(())
}

/// Test 8: Performance - 100K add + 10K sample operations < 1 second
#[test]
fn test_performance_benchmark() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_perf");
    let config = ReplayBufferConfig {
        capacity: 100_000,
        batch_size: 32,
        min_experiences: 1000,
    };

    let buffer = ReplayBuffer::new(path, config)?;

    // Benchmark: 100K additions
    let start = Instant::now();
    for i in 0..100_000 {
        let exp = create_test_experience(i, (i % 3) as u8, (i % 1000) as i32);
        buffer.push(exp)?;
    }
    let add_duration = start.elapsed();

    println!("✅ 100K additions: {:?}", add_duration);

    // Benchmark: 10K samples
    let start = Instant::now();
    for _ in 0..10_000 {
        let _batch = buffer.sample(Some(32))?;
    }
    let sample_duration = start.elapsed();

    println!("✅ 10K samples (batch=32): {:?}", sample_duration);

    let total_duration = add_duration + sample_duration;
    println!("✅ Total time: {:?}", total_duration);

    // Performance target: < 1 second for 110K operations
    assert!(
        total_duration.as_secs_f64() < 1.0,
        "Performance regression: 110K operations took {:?} (target: <1s)",
        total_duration
    );

    // Stats
    let stats = buffer.stats();
    assert_eq!(stats.size, 100_000);
    assert_eq!(stats.experiences_added, 100_000);
    assert_eq!(stats.samples_taken, 10_000);

    Ok(())
}

/// Test 9: Sampling without replacement (no duplicates in batch)
#[test]
fn test_no_duplicate_sampling() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::path::Path::new("/tmp/test_buffer_no_dup");
    let config = ReplayBufferConfig {
        capacity: 50,
        batch_size: 10,
        min_experiences: 10,
    };

    let buffer = ReplayBuffer::new(path, config)?;

    // Add unique experiences
    for i in 0..50 {
        let exp = create_test_experience(i, 0, i as f32);
        buffer.push(exp)?;
    }

    // Sample and verify no duplicates
    for _ in 0..100 {
        let batch = buffer.sample(Some(20))?;

        // Use integer conversion for HashSet (f32 can't be hashed directly)
        let rewards: Vec<i32> = batch
            .experiences
            .iter()
            .map(|e| (e.reward_f32() * 10.0) as i32)
            .collect();
        let unique_rewards: HashSet<i32> = rewards.iter().cloned().collect();

        assert_eq!(
            rewards.len(),
            unique_rewards.len(),
            "Batch should not contain duplicates"
        );
    }

    Ok(())
}

/// Test 10: Thread-safety (concurrent push and sample)
#[test]
fn test_thread_safety() -> Result<(), Box<dyn std::error::Error>> {
    use std::sync::Arc;
    use std::thread;

    let path = std::path::Path::new("/tmp/test_buffer_threads");
    let config = ReplayBufferConfig {
        capacity: 10_000,
        batch_size: 32,
        min_experiences: 100,
    };

    let buffer = Arc::new(ReplayBuffer::new(path, config)?);

    // Pre-fill buffer to enable sampling
    for i in 0..200 {
        let exp = create_test_experience(i, 0, i as f32);
        buffer.push(exp)?;
    }

    // Spawn writer threads
    let mut handles = vec![];
    for thread_id in 0..4 {
        let buffer_clone = Arc::clone(&buffer);
        let handle = thread::spawn(move || {
            for i in 0..1000 {
                let exp = create_test_experience((thread_id * 1000 + i) as u32, 0, i as i32);
                buffer_clone.push(exp).unwrap();
            }
        });
        handles.push(handle);
    }

    // Spawn reader threads
    for _ in 0..2 {
        let buffer_clone = Arc::clone(&buffer);
        let handle = thread::spawn(move || {
            for _ in 0..500 {
                let _batch = buffer_clone.sample(Some(32)).unwrap();
            }
        });
        handles.push(handle);
    }

    // Wait for all threads
    for handle in handles {
        handle.join().unwrap();
    }

    let stats = buffer.stats();
    assert_eq!(stats.size, 4200); // 200 initial + 4*1000 added
    assert_eq!(stats.samples_taken, 1000); // 2*500 samples

    Ok(())
}