//! GPU Pipeline: Pre-uploaded training data for DQN/PPO trainers.
//!
//! Eliminates per-bar heap allocations by uploading all training data
//! to GPU at epoch start. Both DQN and PPO trainers index into these
//! pre-uploaded buffers instead of creating per-step copies.
//!
//! All modules operate on raw `CudaSlice` via cudarc.
//! `DqnGpuData` and `PpoGpuData` hold GPU-resident `CudaSlice<f32>` buffers.

use std::sync::Arc;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
use crate::MLError;

pub mod double_buffer;
pub mod multi_gpu;
pub mod signal_adapter;

pub mod gpu_portfolio;
pub mod gpu_weights;
pub mod gpu_experience_collector;
pub mod gpu_ppo_collector;
pub mod gpu_action_selector;
pub mod gpu_statistics;
pub mod gpu_training_guard;
pub mod gpu_monitoring;
pub mod gpu_backtest_evaluator;
pub mod gpu_curiosity_trainer;
pub mod gpu_walk_forward;
pub mod gpu_dqn_trainer;
pub mod batched_forward;
pub mod batched_backward;
pub mod gpu_her;
pub mod gpu_iql_trainer;
pub mod gpu_iqn_head;
pub mod gpu_attention;
pub mod decision_transformer;
// gpu_replay_buffer moved to ml-dqn crate

/// Maximum bytes allowed for a single GPU upload (2 GB safety limit).
const MAX_UPLOAD_BYTES: usize = 2 * 1024 * 1024 * 1024;

/// Round dimension up to nearest multiple of 8 for tensor core alignment.
///
/// Tensor cores on Ampere (A100, L4, L40S) and Hopper (H100) require
/// matrix dimensions divisible by 8 (FP16/BF16) for optimal throughput.
/// This pads dimensions like state_dim=45 → 48 to enable tensor core paths.
pub fn align_to_tensor_cores(dim: usize) -> usize {
    (dim + 7) & !7
}

/// Estimate VRAM usage in bytes for pre-uploaded f32 data.
pub fn estimate_vram_bytes(num_elements: usize) -> usize {
    num_elements * std::mem::size_of::<f32>()
}

/// Compute optimal (grid_dim, block_dim) for a 1-D kernel launch.
///
/// Selects the largest block size that divides 32 (warp size) and does not
/// exceed the GPU's `max_threads_per_block`. On data-centre GPUs with
/// high SM counts (H100: 132 SMs, A100: 108) the larger block sizes
/// improve occupancy vs. the default 32.
///
/// # Arguments
/// * `n_items` - Total work items (e.g. number of episodes)
/// * `max_threads_per_block` - Device limit (typically 1024). Pass 0 for default (256).
///
/// # Returns
/// `(grid_x, block_x)` tuple suitable for `LaunchConfig`.
pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u32) {
    let max_tpb = if max_threads_per_block == 0 { 256 } else { max_threads_per_block };

    // Clamp to a warp-multiple. For experience-collection kernels each thread
    // runs one full episode, so very large blocks aren't needed. Cap at 256 to
    // keep register pressure manageable while still improving over 32.
    let block = if n_items >= 256 && max_tpb >= 256 {
        256
    } else if n_items >= 128 && max_tpb >= 128 {
        128
    } else if n_items >= 64 && max_tpb >= 64 {
        64
    } else {
        32
    };

    let grid = n_items.saturating_add(block - 1) / block;
    (grid, block)
}

/// Compile CUDA source to native cubin with architecture-aware options and disk caching.
///
/// Compiles to **cubin** (native SASS binary), NOT PTX. This eliminates driver
/// JIT entirely — `cuModuleLoad` loads the cubin directly. TMA instructions
/// (cp.async.bulk on sm_90+) compile natively because nvcc handles them during
/// offline compilation; the driver never sees PTX.
///
/// On first compilation, nvcc compiles the source to cubin and caches the result
/// to `$CARGO_TARGET_DIR/.cubin_cache/` (CI PVC) or `/tmp/.cubin_cache/` (fallback).
/// Subsequent runs with identical source skip compilation entirely by loading
/// the cached cubin file. The cache key is a SHA-256 hash of (arch, source, flags),
/// so any change to the kernel source or network dimensions invalidates the cache.
///
/// The fused experience collector kernel (4490 lines, branching+C51+NoisyNets)
// Re-export precompilation from ml-core so existing callers
// (`crate::cuda_pipeline::compile_ptx_for_device`) keep working.
pub use ml_core::cuda_compile::compile_ptx_for_device;

/// Clone a `CudaSlice<f32>` via DtoD memcpy.
///
/// The caller owns the returned slice independently of the source.
pub fn clone_cuda_slice_f32(
    src: &CudaSlice<f32>,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, crate::MLError> {
    let n = src.len();
    let mut dst = stream.alloc_zeros::<f32>(n).map_err(|e| {
        crate::MLError::ModelError(format!("clone_cuda_slice_f32 alloc: {e}"))
    })?;
    {
        let (src_ptr, _src_guard) = src.device_ptr(stream);
        let (dst_ptr, _dst_guard) = dst.device_ptr_mut(stream);
        let num_bytes = n * std::mem::size_of::<f32>();
        #[allow(unsafe_code)]
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
            ).map_err(|e| crate::MLError::ModelError(format!("clone_cuda_slice_f32 DtoD: {e}")))?;
        }
    }
    Ok(dst)
}

/// Clone a `CudaSlice<u32>` via DtoD memcpy.
pub fn clone_cuda_slice_u32(
    src: &CudaSlice<u32>,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<u32>, crate::MLError> {
    let n = src.len();
    let mut dst = stream.alloc_zeros::<u32>(n).map_err(|e| {
        crate::MLError::ModelError(format!("clone_cuda_slice_u32 alloc: {e}"))
    })?;
    {
        let (src_ptr, _src_guard) = src.device_ptr(stream);
        let (dst_ptr, _dst_guard) = dst.device_ptr_mut(stream);
        let num_bytes = n * std::mem::size_of::<u32>();
        #[allow(unsafe_code)]
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
            ).map_err(|e| crate::MLError::ModelError(format!("clone_cuda_slice_u32 DtoD: {e}")))?;
        }
    }
    Ok(dst)
}

