#![allow(unsafe_code)] // Required for CUDA DtoD memcpy via cudarc.

//! Enum wrapper for runtime replay buffer type switching
//!
//! Enables Prioritized Experience Replay (PER) activation based on configuration.
//! Provides unified API for both uniform and prioritized sampling.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use parking_lot::Mutex;

use crate::dqn::ExperienceReplayBuffer;
use crate::prioritized_replay::PrioritizedReplayBuffer;
use crate::experience::Experience;
use ml_core::cuda_autograd::GpuTensor;
use ml_core::MLError;

/// Batch sample with importance sampling weights
#[derive(Debug)]
pub struct BatchSample {
    pub experiences: Vec<Experience>,
    pub weights: Vec<f32>,  // Importance sampling weights (1.0 for uniform)
    pub indices: Vec<usize>, // Indices for priority updates (empty for uniform)
    pub gpu_batch: Option<GpuBatch>,
}

impl BatchSample {
    /// Create a batch with uniform weights (no importance sampling correction).
    /// Use this when experiences were not sampled via PER.
    pub fn uniform(experiences: Vec<Experience>) -> Self {
        let len = experiences.len();
        Self {
            experiences,
            weights: vec![1.0; len],
            indices: vec![],
            gpu_batch: None,
        }
    }
}

/// Pre-built GPU tensors for a training batch.
/// When present, `compute_gradients()` uses these directly -- no CPU->GPU transfer.
#[derive(Debug)]
pub struct GpuBatch {
    pub states: GpuTensor,      // [batch_size, state_dim] f32 on GPU
    pub actions: GpuTensor,     // [batch_size] u32 on GPU
    pub rewards: GpuTensor,     // [batch_size] f32 on GPU
    pub next_states: GpuTensor, // [batch_size, state_dim] f32 on GPU
    pub dones: GpuTensor,       // [batch_size] f32 on GPU (0.0/1.0)
    pub weights: GpuTensor,     // [batch_size] f32 on GPU (IS weights)
    pub indices: GpuTensor,     // [batch_size] u32 on GPU (buffer indices)
    /// Episode IDs per transition `[batch_size]` i32 on GPU.
    ///
    /// `episode_ids[i] = buffer_index[i] / episode_length`. Required by HER
    /// Future and Final strategies for GPU-native donor selection.
    /// `None` when episode tracking is disabled (HER Random or no HER).
    pub episode_ids: Option<cudarc::driver::CudaSlice<i32>>,
}

/// GPU buffer + CPU staging: `add()` stages on CPU (zero GPU ops),
/// `sample()` batch-flushes staging -> GPU in one DMA before sampling.
pub struct StagedGpuBuffer {
    pub gpu: crate::gpu_replay_buffer::GpuReplayBuffer,
    staging: Vec<Experience>,
}