/// Pre-uploaded GPU training data for DQN trainer.
///
/// Holds market features [N * 42] and target prices [N * 4] as GPU-resident
/// `CudaSlice<f32>` buffers. Portfolio features (3 dims) are computed per-bar
/// by the trainer and concatenated on-device. OFI features (8 dims from MBP-10
/// order book) are optionally uploaded and concatenated to produce [N, 51] states.
pub struct DqnGpuData {
    /// Market features [num_bars * 42] on GPU (f32, row-major)
    pub features: CudaSlice<f32>,
    /// Target prices [num_bars * 4] on GPU (f32, row-major)
    pub targets: CudaSlice<f32>,
    /// OFI features [num_bars * 8] on GPU (f32), from MBP-10 data
    pub ofi_features: Option<CudaSlice<f32>>,
    /// Number of training bars
    pub num_bars: usize,
    /// Feature dimension (42)
    pub feature_dim: usize,
    /// Aligned state dimension (multiple of 8 for tensor core HMMA).
    pub aligned_state_dim: Option<usize>,
}

impl std::fmt::Debug for DqnGpuData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DqnGpuData")
            .field("num_bars", &self.num_bars)
            .field("feature_dim", &self.feature_dim)
            .field("aligned_state_dim", &self.aligned_state_dim)
            .field("has_ofi", &self.ofi_features.is_some())
            .finish()
    }
}

impl DqnGpuData {
    /// Upload DQN training data to GPU as f32 CudaSlice buffers.
    ///
    /// Converts `[f64; 42]` features and `Vec<f64>` targets to f32,
    /// flattens into contiguous arrays, and uploads once via `clone_htod`.
    pub fn upload(
        data: &[([f64; 42], Vec<f64>)],
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let num_bars = data.len();
        if num_bars == 0 {
            return Err(MLError::ModelError("Empty training data".to_owned()));
        }

        let feature_dim = 42;
        let target_dim = 4;

        let estimated_bytes = estimate_vram_bytes(num_bars * (feature_dim + target_dim));
        if estimated_bytes > MAX_UPLOAD_BYTES {
            return Err(MLError::ModelError(format!(
                "Training data too large for GPU upload: {:.1} GB > 2.0 GB limit ({} bars)",
                estimated_bytes as f64 / 1_073_741_824.0,
                num_bars,
            )));
        }

        let mut flat_features = Vec::with_capacity(num_bars * feature_dim);
        for (features, _) in data {
            for &v in features.iter() {
                flat_features.push(v as f32);
            }
        }

        let mut flat_targets = Vec::with_capacity(num_bars * target_dim);
        for (_, targets) in data {
            for i in 0..target_dim {
                flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32);
            }
        }

        let features = stream.clone_htod(&flat_features)
            .map_err(|e| MLError::ModelError(format!("GPU feature upload failed: {e}")))?;

        let targets = stream.clone_htod(&flat_targets)
            .map_err(|e| MLError::ModelError(format!("GPU target upload failed: {e}")))?;