impl StagedGpuBuffer {
    fn flush(&mut self) -> Result<(), MLError> {
        if self.staging.is_empty() { return Ok(()); }
        let cap = self.gpu.capacity();
        if self.staging.len() > cap {
            let drain_count = self.staging.len() - cap;
            self.staging.drain(..drain_count);
        }
        let n = self.staging.len();
        let sd = self.gpu.state_dim();
        let stream = self.gpu.stream().clone();

        let mut states_flat = Vec::with_capacity(n * sd);
        let mut next_flat = Vec::with_capacity(n * sd);
        let mut actions = Vec::with_capacity(n);
        let mut rewards = Vec::with_capacity(n);
        let mut dones = Vec::with_capacity(n);

        for exp in self.staging.drain(..) {
            states_flat.extend_from_slice(&exp.state);
            if exp.state.len() < sd {
                states_flat.resize(states_flat.len() + sd - exp.state.len(), 0.0);
            }
            next_flat.extend_from_slice(&exp.next_state);
            if exp.next_state.len() < sd {
                next_flat.resize(next_flat.len() + sd - exp.next_state.len(), 0.0);
            }
            actions.push(exp.action as u32);
            rewards.push(exp.reward_f32());
            dones.push(if exp.done { 1.0_f32 } else { 0.0 });
        }

        // Upload to GPU CudaSlice buffers (insert_batch casts F32->BF16 internally)
        let mut sg = stream.alloc_zeros::<f32>(n * sd).map_err(|e| MLError::TrainingError(format!("alloc s: {e}")))?;
        let mut ng = stream.alloc_zeros::<f32>(n * sd).map_err(|e| MLError::TrainingError(format!("alloc n: {e}")))?;
        let mut ag = stream.alloc_zeros::<u32>(n).map_err(|e| MLError::TrainingError(format!("alloc a: {e}")))?;
        let mut rg = stream.alloc_zeros::<f32>(n).map_err(|e| MLError::TrainingError(format!("alloc r: {e}")))?;
        let mut dg = stream.alloc_zeros::<f32>(n).map_err(|e| MLError::TrainingError(format!("alloc d: {e}")))?;
        stream.memcpy_htod(&states_flat, &mut sg).map_err(|e| MLError::TrainingError(format!("htod s: {e}")))?;
        stream.memcpy_htod(&next_flat, &mut ng).map_err(|e| MLError::TrainingError(format!("htod n: {e}")))?;
        stream.memcpy_htod(&actions, &mut ag).map_err(|e| MLError::TrainingError(format!("htod a: {e}")))?;
        stream.memcpy_htod(&rewards, &mut rg).map_err(|e| MLError::TrainingError(format!("htod r: {e}")))?;
        stream.memcpy_htod(&dones, &mut dg).map_err(|e| MLError::TrainingError(format!("htod d: {e}")))?;
        self.gpu.insert_batch(&sg, &ng, &ag, &rg, &dg, n)
    }
}

/// Enum wrapper for runtime buffer type switching.
///
/// Clone is cheap (Arc inner types) and enables sharing the buffer between
/// the training thread and a background pre-sample thread for async PER.
#[derive(Clone)]
pub enum ReplayBufferType {
    /// Uniform random sampling (standard DQN)
    Uniform(Arc<Mutex<ExperienceReplayBuffer>>),
    /// Prioritized sampling based on TD errors (Rainbow DQN)
    Prioritized(Arc<PrioritizedReplayBuffer>),
    /// GPU-resident prioritized replay with CPU staging buffer.
    /// `add()` stages on CPU (zero GPU ops). `sample()` flushes staging -> GPU
    /// in one batch DMA, then samples entirely on GPU.
    GpuPrioritized(Arc<Mutex<StagedGpuBuffer>>),
}

impl std::fmt::Debug for ReplayBufferType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Uniform(_) => write!(f, "ReplayBufferType::Uniform"),
            Self::Prioritized(_) => write!(f, "ReplayBufferType::Prioritized"),
            Self::GpuPrioritized(_) => write!(f, "ReplayBufferType::GpuPrioritized"),
        }
    }
}

impl ReplayBufferType {
    /// Create uniform replay buffer
    pub fn new_uniform(capacity: usize) -> Self {
        Self::Uniform(Arc::new(Mutex::new(ExperienceReplayBuffer::new(capacity))))
    }

    /// Create prioritized replay buffer
    pub fn new_prioritized(
        capacity: usize,
        alpha: f64,
        beta: f64,
        beta_max: f64,
        beta_annealing_steps: usize,
    ) -> Result<Self, MLError> {
        use crate::prioritized_replay::{PrioritizedReplayConfig, PrioritizationStrategy};

        let config = PrioritizedReplayConfig {
            capacity,
            alpha: alpha as f32,
            beta: beta as f32,
            initial_priority: 1.0,
            min_priority: 1e-6,
            strategy: PrioritizationStrategy::Proportional,
            beta_max: beta_max as f32,
            beta_annealing_steps,
        };

        let buffer = PrioritizedReplayBuffer::new(config)?;
        Ok(Self::Prioritized(Arc::new(buffer)))
    }