        Ok(Self {
            features,
            targets,
            ofi_features: None,
            num_bars,
            feature_dim,
            aligned_state_dim: None,
        })
    }

    /// Upload OFI features (8 dims per bar) from MBP-10 order book data.
    ///
    /// Must be called after `upload()` with a slice matching `num_bars` length.
    /// Features are: OFI, VPIN, Kyle's Lambda, trade imbalance, etc.
    pub fn upload_ofi(
        &mut self,
        ofi_data: &[[f64; 8]],
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        let n = ofi_data.len().min(self.num_bars);
        if n == 0 {
            return Ok(());
        }

        let mut flat = Vec::with_capacity(self.num_bars * 8);
        for i in 0..self.num_bars {
            if i < n {
                for &v in ofi_data[i].iter() {
                    flat.push(v as f32);
                }
            } else {
                // Zero-pad if OFI data is shorter than market data
                flat.extend_from_slice(&[0.0_f32; 8]);
            }
        }

        let ofi_buf = stream.clone_htod(&flat)
            .map_err(|e| MLError::ModelError(format!("GPU OFI upload failed: {e}")))?;

        self.ofi_features = Some(ofi_buf);
        Ok(())
    }

    /// Set the aligned state dimension for tensor core padding.
    ///
    /// When set, `build_state_tensor` and `build_batch_states` will zero-pad
    /// the output to this width (must be >= raw state dim).
    pub fn set_aligned_state_dim(&mut self, dim: usize) {
        self.aligned_state_dim = Some(dim);
    }

    /// Estimated VRAM usage in bytes.
    pub fn vram_bytes(&self) -> usize {
        let ofi_dim = if self.ofi_features.is_some() { 8 } else { 0 };
        estimate_vram_bytes(self.num_bars * (self.feature_dim + 4 + ofi_dim))
    }

    /// D2D subrange copy: extract `count` f32 elements at `offset` into a new `CudaSlice`.
    pub(crate) fn d2d_subrange(
        src: &CudaSlice<f32>,
        offset: usize,
        count: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        let mut dst = stream.alloc_zeros::<f32>(count).map_err(|e| {
            MLError::ModelError(format!("d2d_subrange alloc: {e}"))
        })?;
        {
            let src_view = src.slice(offset..offset + count);
            let (src_ptr, _sg) = src_view.device_ptr(stream);
            let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
            let nbytes = count * std::mem::size_of::<f32>();
            #[allow(unsafe_code)]
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("d2d_subrange copy: {e}")))?;
            }
        }
        Ok(dst)
    }

    /// Get market features for a single bar as a [42] `CudaSlice` (D2D copy from GPU buffer).
    pub fn bar_features(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
        if bar_idx >= self.num_bars {
            return Err(MLError::ModelError(format!("Bar {bar_idx} out of range (num_bars={})", self.num_bars)));
        }
        Self::d2d_subrange(&self.features, bar_idx * self.feature_dim, self.feature_dim, stream)
    }

    /// Get market features for a batch of bars as a [batch_size * 42] `CudaSlice`.
    pub fn batch_features(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
        let count = count.min(self.num_bars.saturating_sub(start));
        Self::d2d_subrange(&self.features, start * self.feature_dim, count * self.feature_dim, stream)
    }

    /// Get target prices for a single bar as a [4] `CudaSlice`.
    pub fn bar_targets(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
        if bar_idx >= self.num_bars {
            return Err(MLError::ModelError(format!("Bar {bar_idx} target out of range (num_bars={})", self.num_bars)));
        }
        Self::d2d_subrange(&self.targets, bar_idx * 4, 4, stream)
    }

    /// Get target prices for a batch as [batch_size * 4] `CudaSlice`.
    pub fn batch_targets(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
        let count = count.min(self.num_bars.saturating_sub(start));
        Self::d2d_subrange(&self.targets, start * 4, count * 4, stream)
    }

    /// Get target values for a bar as a GPU-resident [4] `CudaSlice<f32>`.
    ///
    /// Tensor order: [current_close_preproc, next_close_preproc, current_close_raw, next_close_raw].
    /// Stays on GPU. Callers that need CPU scalars should extract at their own boundary.
    pub fn bar_target_values(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
        self.bar_targets(bar_idx, stream)
    }

    /// Build a complete state `CudaSlice` by concatenating market, portfolio, and OFI features.
    ///
    /// Returns a contiguous `[state_dim]` buffer, zero-padded to `aligned_state_dim` if set.
    /// All operations are GPU DtoD -- zero CPU download.
    pub fn build_state_tensor(
        &self,
        bar_idx: usize,
        portfolio_features: &[f32; 3],
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        let ofi_dim = if self.ofi_features.is_some() { 8 } else { 0 };
        let raw_dim = self.feature_dim + 3 + ofi_dim;
        let final_dim = self.aligned_state_dim.unwrap_or(raw_dim);

        // Allocate output buffer (zeros for padding region)
        let mut dst = stream.alloc_zeros::<f32>(final_dim)
            .map_err(|e| MLError::ModelError(format!("build_state alloc: {e}")))?;

        // DtoD copy: market features -> dst[0..feature_dim]
        let market_gpu = self.bar_features(bar_idx, stream)?;
        Self::dtod_copy_into(&market_gpu, &mut dst, 0, self.feature_dim, stream)?;

        // HtoD copy: portfolio features -> dst[feature_dim..feature_dim+3] (3 scalars, tiny)
        Self::htod_copy_into(portfolio_features, &mut dst, self.feature_dim, stream)?;

        // DtoD copy: OFI features -> dst[feature_dim+3..feature_dim+11]
        if let Some(ref ofi) = self.ofi_features {
            let ofi_gpu = Self::d2d_subrange(ofi, bar_idx * 8, 8, stream)?;
            Self::dtod_copy_into(&ofi_gpu, &mut dst, self.feature_dim + 3, 8, stream)?;
        }

        Ok(dst)
    }

    /// Build a batch of state `CudaSlice`s (row-interleaved) on GPU.
    ///
    /// Returns `[count * state_dim]` contiguous, zero-padded to `aligned_state_dim` per row.
    /// Market and OFI features use GPU DtoD copies. Only the 3 portfolio scalars come from host.
    pub fn build_batch_states(
        &self,
        start: usize,
        count: usize,
        portfolio_features: &[f32; 3],
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        let count = count.min(self.num_bars.saturating_sub(start));
        if count == 0 {
            return Err(MLError::ModelError("Empty batch for state construction".to_owned()));
        }
        let ofi_dim = if self.ofi_features.is_some() { 8 } else { 0 };
        let raw_dim = self.feature_dim + 3 + ofi_dim;
        let final_dim = self.aligned_state_dim.unwrap_or(raw_dim);

        // Allocate output buffer (zeros for padding region)
        let mut dst = stream.alloc_zeros::<f32>(count * final_dim)
            .map_err(|e| MLError::ModelError(format!("build_batch alloc: {e}")))?;

        // Upload portfolio features once (3 scalars)
        let port_gpu = stream.clone_htod(portfolio_features)
            .map_err(|e| MLError::ModelError(format!("portfolio HtoD: {e}")))?;

        // Get contiguous market features [count * feature_dim]
        let market_gpu = self.batch_features(start, count, stream)?;

        // Get contiguous OFI features [count * 8] if present
        let ofi_gpu = if let Some(ref ofi) = self.ofi_features {
            Some(Self::d2d_subrange(ofi, start * 8, count * 8, stream)?)
        } else {
            None
        };

        // For each row, DtoD copy: market, portfolio, OFI into the interleaved output
        for i in 0..count {
            let row_offset = i * final_dim;

            // Market features: contiguous in market_gpu
            Self::dtod_copy_into_at_offset(
                &market_gpu, i * self.feature_dim,
                &mut dst, row_offset,
                self.feature_dim, stream,
            )?;

            // Portfolio features: same 3 values for every row
            Self::dtod_copy_into(&port_gpu, &mut dst, row_offset + self.feature_dim, 3, stream)?;

            // OFI features: contiguous in ofi_gpu
            if let Some(ref ofi) = ofi_gpu {
                Self::dtod_copy_into_at_offset(
                    ofi, i * 8,
                    &mut dst, row_offset + self.feature_dim + 3,
                    8, stream,
                )?;
            }
        }

        Ok(dst)
    }

    /// DtoD copy `count` f32 elements from `src` (offset 0) into `dst` at element offset.
    #[allow(unsafe_code)]
    fn dtod_copy_into(
        src: &CudaSlice<f32>,
        dst: &mut CudaSlice<f32>,
        dst_offset_elems: usize,
        count: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        Self::dtod_copy_into_at_offset(src, 0, dst, dst_offset_elems, count, stream)
    }

    /// DtoD copy `count` f32 elements from `src[src_offset..]` into `dst[dst_offset..]`.
    ///
    /// Uses raw pointer arithmetic to write at an offset within the destination
    /// buffer. The source is sliced via `CudaSlice::slice()` for safety.
    #[allow(unsafe_code)]
    fn dtod_copy_into_at_offset(
        src: &CudaSlice<f32>,
        src_offset_elems: usize,
        dst: &mut CudaSlice<f32>,
        dst_offset_elems: usize,
        count: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        if count == 0 {
            return Ok(());
        }
        let nbytes = count * std::mem::size_of::<f32>();
        let src_view = src.slice(src_offset_elems..src_offset_elems + count);
        let (src_ptr, _sg) = src_view.device_ptr(stream);
        // Get base pointer of dst, then offset by dst_offset_elems * sizeof(f32)
        let (dst_base, _dg) = dst.device_ptr_mut(stream);
        let dst_ptr = dst_base + (dst_offset_elems * std::mem::size_of::<f32>()) as u64;
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, nbytes, stream.cu_stream(),
            ).map_err(|e| MLError::ModelError(format!("dtod_copy_into: {e}")))?;
        }
        Ok(())
    }

    /// Upload a small host slice into a subregion of a GPU buffer.
    ///
    /// Uploads `src` to a temporary GPU buffer, then DtoD copies into the
    /// destination at the given offset. Avoids needing `slice_mut` on `CudaSlice`.
    #[allow(unsafe_code)]
    fn htod_copy_into(
        src: &[f32],
        dst: &mut CudaSlice<f32>,
        dst_offset_elems: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        let tmp = stream.clone_htod(src)
            .map_err(|e| MLError::ModelError(format!("htod_copy_into upload: {e}")))?;
        Self::dtod_copy_into(&tmp, dst, dst_offset_elems, src.len(), stream)
    }
}

/// Reusable CPU staging buffers for GPU upload.
///
/// Pre-allocates flat `Vec<f32>` staging arrays sized for the maximum
/// expected fold. Reusing these across walk-forward folds avoids
/// repeated heap allocation/deallocation, reducing GC pauses and
/// memory fragmentation — especially important on GPUs with limited
/// VRAM (L4 24 GB) or large datasets.
///
/// # Usage
/// ```ignore
/// let mut pool = GpuBufferPool::new(300_000, 51, 4);
/// for fold in folds {
///     let gpu_data = pool.upload_dqn(&fold_data, &device)?;
///     // ... train on gpu_data ...
/// }
/// ```
#[derive(Debug)]
pub struct GpuBufferPool {
    /// Staging buffer for features [max_bars * feature_dim]
    feature_buf: Vec<f32>,
    /// Staging buffer for targets [max_bars * target_dim]
    target_buf: Vec<f32>,
    /// Maximum bars this pool can handle
    pub max_bars: usize,
    /// Feature dimension (51 for DQN)
    pub feature_dim: usize,
    /// Target dimension (4 for DQN)
    pub target_dim: usize,
}

impl GpuBufferPool {
    /// Create a pool pre-allocated for `max_bars` training bars.
    pub fn new(max_bars: usize, feature_dim: usize, target_dim: usize) -> Self {
        Self {
            feature_buf: vec![0.0_f32; max_bars * feature_dim],
            target_buf: vec![0.0_f32; max_bars * target_dim],
            max_bars,
            feature_dim,
            target_dim,
        }
    }

    /// Upload DQN data to GPU, reusing the pre-allocated staging buffers.
    ///
    /// Data is copied into the staging `Vec<f32>` (no heap alloc if within
    /// `max_bars`), then uploaded to GPU in a single `Tensor::from_vec` call.
    pub fn upload_dqn(
        &mut self,
        data: &[([f64; 42], Vec<f64>)],
        stream: &Arc<CudaStream>,
    ) -> Result<DqnGpuData, MLError> {
        let num_bars = data.len();
        if num_bars == 0 {
            return Err(MLError::ModelError("Empty training data".to_owned()));
        }

        let feat_len = num_bars * self.feature_dim;
        let targ_len = num_bars * self.target_dim;

        // Grow staging buffers only if the fold exceeds the pre-allocated size.
        if feat_len > self.feature_buf.len() {
            self.feature_buf.resize(feat_len, 0.0);
        }
        if targ_len > self.target_buf.len() {
            self.target_buf.resize(targ_len, 0.0);
        }

        // Fill staging buffers (zero-alloc copy).
        for (i, (features, targets)) in data.iter().enumerate() {
            let f_start = i * self.feature_dim;
            for (j, &v) in features.iter().enumerate() {
                self.feature_buf[f_start + j] = v as f32;
            }

            let t_start = i * self.target_dim;
            for j in 0..self.target_dim {
                self.target_buf[t_start + j] = targets.get(j).copied().unwrap_or(0.0) as f32;
            }
        }

        // Upload the used slice to GPU via CudaStream::clone_htod (f32, no BF16 cast).
        let features = stream.clone_htod(&self.feature_buf[..feat_len])
            .map_err(|e| MLError::ModelError(format!("GPU feature upload failed: {e}")))?;

        let targets = stream.clone_htod(&self.target_buf[..targ_len])
            .map_err(|e| MLError::ModelError(format!("GPU target upload failed: {e}")))?;

        Ok(DqnGpuData {
            features,
            targets,
            ofi_features: None,
            num_bars,
            feature_dim: self.feature_dim,
            aligned_state_dim: None,
        })
    }

    /// Total bytes of pre-allocated CPU staging memory.
    pub fn staging_bytes(&self) -> usize {
        (self.feature_buf.len() + self.target_buf.len()) * std::mem::size_of::<f32>()
    }
}

/// Pre-uploaded GPU training data for PPO trainer.
///
/// Holds all rollout states as a single [N * state_dim] `CudaSlice<f32>`.
pub struct PpoGpuData {
    /// All states [num_steps * state_dim] on GPU (f32, row-major)
    pub states: CudaSlice<f32>,
    /// Number of training steps
    pub num_steps: usize,
    /// State dimension
    pub state_dim: usize,
}

impl std::fmt::Debug for PpoGpuData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PpoGpuData")
            .field("num_steps", &self.num_steps)
            .field("state_dim", &self.state_dim)
            .finish()
    }
}