    /// Create GPU-resident prioritized replay buffer.
    ///
    /// All experience data lives as contiguous GPU tensors. Sampling and
    /// priority updates happen entirely on GPU -- zero CPU round-trips.
    pub fn new_gpu_prioritized(
        capacity: usize,
        state_dim: usize,
        alpha: f64,
        beta: f64,
        beta_max: f64,
        beta_annealing_steps: usize,
        max_memory_bytes: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        use crate::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};

        let config = GpuReplayBufferConfig {
            capacity,
            state_dim,
            alpha: alpha as f32,
            beta_start: beta as f32,
            beta_max: beta_max as f32,
            beta_annealing_steps,
            epsilon: 1e-6,
            max_memory_bytes,
            max_batch_size: 1024,
        };

        let buffer = GpuReplayBuffer::new(config, stream)?;
        Ok(Self::GpuPrioritized(Arc::new(Mutex::new(StagedGpuBuffer {
            gpu: buffer,
            staging: Vec::new(),
        }))))
    }

    /// Attempt GPU PER allocation with adaptive capacity halving on OOM.
    /// Returns the buffer on success, or a hard error on exhaustion.
    /// No CPU PER fallback -- GPU PER is mandatory on CUDA.
    pub fn try_gpu_with_halving(
        capacity: usize,
        state_dim: usize,
        alpha: f64,
        beta: f64,
        beta_max: f64,
        beta_annealing_steps: usize,
        max_memory_bytes: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        const MIN_GPU_CAPACITY: usize = 1024;
        let mut try_cap = capacity;

        while try_cap >= MIN_GPU_CAPACITY {
            match Self::new_gpu_prioritized(try_cap, state_dim, alpha, beta, beta_max, beta_annealing_steps, max_memory_bytes, stream) {
                Ok(buf) => {
                    if try_cap < capacity {
                        tracing::warn!(
                            "GPU PER replay buffer allocated at reduced capacity ({} -> {}, {} state_dim)",
                            capacity, try_cap, state_dim
                        );
                    } else {
                        tracing::info!(
                            "GPU PER replay buffer allocated ({} capacity, {} state_dim)",
                            try_cap, state_dim
                        );
                    }
                    return Ok(buf);
                }
                Err(e) => {
                    tracing::warn!(
                        "GPU PER at capacity {} failed ({}), retrying at {}",
                        try_cap, e, try_cap / 2
                    );
                    try_cap /= 2;
                    if try_cap < MIN_GPU_CAPACITY {
                        return Err(e);
                    }
                }
            }
        }

        Err(MLError::ModelError("GPU PER capacity below floor".into()))
    }

    /// Sample a batch from the buffer
    pub fn sample(&self, batch_size: usize) -> Result<BatchSample, MLError> {
        match self {
            Self::Uniform(buffer) => {
                let buffer = buffer.lock();
                let experiences = buffer.sample(batch_size)?;
                Ok(BatchSample {
                    weights: vec![1.0; experiences.len()], // Uniform weights
                    indices: vec![], // Not used for uniform
                    experiences,
                    gpu_batch: None,
                })
            }
            Self::Prioritized(buffer) => {
                let (experiences, weights, indices) = buffer.sample(batch_size)?;
                Ok(BatchSample {
                    experiences,
                    weights,
                    indices,
                    gpu_batch: None,
                })
            }
            Self::GpuPrioritized(buffer) => {
                let mut buf = buffer.lock();
                buf.flush()?;
                let stream = buf.gpu.stream().clone();
                let gpu_batch_slices = buf.gpu.sample_proportional(batch_size)?;
                let gpu_batch = gpu_batch_slices.into_gpu_batch(&stream)?;
                Ok(BatchSample {
                    experiences: vec![], // Empty -- data is on GPU
                    weights: vec![],     // Empty -- weights on GPU
                    indices: vec![],     // Empty -- indices on GPU
                    gpu_batch: Some(gpu_batch),
                })
            }
        }
    }

    /// Add experience to buffer
    pub fn add(&self, experience: Experience) -> Result<(), MLError> {
        match self {
            Self::Uniform(buffer) => {
                let mut buffer = buffer.lock();
                buffer.push(experience);
                Ok(())
            }
            Self::Prioritized(buffer) => {
                buffer.push(experience)
            }
            Self::GpuPrioritized(buf) => {
                buf.lock().staging.push(experience);
                Ok(())
            }
        }
    }

    /// Add a batch of experiences to buffer with a single lock acquisition.
    pub fn add_batch(&self, experiences: Vec<Experience>) -> Result<(), MLError> {
        match self {
            Self::Uniform(buffer) => {
                let mut buffer = buffer.lock();
                for exp in experiences {
                    buffer.push(exp);
                }
                Ok(())
            }
            Self::Prioritized(buffer) => {
                for exp in experiences {
                    buffer.push(exp)?;
                }
                Ok(())
            }
            Self::GpuPrioritized(buf) => {
                buf.lock().staging.extend(experiences);
                Ok(())
            }
        }
    }

    /// Update priorities (PER only)
    pub fn update_priorities(&self, indices: &[usize], td_errors: &[f32]) -> Result<(), MLError> {
        match self {
            Self::Uniform(_) => Ok(()), // No-op for uniform
            Self::Prioritized(buffer) => {
                // Convert TD errors to priorities (absolute value)
                let priorities: Vec<f32> = td_errors.iter().map(|&td| td.abs()).collect();
                buffer.update_priorities(indices, &priorities)
            }
            Self::GpuPrioritized(_) => {
                // Use update_priorities_gpu() with tensor inputs instead
                Ok(())
            }
        }
    }

    /// Update priorities from GPU-resident TD error tensors (`GpuPrioritized` only).
    ///
    /// Takes `GpuTensor` indices and TD errors directly -- no `to_vec1()` needed.
    ///
    /// Indices are expected to hold f32-encoded u32 values (cast to u32 internally).
    pub fn update_priorities_gpu(
        &self,
        indices: &GpuTensor,
        td_errors: &GpuTensor,
    ) -> Result<(), MLError> {
        match self {
            Self::GpuPrioritized(buffer) => {
                let mut buf = buffer.lock();
                let bs = td_errors.shape().first().copied().unwrap_or(0);
                if bs == 0 { return Ok(()); }
                // Convert f32-encoded indices to CudaSlice<u32> via GPU cast kernel (zero CPU download)
                let stream = buf.gpu.stream().clone();
                let cast_kernels = crate::gpu_replay_buffer::get_cast_kernels_f32_to_u32(&stream)?;
                let idx_gpu = cast_kernels.cast_f32_to_u32(indices.data(), bs, &stream)?;
                buf.gpu.update_priorities_gpu(&idx_gpu, td_errors.data(), bs)
            }
            Self::Uniform(_) | Self::Prioritized(_) => Ok(()), // No-op for non-GPU buffers
        }
    }

    /// Flush GPU-accumulated max priority to CPU (single scalar readback).
    ///
    /// Call once per epoch after all `update_priorities_gpu` calls.
    /// No-op for non-GPU buffers.
    pub fn flush_max_priority(&self) -> Result<(), MLError> {
        match self {
            Self::GpuPrioritized(buffer) => {
                buffer.lock().gpu.flush_max_priority()
            }
            Self::Uniform(_) | Self::Prioritized(_) => Ok(()),
        }
    }

    /// Direct reference to the GPU-resident priorities tensor (`GpuPrioritized` only).
    ///
    /// Returns `None` for non-GPU buffers. Used by fused CUDA training to pass
    /// the priorities tensor to `GpuDqnTrainer::update_priorities_cuda()`.
    pub fn priorities_tensor(&self) -> Option<GpuTensor> {
        match self {
            Self::GpuPrioritized(buffer) => {
                let buf = buffer.lock();
                let slice = buf.gpu.priorities_slice();
                let n = slice.len();
                let stream = buf.gpu.stream();
                // GPU DtoD clone (zero CPU download)
                let gpu_tensor = GpuTensor::new(
                    {
                        // DtoD copy of the raw CudaSlice<f32>
                        use cudarc::driver::{DevicePtr, DevicePtrMut};
                        let mut dst = stream.alloc_zeros::<f32>(n).ok()?;
                        let num_bytes = n * std::mem::size_of::<f32>();
                        let src_ptr = {
                            let (ptr, guard) = slice.device_ptr(stream);
                            let _no_drop = std::mem::ManuallyDrop::new(guard);
                            ptr
                        };
                        let dst_ptr = {
                            let (ptr, guard) = dst.device_ptr_mut(stream);
                            let _no_drop = std::mem::ManuallyDrop::new(guard);
                            ptr
                        };
                        // SAFETY: dst_ptr and src_ptr are valid device pointers. num_bytes = n * sizeof(f32).
                        unsafe {
                            cudarc::driver::result::memcpy_dtod_async(
                                dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
                            ).ok()?;
                        }
                        dst
                    },
                    vec![n],
                ).ok()?;
                Some(gpu_tensor)
            }
            Self::Uniform(_) | Self::Prioritized(_) => None,
        }
    }

    /// PER alpha/epsilon for GPU-native priority update kernel.
    ///
    /// Returns `(alpha, epsilon)` or `None` for non-GPU buffers.
    pub fn per_alpha_epsilon(&self) -> Option<(f32, f32)> {
        match self {
            Self::GpuPrioritized(buffer) => {
                let buf = buffer.lock();
                Some((buf.gpu.alpha(), buf.gpu.epsilon()))
            }
            Self::Uniform(_) | Self::Prioritized(_) => None,
        }
    }

    /// Set `max_priority` from a CPU scalar (epoch-boundary flush from CUDA kernel).
    pub fn apply_max_priority_scalar(&self, max_prio: f32) -> Result<(), MLError> {
        match self {
            Self::GpuPrioritized(buffer) => {
                buffer.lock().gpu.apply_max_priority_scalar(max_prio)
            }
            Self::Uniform(_) | Self::Prioritized(_) => Ok(()),
        }
    }

    /// Step training counter for beta annealing (PER only)
    pub fn step(&self) {
        match self {
            Self::Prioritized(buffer) => buffer.step(),
            Self::GpuPrioritized(buffer) => buffer.lock().gpu.step(),
            _ => {}
        }
    }

    /// M2: Get indices of stale experiences for priority refresh (PER only).
    pub fn get_stale_indices(&self, max_age: usize, limit: usize) -> Vec<usize> {
        match self {
            Self::Prioritized(buffer) => buffer.get_stale_indices(max_age, limit),
            Self::Uniform(_) => vec![],
            Self::GpuPrioritized(_) => vec![],
        }
    }

    /// M2: Get experiences at specific indices for priority recomputation (PER only).
    pub fn get_experiences_at(&self, indices: &[usize]) -> Vec<Option<Experience>> {
        match self {
            Self::Prioritized(buffer) => buffer.get_experiences_at(indices),
            Self::Uniform(_) => vec![],
            Self::GpuPrioritized(_) => vec![],
        }
    }

    /// Get current beta value (PER only)
    pub fn current_beta(&self) -> f32 {
        match self {
            Self::Uniform(_) => 1.0,
            Self::Prioritized(buffer) => buffer.current_beta(),
            Self::GpuPrioritized(buffer) => buffer.lock().gpu.current_beta(),
        }
    }

    /// Get buffer size
    pub fn len(&self) -> usize {
        match self {
            Self::Uniform(buffer) => buffer.lock().len(),
            Self::Prioritized(buffer) => buffer.len(),
            Self::GpuPrioritized(buffer) => {
                let buf = buffer.lock();
                buf.gpu.len() + buf.staging.len()
            }
        }
    }

    /// Check if buffer is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// BUG #38 FIX: Clear all experiences from the buffer
    pub fn clear(&self) -> Result<(), MLError> {
        match self {
            Self::Uniform(buffer) => {
                let mut buffer = buffer.lock();
                buffer.clear();
            }
            Self::Prioritized(buffer) => {
                buffer.clear();
            }
            Self::GpuPrioritized(buffer) => {
                let mut buf = buffer.lock();
                buf.gpu.clear()?;
                buf.staging.clear();
            }
        }
        Ok(())
    }

    /// Check if buffer can sample given batch size
    pub fn can_sample(&self, batch_size: usize) -> bool {
        match self {
            Self::Uniform(buffer) => buffer.lock().len() >= batch_size,
            Self::Prioritized(buffer) => buffer.can_sample(batch_size),
            Self::GpuPrioritized(buffer) => {
                let buf = buffer.lock();
                buf.gpu.len() + buf.staging.len() >= batch_size
            }
        }
    }

    /// Check if using prioritized replay
    pub const fn is_prioritized(&self) -> bool {
        matches!(self, Self::Prioritized(_))
    }

    /// Check if using GPU-resident prioritized replay.
    pub const fn is_gpu_prioritized(&self) -> bool {
        matches!(self, Self::GpuPrioritized(_))
    }

    /// Adaptive buffer sizing based on epsilon decay
    pub fn adaptive_resize(&mut self, epsilon: f64, max_capacity: usize) -> Result<bool, MLError> {
        const BASE_CAPACITY: usize = 10_000;
        const RESIZE_THRESHOLDS: [f64; 5] = [0.9, 0.7, 0.5, 0.3, 0.1];

        // Calculate new capacity using growth formula
        let growth_factor = 1.0 - epsilon;
        let new_capacity = BASE_CAPACITY +
            ((max_capacity - BASE_CAPACITY) as f64 * growth_factor) as usize;

        // Get current capacity
        let current_capacity = self.get_capacity();

        // Only grow buffer (never shrink to prevent data loss)
        if new_capacity <= current_capacity {
            return Ok(false);
        }

        // Check if we crossed a threshold (avoid micro-resizes)
        let should_resize = RESIZE_THRESHOLDS.iter().any(|&threshold| {
            epsilon <= threshold && current_capacity < self.capacity_at_threshold(threshold, max_capacity)
        });

        if !should_resize {
            return Ok(false);
        }

        // Perform resize based on buffer type
        match self {
            Self::Uniform(ref buffer) => {
                let mut buf = buffer.lock();
                buf.resize(new_capacity)?;
            }
            Self::Prioritized(_) => {
                // PER resize not yet implemented (requires rebuilding sum tree)
                return Ok(false);
            }
            Self::GpuPrioritized(_) => {
                // GPU buffer is fixed-capacity (pre-allocated tensors)
                return Ok(false);
            }
        }

        tracing::info!(
            "Adaptive buffer resize: {} -> {} samples (eps={:.3}, growth={:.1}%)",
            current_capacity,
            new_capacity,
            epsilon,
            ((new_capacity - current_capacity) as f64 / current_capacity as f64) * 100.0
        );

        Ok(true)
    }

    /// Get current buffer capacity
    pub fn get_capacity(&self) -> usize {
        match self {
            Self::Uniform(buffer) => buffer.lock().capacity(),
            Self::Prioritized(buffer) => buffer.capacity(),
            Self::GpuPrioritized(buffer) => buffer.lock().gpu.capacity(),
        }
    }

    /// Get a locked reference to the staged GPU buffer (`GpuPrioritized` only).
    pub fn as_gpu_buffer(&self) -> Option<parking_lot::MutexGuard<'_, StagedGpuBuffer>> {
        match self {
            Self::GpuPrioritized(buffer) => Some(buffer.lock()),
            Self::Uniform(_) | Self::Prioritized(_) => None,
        }
    }

    /// Calculate capacity at a given epsilon threshold
    pub fn capacity_at_threshold(&self, threshold: f64, max_capacity: usize) -> usize {
        const BASE_CAPACITY: usize = 10_000;
        let growth_factor = 1.0 - threshold;
        BASE_CAPACITY + ((max_capacity - BASE_CAPACITY) as f64 * growth_factor) as usize
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_test_experience(reward: f32) -> Experience {
        let state = vec![0.1, 0.2, 0.3, 0.4, 0.5];
        let next_state = vec![0.2, 0.3, 0.4, 0.5, 0.6];
        Experience::new(state, 0, reward, next_state, false)
    }

    #[test]
    fn test_uniform_buffer_creation() {
        let buffer = ReplayBufferType::new_uniform(1000);
        assert!(!buffer.is_prioritized());
        assert_eq!(buffer.len(), 0);
        assert!(buffer.is_empty());
    }

    #[test]
    fn test_prioritized_buffer_creation() -> Result<(), MLError> {
        let buffer = ReplayBufferType::new_prioritized(1000, 0.6, 0.4, 1.0, 10000)?;
        assert!(buffer.is_prioritized());
        assert_eq!(buffer.len(), 0);
        assert!(buffer.is_empty());
        Ok(())
    }

    #[test]
    fn test_uniform_add_and_sample() -> Result<(), MLError> {
        let buffer = ReplayBufferType::new_uniform(100);

        // Add experiences
        for i in 0..50 {
            buffer.add(create_test_experience(i as f32))?;
        }

        assert_eq!(buffer.len(), 50);
        assert!(buffer.can_sample(32));

        // Sample batch
        let batch = buffer.sample(32)?;
        assert_eq!(batch.experiences.len(), 32);
        assert_eq!(batch.weights.len(), 32);
        assert!(batch.indices.is_empty()); // No indices for uniform

        // All weights should be 1.0 for uniform
        assert!(batch.weights.iter().all(|&w| (w - 1.0).abs() < 1e-6));

        Ok(())
    }

    #[test]
    fn test_prioritized_add_and_sample() -> Result<(), MLError> {
        let buffer = ReplayBufferType::new_prioritized(100, 0.6, 0.4, 1.0, 10000)?;

        // Add experiences
        for i in 0..50 {
            buffer.add(create_test_experience(i as f32))?;
        }

        assert_eq!(buffer.len(), 50);
        assert!(buffer.can_sample(32));

        // Sample batch
        let batch = buffer.sample(32)?;
        assert_eq!(batch.experiences.len(), 32);
        assert_eq!(batch.weights.len(), 32);
        assert_eq!(batch.indices.len(), 32); // Indices provided for PER

        // Weights should be positive
        assert!(batch.weights.iter().all(|&w| w > 0.0));

        Ok(())
    }

    #[test]
    fn test_priority_updates() -> Result<(), MLError> {
        let buffer = ReplayBufferType::new_prioritized(100, 0.6, 0.4, 1.0, 10000)?;

        // Add experiences
        for i in 0..50 {
            buffer.add(create_test_experience(i as f32))?;
        }

        // Sample and update priorities
        let batch = buffer.sample(10)?;
        let td_errors = vec![5.0, 2.0, 8.0, 1.0, 3.0, 0.5, 10.0, 0.1, 4.0, 6.0];
        buffer.update_priorities(&batch.indices, &td_errors)?;

        // No error should occur
        Ok(())
    }

    #[test]
    fn test_beta_annealing() -> Result<(), MLError> {
        let buffer = ReplayBufferType::new_prioritized(100, 0.6, 0.4, 1.0, 1000)?;

        // Initial beta
        let initial_beta = buffer.current_beta();
        assert!((initial_beta - 0.4).abs() < 1e-6);

        // Step 500 times
        for _ in 0..500 {
            buffer.step();
        }

        // Beta should have increased
        let mid_beta = buffer.current_beta();
        assert!(mid_beta > initial_beta);
        assert!(mid_beta < 1.0);

        Ok(())
    }
}