impl PpoGpuData {
    /// Upload PPO market data to GPU as a single contiguous `CudaSlice<f32>`.
    pub fn upload(market_data: &[Vec<f32>], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let num_steps = market_data.len();
        if num_steps == 0 {
            return Err(MLError::ModelError("Empty market data".to_owned()));
        }

        let state_dim = market_data[0].len();

        let estimated_bytes = estimate_vram_bytes(num_steps * state_dim);
        if estimated_bytes > MAX_UPLOAD_BYTES {
            return Err(MLError::ModelError(format!(
                "Market data too large for GPU upload: {:.1} GB > 2.0 GB limit ({} steps × {} dims)",
                estimated_bytes as f64 / 1_073_741_824.0,
                num_steps,
                state_dim,
            )));
        }

        let mut flat_states = Vec::with_capacity(num_steps * state_dim);
        for state in market_data {
            if state.len() != state_dim {
                return Err(MLError::ModelError(format!(
                    "Inconsistent state dim: expected {state_dim}, got {}",
                    state.len()
                )));
            }
            flat_states.extend_from_slice(state);
        }

        let states = stream.clone_htod(&flat_states)
            .map_err(|e| MLError::ModelError(format!("GPU state upload failed: {e}")))?;

        Ok(Self {
            states,
            num_steps,
            state_dim,
        })
    }

    /// Estimated VRAM usage in bytes.
    pub fn vram_bytes(&self) -> usize {
        estimate_vram_bytes(self.num_steps * self.state_dim)
    }

    /// Get a single step's state as a [state_dim] `CudaSlice` (D2D copy from GPU buffer).
    pub fn step_state(&self, step_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
        if step_idx >= self.num_steps {
            return Err(MLError::ModelError(format!("Step {step_idx} out of range (num_steps={})", self.num_steps)));
        }
        DqnGpuData::d2d_subrange(&self.states, step_idx * self.state_dim, self.state_dim, stream)
    }

    /// Get a batch of states as [batch_size * state_dim] `CudaSlice`.
    pub fn batch_states(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
        let count = count.min(self.num_steps.saturating_sub(start));
        DqnGpuData::d2d_subrange(&self.states, start * self.state_dim, count * self.state_dim, stream)
    }
}

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

    /// Shared CUDA stream for all cuda_pipeline tests.
    /// Reuses one context/stream instead of creating a new one per test,
    /// preventing handle pool exhaustion across test binaries.
    static SHARED_STREAM: OnceLock<Arc<CudaStream>> = OnceLock::new();

    fn cuda_stream() -> Arc<CudaStream> {
        Arc::clone(SHARED_STREAM.get_or_init(|| {
            let ctx = cudarc::driver::CudaContext::new(0)
                .expect("CUDA context required for cuda_pipeline tests");
            ctx.new_stream()
                .expect("CUDA stream required for cuda_pipeline tests")
        }))
    }

    #[test]
    fn test_dqn_gpu_data_upload() {
        let stream = cuda_stream();
        let data: Vec<([f64; 42], Vec<f64>)> = (0..100)
            .map(|i| {
                let features = [i as f64 * 0.01; 42];
                let targets = vec![100.0 + i as f64, 101.0 + i as f64, 100.5, 101.5];
                (features, targets)
            })
            .collect();

        let gpu_data = DqnGpuData::upload(&data, &stream).expect("upload");
        assert_eq!(gpu_data.num_bars, 100);
        assert_eq!(gpu_data.feature_dim, 42);

        let bar0 = gpu_data.bar_features(0, &stream).expect("bar0");
        assert_eq!(bar0.len(), 42);

        let batch = gpu_data.batch_features(0, 10, &stream).expect("batch");
        assert_eq!(batch.len(), 10 * 42);

        // Download target values and verify on CPU
        let targets_gpu = gpu_data.bar_target_values(0, &stream).expect("targets");
        let mut targets_host = vec![0.0_f32; 4];
        stream.memcpy_dtoh(&targets_gpu, &mut targets_host).expect("DtoH"); // test readback
        assert!((targets_host.first().copied().unwrap_or(0.0) - 100.0).abs() < 0.01);
    }

    #[test]
    fn test_dqn_build_state_tensor() {
        let stream = cuda_stream();
        let data: Vec<([f64; 42], Vec<f64>)> = vec![
            ([1.0; 42], vec![100.0, 101.0, 100.5, 101.5]),
        ];

        let gpu_data = DqnGpuData::upload(&data, &stream).expect("upload");
        let portfolio = [0.95_f32, 0.5, 0.0001];
        let state = gpu_data.build_state_tensor(0, &portfolio, &stream).expect("state");
        assert_eq!(state.len(), 45); // 42 market + 3 portfolio
    }

    #[test]
    fn test_tensor_core_alignment() {
        assert_eq!(align_to_tensor_cores(54), 56);
        assert_eq!(align_to_tensor_cores(64), 64); // already aligned
        assert_eq!(align_to_tensor_cores(51), 56); // pad up
        assert_eq!(align_to_tensor_cores(1), 8);
        assert_eq!(align_to_tensor_cores(8), 8);
        assert_eq!(align_to_tensor_cores(0), 0);
    }

    #[test]
    fn test_optimal_launch_dims() {
        // Small: 16 episodes → block=32, grid=1
        let (g, b) = optimal_launch_dims(16, 0);
        assert_eq!(b, 32);
        assert_eq!(g, 1);

        // Medium: 128 episodes → block=128, grid=1
        let (g, b) = optimal_launch_dims(128, 0);
        assert_eq!(b, 128);
        assert_eq!(g, 1);

        // Large: 256 episodes → block=256, grid=1
        let (g, b) = optimal_launch_dims(256, 0);
        assert_eq!(b, 256);
        assert_eq!(g, 1);

        // Very large: 512 episodes → block=256, grid=2
        let (g, b) = optimal_launch_dims(512, 0);
        assert_eq!(b, 256);
        assert_eq!(g, 2);

        // Constrained GPU: max_tpb=64
        let (g, b) = optimal_launch_dims(256, 64);
        assert_eq!(b, 64);
        assert_eq!(g, 4);

        // Edge: 1 episode
        let (g, b) = optimal_launch_dims(1, 0);
        assert_eq!(b, 32);
        assert_eq!(g, 1);
    }

    #[test]
    fn test_ppo_gpu_data_upload() {
        let stream = cuda_stream();
        let market_data: Vec<Vec<f32>> = (0..200)
            .map(|i| vec![i as f32 * 0.01; 45])
            .collect();

        let gpu_data = PpoGpuData::upload(&market_data, &stream).expect("upload");
        assert_eq!(gpu_data.num_steps, 200);
        // PpoGpuData::upload uses raw state_dim (no padding)
        assert_eq!(gpu_data.state_dim, 45);

        let step0 = gpu_data.step_state(0, &stream).expect("step0");
        assert_eq!(step0.len(), 45);

        let batch = gpu_data.batch_states(10, 20, &stream).expect("batch");
        assert_eq!(batch.len(), 20 * 45);
    }

    #[test]
    fn test_ppo_gpu_data_inconsistent_dims() {
        let stream = cuda_stream();
        let market_data = vec![
            vec![1.0_f32; 45],
            vec![2.0_f32; 42],
        ];
        let result = PpoGpuData::upload(&market_data, &stream);
        assert!(result.is_err());
    }

    #[test]
    fn test_dqn_gpu_data_empty() {
        let stream = cuda_stream();
        let data: Vec<([f64; 42], Vec<f64>)> = vec![];
        let result = DqnGpuData::upload(&data, &stream);
        assert!(result.is_err());
    }

    #[test]
    fn test_ppo_gpu_data_empty() {
        let stream = cuda_stream();
        let data: Vec<Vec<f32>> = vec![];
        let result = PpoGpuData::upload(&data, &stream);
        assert!(result.is_err());
    }

    #[test]
    fn test_vram_estimation() {
        // 300K bars × 55 features × 4 bytes = ~66 MB
        let bytes = estimate_vram_bytes(300_000 * 55);
        assert!(bytes < 100_000_000); // Under 100 MB
        assert!(bytes > 60_000_000);  // Over 60 MB
    }

    #[test]
    fn test_dqn_boundary_access() {
        let stream = cuda_stream();
        let data: Vec<([f64; 42], Vec<f64>)> = vec![
            ([0.5; 42], vec![100.0, 101.0, 100.5, 101.5]),
        ];
        let gpu_data = DqnGpuData::upload(&data, &stream).expect("upload");

        assert!(gpu_data.bar_features(0, &stream).is_ok());
        assert!(gpu_data.bar_features(1, &stream).is_err());

        let batch = gpu_data.batch_features(0, 100, &stream).expect("batch");
        assert_eq!(batch.len(), 42); // clamped to 1 bar
    }

    #[test]
    fn test_dqn_build_batch_states() {
        let stream = cuda_stream();
        let data: Vec<([f64; 42], Vec<f64>)> = (0..100)
            .map(|i| {
                let features = [i as f64 * 0.01; 42];
                let targets = vec![100.0, 101.0, 100.5, 101.5];
                (features, targets)
            })
            .collect();

        let gpu_data = DqnGpuData::upload(&data, &stream).expect("upload");
        let portfolio = [0.95_f32, 0.5, 0.0001];

        // Build batch of 32 states starting at bar 10
        let batch = gpu_data.build_batch_states(10, 32, &portfolio, &stream).expect("batch");
        assert_eq!(batch.len(), 32 * 45); // 42 market + 3 portfolio

        // Download and verify portfolio features in row 0
        let mut host = vec![0.0_f32; 32 * 45];
        stream.memcpy_dtoh(&batch, &mut host).expect("DtoH"); // test readback

        // Row 0: features[42..45] should be portfolio [0.95, 0.5, 0.0001]
        let pf = &host[42..45];
        assert!((pf[0] - 0.95).abs() < 1e-4, "portfolio[0]={}", pf[0]);
        assert!((pf[1] - 0.5).abs() < 1e-4, "portfolio[1]={}", pf[1]);

        // Row 31: same portfolio features
        let row31_start = 31 * 45;
        let pf31 = &host[row31_start + 42..row31_start + 45];
        assert!((pf31[0] - 0.95).abs() < 1e-4);

        // Market features should differ between rows 0 and 31
        assert!((host[0] - host[row31_start]).abs() > 1e-6, "market features should differ");
    }

    #[test]
    fn test_dqn_build_batch_states_clamping() {
        let stream = cuda_stream();
        let data: Vec<([f64; 42], Vec<f64>)> = vec![
            ([1.0; 42], vec![100.0, 101.0, 100.5, 101.5]),
            ([2.0; 42], vec![100.0, 101.0, 100.5, 101.5]),
        ];
        let gpu_data = DqnGpuData::upload(&data, &stream).expect("upload");
        let portfolio = [1.0_f32, 0.0, 0.0001];

        // Request more bars than available -> clamped to 2
        let batch = gpu_data.build_batch_states(0, 100, &portfolio, &stream).expect("batch");
        assert_eq!(batch.len(), 2 * 45);

        // Start beyond range -> error
        let err = gpu_data.build_batch_states(5, 10, &portfolio, &stream);
        assert!(err.is_err());
    }

    #[test]
    fn test_gpu_buffer_pool_basic() {
        let stream = cuda_stream();
        let mut pool = GpuBufferPool::new(1000, 42, 4);
        assert_eq!(pool.max_bars, 1000);
        assert_eq!(pool.staging_bytes(), (1000 * 42 + 1000 * 4) * 4);

        let data: Vec<([f64; 42], Vec<f64>)> = (0..50)
            .map(|i| ([i as f64 * 0.01; 42], vec![100.0, 101.0, 100.5, 101.5]))
            .collect();
        let gpu_data = pool.upload_dqn(&data, &stream).expect("upload");
        assert_eq!(gpu_data.num_bars, 50);
        assert_eq!(gpu_data.feature_dim, 42);
    }

    #[test]
    fn test_gpu_buffer_pool_reuse() {
        let stream = cuda_stream();
        let mut pool = GpuBufferPool::new(500, 42, 4);

        let data1: Vec<([f64; 42], Vec<f64>)> = (0..100)
            .map(|i| ([i as f64 * 0.01; 42], vec![100.0, 101.0, 100.5, 101.5]))
            .collect();
        let g1 = pool.upload_dqn(&data1, &stream).expect("upload1");
        assert_eq!(g1.num_bars, 100);

        let data2: Vec<([f64; 42], Vec<f64>)> = (0..200)
            .map(|i| ([i as f64 * 0.02; 42], vec![200.0, 201.0, 200.5, 201.5]))
            .collect();
        let g2 = pool.upload_dqn(&data2, &stream).expect("upload2");
        assert_eq!(g2.num_bars, 200);

        // Download target values and verify
        let t = g2.bar_target_values(0, &stream).expect("targets");
        let mut host = vec![0.0_f32; 4];
        stream.memcpy_dtoh(&t, &mut host).expect("DtoH"); // test readback
        assert!((host[0] - 200.0).abs() < 0.01);
    }

    #[test]
    fn test_gpu_buffer_pool_grow() {
        let stream = cuda_stream();
        let mut pool = GpuBufferPool::new(10, 42, 4);

        let data: Vec<([f64; 42], Vec<f64>)> = (0..50)
            .map(|i| ([i as f64; 42], vec![1.0, 2.0, 3.0, 4.0]))
            .collect();
        let gpu_data = pool.upload_dqn(&data, &stream).expect("upload");
        assert_eq!(gpu_data.num_bars, 50);
        assert!(pool.feature_buf.len() >= 50 * 42);
    }

    #[test]
    fn test_gpu_buffer_pool_empty() {
        let stream = cuda_stream();
        let mut pool = GpuBufferPool::new(100, 42, 4);
        let data: Vec<([f64; 42], Vec<f64>)> = vec![];
        assert!(pool.upload_dqn(&data, &stream).is_err());
    }

    #[test]
    fn test_kernel_source_contains_entry_point() {
        let src = include_str!("experience_kernels.cu");
        assert!(
            src.contains("experience_state_gather"),
            "Kernel source must contain state_gather entry point"
        );
        assert!(
            src.contains("experience_action_select"),
            "Kernel source must contain action_select entry point"
        );
        assert!(
            src.contains("experience_env_step"),
            "Kernel source must contain env_step entry point"
        );
        assert!(
            src.contains("extern \"C\""),
            "Kernel source must declare extern C linkage"
        );
    }

    #[test]
    fn test_common_header_contains_shared_functions() {
        let src = include_str!("common_device_functions.cuh");
        assert!(src.contains("gpu_random"), "Missing gpu_random");
        assert!(src.contains("matvec_leaky_relu"), "Missing matvec_leaky_relu");
        assert!(src.contains("action_to_exposure"), "Missing action_to_exposure");
        assert!(src.contains("barrier_init"), "Missing barrier_init");
        assert!(src.contains("diversity_entropy"), "Missing diversity_entropy");
        assert!(src.contains("curiosity_inference"), "Missing curiosity_inference");
    }

    #[test]
    fn test_experience_config_defaults() {
        use super::gpu_experience_collector::ExperienceCollectorConfig;

        let cfg = ExperienceCollectorConfig::default();
        assert_eq!(cfg.n_episodes, 2048);
        assert_eq!(cfg.timesteps_per_episode, 500);
        assert_eq!(cfg.total_experiences(), 2048 * 500);
        assert!((cfg.epsilon - 0.1).abs() < f32::EPSILON);
        assert!((cfg.gamma - 0.99).abs() < f32::EPSILON);
        assert!((cfg.barrier_profit_mult - 1.02).abs() < f32::EPSILON);
        assert!((cfg.barrier_loss_mult - 0.98).abs() < f32::EPSILON);
    }

    #[test]
    fn test_ppo_kernel_source_contains_all_functions() {
        let src = include_str!("ppo_experience_kernel.cu");
        assert!(src.contains("ppo_actor_forward"), "Missing ppo_actor_forward");
        assert!(src.contains("softmax_sample"), "Missing softmax_sample");
        assert!(src.contains("ppo_critic_forward"), "Missing ppo_critic_forward");
        assert!(src.contains("compute_gae_backward"), "Missing compute_gae_backward");
        assert!(src.contains("ppo_full_experience_kernel"), "Missing kernel entry point");
        assert!(src.contains("extern \"C\""), "Missing extern C linkage");
    }

    #[test]
    fn test_ppo_kernel_source_concatenation() {
        let common = include_str!("common_device_functions.cuh");
        let kernel = include_str!("ppo_experience_kernel.cu");
        let full = format!("{}\n{}", common, kernel);
        assert!(full.contains("gpu_random"));
        assert!(full.contains("matvec_leaky_relu"));
        assert!(full.contains("ppo_full_experience_kernel"));
        assert!(full.contains("compute_gae_backward"));
        assert!(full.contains("softmax_sample"));
    }

    #[test]
    fn test_ppo_collector_config_defaults() {
        use super::gpu_ppo_collector::PpoCollectorConfig;
        let cfg = PpoCollectorConfig::default();
        assert_eq!(cfg.n_episodes, 128);
        assert_eq!(cfg.timesteps_per_episode, 500);
        assert_eq!(cfg.total_experiences(), 64_000);
        assert!((cfg.gamma - 0.99).abs() < f32::EPSILON);
        assert!((cfg.gae_lambda - 0.95).abs() < f32::EPSILON);
    }

    // ── Phase 3: Batch struct layout & index-math validation ──────────

    #[test]
    fn test_experience_batch_index_math() {
        // Pure index-math validation — no GPU types needed
        let n_episodes = 4_usize;
        let timesteps = 10_usize;
        let state_dim = 72_usize;
        let total = n_episodes * timesteps;

        let states: Vec<f32> = (0..total * state_dim).map(|i| i as f32 * 0.001).collect();
        let actions: Vec<i32> = (0..total).map(|i| (i % 5) as i32).collect();
        let done_flags: Vec<i32> = (0..total)
            .map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 })
            .collect();

        assert_eq!(states.len(), total * state_dim);
        assert_eq!(actions.len(), total);

        for ep in 0..n_episodes {
            for t in 0..timesteps {
                let idx = ep * timesteps + t;
                assert!(idx < total);
                let state_end = idx * state_dim + state_dim;
                assert!(state_end <= states.len());
                let a = actions.get(idx).copied().unwrap_or(-1);
                assert!((0..5).contains(&a));
            }
            let last_idx = ep * timesteps + timesteps - 1;
            assert_eq!(done_flags.get(last_idx).copied().unwrap_or(0), 1);
        }
    }

    #[test]
    fn test_ppo_experience_batch_index_math() {
        use super::gpu_ppo_collector::PpoExperienceBatch;

        let n_episodes = 4;
        let timesteps = 10;
        let state_dim = 72;
        let total = n_episodes * timesteps;

        // PpoExperienceBatch is now GPU-resident. Upload test data to GPU.
        let ctx = cudarc::driver::CudaContext::new(0).expect("CUDA required");
        let stream = ctx.default_stream();

        let states_host: Vec<f32> = (0..total * state_dim).map(|i| i as f32 * 0.001).collect();
        let advantages_host: Vec<f32> = (0..total).map(|i| i as f32 * 0.1 - 2.0).collect();
        let returns_host: Vec<f32> = (0..total).map(|i| i as f32 * 0.5).collect();

        let batch = PpoExperienceBatch {
            states: stream.clone_htod(&states_host).unwrap(),
            actions: stream.clone_htod(&(0..total).map(|i| (i % 45) as i32).collect::<Vec<_>>()).unwrap(),
            log_probs: stream.clone_htod(&vec![-1.5_f32; total]).unwrap(),
            advantages: stream.clone_htod(&advantages_host).unwrap(),
            returns: stream.clone_htod(&returns_host).unwrap(),
            done_flags: stream.clone_htod(&(0..total).map(|i| if i % timesteps == timesteps - 1 { 1_i32 } else { 0 }).collect::<Vec<_>>()).unwrap(),
            n_episodes,
            timesteps,
            state_dim,
            stream: stream.clone(),
        };

        let states = batch.download_states().unwrap();
        let returns = batch.download_returns().unwrap();
        let advantages = batch.download_advantages().unwrap();

        assert_eq!(states.len(), total * state_dim);
        assert_eq!(batch.total(), total);

        // Verify V(s) = R(s) - A(s) reconstruction
        for i in 0..total {
            let r = returns.get(i).copied().unwrap_or(0.0);
            let a = advantages.get(i).copied().unwrap_or(0.0);
            assert!((r - a).is_finite(), "V(s) should be finite at idx {}", i);
        }

        // Verify state chunking
        let chunks: Vec<&[f32]> = states.chunks(state_dim).collect();
        assert_eq!(chunks.len(), total);
    }

    #[test]
    fn test_experience_batch_next_state_boundary() {
        // Pure index-math validation for next_state shift logic
        let n_episodes = 2_usize;
        let timesteps = 5_usize;
        let state_dim = 72_usize;
        let total = n_episodes * timesteps;

        let states: Vec<f32> = (0..total)
            .flat_map(|idx| std::iter::repeat(idx as f32).take(state_dim))
            .collect();
        let done_flags: Vec<i32> = (0..total)
            .map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 })
            .collect();

        // Verify next_state shift: state[t+1] within episode, self at terminal
        for ep in 0..n_episodes {
            for t in 0..timesteps {
                let idx = ep * timesteps + t;
                let done = done_flags.get(idx).copied().unwrap_or(0) != 0;
                if !done && t + 1 < timesteps {
                    let next_idx = ep * timesteps + t + 1;
                    let s_val = states.get(idx * state_dim).copied().unwrap_or(0.0);
                    let ns_val = states.get(next_idx * state_dim).copied().unwrap_or(0.0);
                    assert!(
                        (ns_val - s_val - 1.0).abs() < f32::EPSILON,
                        "next_state should be current + 1 step"
                    );
                }
            }
        }
    }

    #[test]
    fn test_ppo_batch_gae_consistency() {
        use super::gpu_ppo_collector::PpoExperienceBatch;

        let n_episodes = 2;
        let timesteps = 10;
        let total = n_episodes * timesteps;
        let state_dim = 72;

        let ctx = cudarc::driver::CudaContext::new(0).expect("CUDA required");
        let stream = ctx.default_stream();

        let mut advantages_host = Vec::with_capacity(total);
        for _ep in 0..n_episodes {
            for t in 0..timesteps {
                let a = (timesteps as f32 / 2.0 - t as f32) * 0.1;
                advantages_host.push(a);
            }
        }
        let returns_host: Vec<f32> = advantages_host.iter().map(|&a| a + 5.0).collect();

        let batch = PpoExperienceBatch {
            states: stream.clone_htod(&vec![0.0_f32; total * state_dim]).unwrap(),
            actions: stream.clone_htod(&vec![1_i32; total]).unwrap(),
            log_probs: stream.clone_htod(&vec![-1.0_f32; total]).unwrap(),
            advantages: stream.clone_htod(&advantages_host).unwrap(),
            returns: stream.clone_htod(&returns_host).unwrap(),
            done_flags: stream.clone_htod(&(0..total).map(|i| if i % timesteps == timesteps - 1 { 1_i32 } else { 0 }).collect::<Vec<_>>()).unwrap(),
            n_episodes,
            timesteps,
            state_dim,
            stream: stream.clone(),
        };

        let returns = batch.download_returns().unwrap();
        let advantages = batch.download_advantages().unwrap();

        for i in 0..total {
            let r = returns.get(i).copied().unwrap_or(0.0);
            let a = advantages.get(i).copied().unwrap_or(0.0);
            assert!(
                (r - a - 5.0).abs() < f32::EPSILON,
                "Value at {} should be ~5.0, got {}",
                i,
                r - a
            );
        }
    }
}
