#![allow(unsafe_code)] // Required for CUDA kernel launches

//! GPU-Resident Replay Buffer -- pure cudarc `CudaSlice` internals.
//!
//! PER sampling uses an O(log n) GPU segment tree instead of O(n) prefix sum.
//! The segment tree is a flat `CudaSlice<f32>` of size `2 * capacity_pow2`:
//! root at index 1, leaves at `capacity_pow2..2*capacity_pow2`. Internal
//! nodes store sums of children. Sampling is parallel root-to-leaf traversal
//! with Philox RNG (1 kernel launch), replacing the old 5-kernel pipeline
//! (pow_alpha + prefix_sum + threshold_gen + searchsorted + i64_cast).
//!
//! Output `GpuBatchSlices` wraps gathered data as raw `CudaSlice` buffers
//! for downstream neural network consumption.

use std::sync::{Arc, OnceLock};

use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use ml_core::nvtx::NvtxRange;
use ml_core::MLError;

// ---------------------------------------------------------------------------
// GPU batch output (CudaSlice-based, no Candle Tensor dependency)
// ---------------------------------------------------------------------------

/// Pre-built GPU batch for training. All fields are raw `CudaSlice` on GPU.
#[allow(missing_debug_implementations)]
pub struct GpuBatchSlices {
    pub states: CudaSlice<u16>,       // [batch_size * state_dim] bf16 on GPU
    pub next_states: CudaSlice<u16>,  // [batch_size * state_dim] bf16 on GPU
    pub actions: CudaSlice<u32>,      // [batch_size] u32 on GPU
    pub rewards: CudaSlice<f32>,      // [batch_size] f32 on GPU
    pub dones: CudaSlice<f32>,        // [batch_size] f32 on GPU (0.0/1.0)
    pub weights: CudaSlice<f32>,      // [batch_size] f32 on GPU (IS weights)
    pub indices: CudaSlice<u32>,      // [batch_size] u32 on GPU (buffer indices)
    /// Episode IDs for sampled transitions `[batch_size]` i32 on GPU.
    /// Used by HER Future/Final strategies for episode-aware donor selection.
    pub episode_ids: Option<CudaSlice<i32>>,
    pub batch_size: usize,
    pub state_dim: usize,
}

impl GpuBatchSlices {
    /// Convert to `GpuBatch` (GpuTensor-based) via GPU-only bf16->f32 cast kernels.
    ///
    /// Zero CPU downloads -- all type conversions happen on GPU via custom CUDA kernels.
    pub fn into_gpu_batch(self, stream: &Arc<CudaStream>) -> Result<crate::replay_buffer_type::GpuBatch, MLError> {
        use ml_core::cuda_autograd::GpuTensor;

        let bs = self.batch_size;
        let sd = self.state_dim;

        let kernels = get_cast_kernels(stream)?;

        // bf16 states -> f32 via GPU cast kernel (zero host roundtrip)
        let states = bf16_slice_to_gpu_tensor_gpu(&self.states, vec![bs, sd], stream, kernels)?;
        let next_states = bf16_slice_to_gpu_tensor_gpu(&self.next_states, vec![bs, sd], stream, kernels)?;

        // u32 actions/indices -> f32 via GPU cast kernel (zero host roundtrip)
        let actions = u32_slice_to_gpu_tensor_gpu(&self.actions, vec![bs], stream, kernels)?;
        let indices = u32_slice_to_gpu_tensor_gpu(&self.indices, vec![bs], stream, kernels)?;

        // f32 slices -> GpuTensor (trivial wrap)
        let rewards = GpuTensor::from_flat(self.rewards);
        let dones = GpuTensor::from_flat(self.dones);
        let weights = GpuTensor::from_flat(self.weights);

        Ok(crate::replay_buffer_type::GpuBatch {
            states,
            actions,
            rewards,
            next_states,
            dones,
            weights,
            indices,
            episode_ids: self.episode_ids,
        })
    }
}

// ---------------------------------------------------------------------------
// Cast kernel cache for bf16->f32 and u32->f32 GPU-only conversions
// ---------------------------------------------------------------------------

struct CastKernels {
    bf16_to_f32: CudaFunction,
    u32_to_f32: CudaFunction,
    f32_to_u32: CudaFunction,
}

static CAST_KERNELS: OnceLock<Result<CastKernels, String>> = OnceLock::new();

fn get_cast_kernels(stream: &Arc<CudaStream>) -> Result<&'static CastKernels, MLError> {
    let result = CAST_KERNELS.get_or_init(|| {
        let ctx = stream.context();
        let src = include_str!("replay_buffer_kernels.cu");
        let ptx = ml_core::cuda_compile::compile_ptx_for_device(src, ctx)
            .map_err(|e| format!("cast kernels compile: {e}"))?;
        let module = ctx.load_module(ptx)
            .map_err(|e| format!("cast module: {e}"))?;
        Ok(CastKernels {
            bf16_to_f32: module.load_function("bf16_to_f32_cast")
                .map_err(|e| format!("bf16_to_f32: {e}"))?,
            u32_to_f32: module.load_function("u32_to_f32_cast")
                .map_err(|e| format!("u32_to_f32: {e}"))?,
            f32_to_u32: module.load_function("f32_idx_to_u32")
                .map_err(|e| format!("f32_to_u32: {e}"))?,
        })
    });
    match result {
        Ok(k) => Ok(k),
        Err(e) => Err(MLError::ModelError(e.clone())),
    }
}

/// Public API for f32->u32 GPU cast (used by `replay_buffer_type` for index conversion).
pub struct F32ToU32Caster {
    kernel: &'static CudaFunction,
}

impl F32ToU32Caster {
    /// Cast f32-encoded indices to u32 on GPU (zero CPU download).
    pub fn cast_f32_to_u32(
        &self,
        src: &CudaSlice<f32>,
        n: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<u32>, MLError> {
        let out = stream.alloc_zeros::<u32>(n)
            .map_err(|e| MLError::ModelError(format!("f32->u32 alloc: {e}")))?;
        let ni = n as i32;
        // SAFETY: out, src are valid device allocations of at least n elements.
        unsafe {
            stream.launch_builder(self.kernel)
                .arg(&out).arg(src).arg(&ni)
                .launch(lcfg(n))
                .map_err(|e| MLError::ModelError(format!("f32_to_u32 kernel: {e}")))?;
        }
        Ok(out)
    }
}

/// Get the f32->u32 GPU caster (compiles kernel on first call, cached thereafter).
pub fn get_cast_kernels_f32_to_u32(stream: &Arc<CudaStream>) -> Result<F32ToU32Caster, MLError> {
    let kernels = get_cast_kernels(stream)?;
    Ok(F32ToU32Caster { kernel: &kernels.f32_to_u32 })
}

/// Convert a bf16 `CudaSlice<u16>` to f32 `GpuTensor` via GPU cast kernel (zero CPU download).
fn bf16_slice_to_gpu_tensor_gpu(
    src: &CudaSlice<u16>,
    shape: Vec<usize>,
    stream: &Arc<CudaStream>,
    kernels: &CastKernels,
) -> Result<ml_core::cuda_autograd::GpuTensor, MLError> {
    use ml_core::cuda_autograd::GpuTensor;

    let n = src.len();
    let out = stream.alloc_zeros::<f32>(n)
        .map_err(|e| MLError::ModelError(format!("bf16->f32 alloc: {e}")))?;
    let ni = n as i32;
    // SAFETY: out, src are valid device allocations of at least n elements.
    unsafe {
        stream.launch_builder(&kernels.bf16_to_f32)
            .arg(&out).arg(src).arg(&ni)
            .launch(lcfg(n))
            .map_err(|e| MLError::ModelError(format!("bf16_to_f32 kernel: {e}")))?;
    }
    GpuTensor::new(out, shape)
}

/// Convert a u32 `CudaSlice<u32>` to f32 `GpuTensor` via GPU cast kernel (zero CPU download).
fn u32_slice_to_gpu_tensor_gpu(
    src: &CudaSlice<u32>,
    shape: Vec<usize>,
    stream: &Arc<CudaStream>,
    kernels: &CastKernels,
) -> Result<ml_core::cuda_autograd::GpuTensor, MLError> {
    use ml_core::cuda_autograd::GpuTensor;

    let n = src.len();
    let out = stream.alloc_zeros::<f32>(n)
        .map_err(|e| MLError::ModelError(format!("u32->f32 alloc: {e}")))?;
    let ni = n as i32;
    // SAFETY: out, src are valid device allocations of at least n elements.
    unsafe {
        stream.launch_builder(&kernels.u32_to_f32)
            .arg(&out).arg(src).arg(&ni)
            .launch(lcfg(n))
            .map_err(|e| MLError::ModelError(format!("u32_to_f32 kernel: {e}")))?;
    }
    GpuTensor::new(out, shape)
}

// ---------------------------------------------------------------------------
// Compiled kernel cache
// ---------------------------------------------------------------------------

struct ReplayKernels {
    scatter_insert_f32: CudaFunction,
    scatter_insert_u32: CudaFunction,
    scatter_insert_bf16: CudaFunction,
    gather_f32: CudaFunction,
    gather_u32: CudaFunction,
    gather_bf16_rows: CudaFunction,
    is_weights_f32: CudaFunction,
    normalize_weights_f32: CudaFunction,
    fill_from_gpu_f32: CudaFunction,
    reduce_max_f32: CudaFunction,
    i64_to_u32: CudaFunction,
    f32_to_bf16_cast: CudaFunction,
    max_of_two_f32: CudaFunction,
    // Segment tree kernels (replace prefix_sum + searchsorted + pow_alpha pipeline)
    seg_tree_update: CudaFunction,
    seg_tree_insert: CudaFunction,
    seg_tree_sample: CudaFunction,
    seg_tree_gather_prios: CudaFunction,
}

impl ReplayKernels {
    fn compile(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let ctx = stream.context();
        let rb_src = include_str!("replay_buffer_kernels.cu");
        let rb_ptx = ml_core::cuda_compile::compile_ptx_for_device(rb_src, ctx)
            .map_err(|e| MLError::ModelError(format!("rb kernels compile: {e}")))?;
        let rb_mod = ctx.load_module(rb_ptx)
            .map_err(|e| MLError::ModelError(format!("rb module: {e}")))?;
        let ld = |n: &str| -> Result<CudaFunction, MLError> {
            rb_mod.load_function(n).map_err(|e| MLError::ModelError(format!("load {n}: {e}")))
        };
        // Segment tree kernels (replace prefix_sum + searchsorted + pow_alpha)
        let st_ptx = ml_core::cuda_compile::compile_ptx_for_device(
            include_str!("seg_tree_kernel.cu"), ctx,
        ).map_err(|e| MLError::ModelError(format!("seg_tree compile: {e}")))?;
        let st_mod = ctx.load_module(st_ptx)
            .map_err(|e| MLError::ModelError(format!("seg_tree module: {e}")))?;

        Ok(Self {
            scatter_insert_f32: ld("scatter_insert_f32")?,
            scatter_insert_u32: ld("scatter_insert_u32")?,
            scatter_insert_bf16: ld("scatter_insert_bf16")?,
            gather_f32: ld("gather_f32")?, gather_u32: ld("gather_u32")?,
            gather_bf16_rows: ld("gather_bf16_rows")?,
            is_weights_f32: ld("is_weights_f32")?,
            normalize_weights_f32: ld("normalize_weights_f32")?,
            fill_from_gpu_f32: ld("fill_from_gpu_f32")?,
            reduce_max_f32: ld("reduce_max_f32")?,
            i64_to_u32: ld("i64_to_u32")?,
            f32_to_bf16_cast: ld("f32_to_bf16_cast")?,
            max_of_two_f32: ld("max_of_two_f32")?,
            seg_tree_update: st_mod.load_function("seg_tree_update")
                .map_err(|e| MLError::ModelError(format!("st_update fn: {e}")))?,
            seg_tree_insert: st_mod.load_function("seg_tree_insert")
                .map_err(|e| MLError::ModelError(format!("st_insert fn: {e}")))?,
            seg_tree_sample: st_mod.load_function("seg_tree_sample")
                .map_err(|e| MLError::ModelError(format!("st_sample fn: {e}")))?,
            seg_tree_gather_prios: st_mod.load_function("seg_tree_gather_prios")
                .map_err(|e| MLError::ModelError(format!("st_gather fn: {e}")))?,
        })
    }
}

fn lcfg(n: usize) -> LaunchConfig {
    let t = 256_u32; let b = (n as u32).div_ceil(t);
    LaunchConfig { grid_dim: (b.max(1), 1, 1), block_dim: (t, 1, 1), shared_mem_bytes: 0 }
}

// ---------------------------------------------------------------------------
// Config + Buffer
// ---------------------------------------------------------------------------

#[allow(clippy::module_name_repetitions)]
#[derive(Debug, Clone)]
pub struct GpuReplayBufferConfig {
    pub capacity: usize, pub state_dim: usize, pub alpha: f32,
    pub beta_start: f32, pub beta_max: f32, pub beta_annealing_steps: usize,
    pub epsilon: f32, pub max_memory_bytes: usize,
    /// Maximum batch size for pre-allocated sampling buffers. Defaults to 1024.
    pub max_batch_size: usize,
}

pub struct GpuReplayBuffer {
    config: GpuReplayBufferConfig,
    stream: Arc<CudaStream>,
    kernels: ReplayKernels,
    states: CudaSlice<u16>, next_states: CudaSlice<u16>,
    actions: CudaSlice<u32>, rewards: CudaSlice<f32>,
    dones: CudaSlice<f32>, priorities: CudaSlice<f32>,
    /// Episode IDs per buffer slot `[capacity]` i32 on GPU.
    /// `episode_ids[i] = i / episode_length`. Written during `insert_batch_with_episode_ids`.
    episode_ids: CudaSlice<i32>,
    write_cursor: usize, size: usize,
    max_priority: CudaSlice<f32>,
    pending_max_priority: Option<CudaSlice<f32>>,
    current_step: usize,
    // Segment tree: flat array [2 * capacity_pow2], root at [1], leaves at [cap_pow2..2*cap_pow2]
    seg_tree: CudaSlice<f32>,
    capacity_pow2: usize,
    // Pre-allocated PER sampling buffers (zero cuMemAlloc after warmup)
    sample_indices_i64: CudaSlice<i64>,
    sample_indices_u32: CudaSlice<u32>,
    sample_states: CudaSlice<u16>,
    sample_next_states: CudaSlice<u16>,
    sample_actions: CudaSlice<u32>,
    sample_rewards: CudaSlice<f32>,
    sample_dones: CudaSlice<f32>,
    sample_priorities: CudaSlice<f32>,
    sample_weights: CudaSlice<f32>,
    sample_max_weight: CudaSlice<f32>,
    sample_episode_ids: CudaSlice<i32>,
    total_sum_buf: CudaSlice<f32>,
    rng_step: u32,
}

impl Drop for GpuReplayBuffer {
    fn drop(&mut self) {
        // Synchronize stream before CudaSlice fields drop.
        // The segment tree, PER sampling, and priority update kernels may have
        // pending work. Without sync, cuMemFree races with async kernel writes.
        #[allow(unsafe_code)]
        unsafe {
            cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
        }
    }
}

impl GpuReplayBuffer {
    pub fn new(config: GpuReplayBufferConfig, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let (cap, sd) = (config.capacity, config.state_dim);
        let mbs = if config.max_batch_size == 0 { 1024 } else { config.max_batch_size };
        let need = 2 * cap * sd * 2 + 5 * cap * 4;
        if need > config.max_memory_bytes {
            #[allow(clippy::integer_division)]
            return Err(MLError::ModelError(format!(
                "GPU replay buffer needs {} MB (limit {} MB)",
                need / (1024 * 1024), config.max_memory_bytes / (1024 * 1024),
            )));
        }
        let k = ReplayKernels::compile(stream)?;
        let s = a16(stream, cap * sd, "s")?;
        let ns = a16(stream, cap * sd, "ns")?;
        let a = a32u(stream, cap, "a")?;
        let r = a32f(stream, cap, "r")?;
        let d = a32f(stream, cap, "d")?;
        let p = a32f(stream, cap, "p")?;
        let mut mp = a32f(stream, 1, "mp")?;
        stream.memcpy_htod(&[1.0_f32], &mut mp).map_err(|e| MLError::ModelError(format!("mp: {e}")))?;

        // Segment tree: flat array [2 * capacity_pow2], all zeros initially.
        // Leaves at [cap_pow2..2*cap_pow2], root at [1]. Power-of-2 for balanced tree.
        let cap_pow2 = cap.next_power_of_two();
        let seg = a32f(stream, 2 * cap_pow2, "seg_tree")?;

        // Pre-allocate PER sampling buffers (zero cuMemAlloc after warmup)
        let si64 = stream.alloc_zeros::<i64>(mbs).map_err(|e| MLError::ModelError(format!("alloc s_i64: {e}")))?;
        let su32 = a32u(stream, mbs, "s_idx")?;
        let ss = a16(stream, mbs * sd, "s_states")?;
        let sns = a16(stream, mbs * sd, "s_nstates")?;
        let sa = a32u(stream, mbs, "s_act")?;
        let sr = a32f(stream, mbs, "s_rew")?;
        let sdn = a32f(stream, mbs, "s_done")?;
        let sp = a32f(stream, mbs, "s_pri")?;
        let sw = a32f(stream, mbs, "s_wt")?;
        let smw = a32f(stream, 1, "s_mw")?;
        let tsb = a32f(stream, 1, "ts_buf")?;
        let ep_ids = a32i(stream, cap, "episode_ids")?;
        let s_ep = a32i(stream, mbs, "s_episode_ids")?;

        Ok(Self {
            config, stream: Arc::clone(stream), kernels: k,
            states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p,
            episode_ids: ep_ids,
            write_cursor: 0, size: 0, max_priority: mp,
            pending_max_priority: None, current_step: 0,
            seg_tree: seg, capacity_pow2: cap_pow2,
            sample_indices_i64: si64,
            sample_indices_u32: su32, sample_states: ss,
            sample_next_states: sns, sample_actions: sa,
            sample_rewards: sr, sample_dones: sdn,
            sample_priorities: sp, sample_weights: sw,
            sample_episode_ids: s_ep,
            sample_max_weight: smw, total_sum_buf: tsb,
            rng_step: 0,
        })
    }

    pub const fn len(&self) -> usize { self.size }
    pub const fn capacity(&self) -> usize { self.config.capacity }
    pub const fn is_empty(&self) -> bool { self.size == 0 }
    pub const fn can_sample(&self, bs: usize) -> bool { self.size >= bs }
    pub fn current_beta(&self) -> f32 {
        if self.config.beta_annealing_steps == 0 { return self.config.beta_max; }
        let p = ((self.current_step as f32) / (self.config.beta_annealing_steps as f32)).min(1.0);
        self.config.beta_start + (self.config.beta_max - self.config.beta_start) * p
    }
    pub const fn step(&mut self) { self.current_step = self.current_step.saturating_add(1); }
    pub fn clear(&mut self) -> Result<(), MLError> {
        self.write_cursor = 0; self.size = 0;
        self.stream.memcpy_htod(&[1.0_f32], &mut self.max_priority).map_err(|e| MLError::ModelError(format!("{e}")))?;
        self.stream.memset_zeros(&mut self.seg_tree).map_err(|e| MLError::ModelError(format!("seg_tree clear: {e}")))?;
        self.pending_max_priority = None; self.current_step = 0; Ok(())
    }
    pub const fn stream(&self) -> &Arc<CudaStream> { &self.stream }
    pub const fn alpha(&self) -> f32 { self.config.alpha }
    pub const fn epsilon(&self) -> f32 { self.config.epsilon }
    pub const fn state_dim(&self) -> usize { self.config.state_dim }

    pub fn insert_batch(&mut self, sf: &CudaSlice<f32>, nf: &CudaSlice<f32>,
        ac: &CudaSlice<u32>, rw: &CudaSlice<f32>, dn: &CudaSlice<f32>, bs: usize,
    ) -> Result<(), MLError> {
        if bs == 0 { return Ok(()); }
        let (cap, sd) = (self.config.capacity, self.config.state_dim);
        let (eff, off) = if bs > cap { (cap, bs - cap) } else { (bs, 0) };
        let el = eff * sd;
        let mut b_s = a16(&self.stream, el, "ib_s")?;
        let mut b_n = a16(&self.stream, el, "ib_n")?;
        let ni = el as i32;
        let ss = if off > 0 { sf.slice(off * sd..) } else { sf.slice(0..) };
        let sn = if off > 0 { nf.slice(off * sd..) } else { nf.slice(0..) };
        // SAFETY: b_s, ss, b_n, sn are valid device allocations. Kernel indices are bounds-checked.
        unsafe {
            self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
                .arg(&mut b_s).arg(&ss).arg(&ni).launch(lcfg(el))
                .map_err(|e| MLError::ModelError(format!("cast s: {e}")))?;
        }
        // SAFETY: same as above for next-state cast.
        unsafe {
            self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
                .arg(&mut b_n).arg(&sn).arg(&ni).launch(lcfg(el))
                .map_err(|e| MLError::ModelError(format!("cast n: {e}")))?;
        }
        let (ci, cpi, sdi, bsi) = (self.write_cursor as i32, cap as i32, sd as i32, eff as i32);
        // SAFETY: states, b_s, next_states, b_n are valid device allocations. Indices within capacity.
        unsafe {
            self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
                .arg(&self.states).arg(&b_s).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
                .launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc s: {e}")))?;
        }
        // SAFETY: same as above for next-state scatter.
        unsafe {
            self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
                .arg(&self.next_states).arg(&b_n).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
                .launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc n: {e}")))?;
        }
        let sa = if off > 0 { ac.slice(off..) } else { ac.slice(0..) };
        let sr = if off > 0 { rw.slice(off..) } else { rw.slice(0..) };
        let sd2 = if off > 0 { dn.slice(off..) } else { dn.slice(0..) };
        // SAFETY: actions, rewards, dones, priorities are valid device allocations within capacity.
        unsafe {
            self.stream.launch_builder(&self.kernels.scatter_insert_u32)
                .arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?;
        }
        // SAFETY: same context, rewards buffer valid.
        unsafe {
            self.stream.launch_builder(&self.kernels.scatter_insert_f32)
                .arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r: {e}")))?;
        }
        // SAFETY: same context, dones buffer valid.
        unsafe {
            self.stream.launch_builder(&self.kernels.scatter_insert_f32)
                .arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d: {e}")))?;
        }
        let mut pt = a32f(&self.stream, eff, "pt")?;
        // SAFETY: GPU-only broadcast of max_priority scalar (zero CPU readback).
        unsafe {
            self.stream.launch_builder(&self.kernels.fill_from_gpu_f32)
                .arg(&mut pt).arg(&self.max_priority).arg(&bsi)
                .launch(lcfg(eff))
                .map_err(|e| MLError::ModelError(format!("fill mp: {e}")))?;
        }
        // SAFETY: priorities buffer valid, pt filled above.
        unsafe {
            self.stream.launch_builder(&self.kernels.scatter_insert_f32)
                .arg(&self.priorities).arg(&pt).arg(&ci).arg(&cpi).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc p: {e}")))?;
        }
        // Update segment tree leaves for inserted indices.
        // Build index array on CPU (small: batch_size <= 1024 typically) and upload.
        let insert_indices: Vec<u32> = (0..eff)
            .map(|j| ((self.write_cursor + j) % cap) as u32)
            .collect();
        let mut idx_buf = a32u(&self.stream, eff, "ib_idx")?;
        self.stream.memcpy_htod(&insert_indices, &mut idx_buf)
            .map_err(|e| MLError::ModelError(format!("ib idx htod: {e}")))?;
        let al = self.config.alpha;
        let cap_pow2_i = self.capacity_pow2 as i32;
        // SAFETY: seg_tree, idx_buf, pt are valid. seg_tree has 2*cap_pow2 elements.
        unsafe {
            self.stream.launch_builder(&self.kernels.seg_tree_insert)
                .arg(&self.seg_tree).arg(&idx_buf).arg(&pt).arg(&al)
                .arg(&cap_pow2_i).arg(&bsi)
                .launch(lcfg(eff))
                .map_err(|e| MLError::ModelError(format!("st insert: {e}")))?;
        }
        // Write episode IDs for inserted positions: simple sequential IDs (slot index).
        // For flat replay buffers, each buffer slot is its own "episode" (L=1).
        // HER Future/Final strategies use these to identify episode boundaries.
        let ep_ids_host: Vec<i32> = (0..eff)
            .map(|j| ((self.write_cursor + j) % cap) as i32)
            .collect();
        let mut ep_buf = a32i(&self.stream, eff, "ib_ep")?;
        self.stream.memcpy_htod(&ep_ids_host, &mut ep_buf)
            .map_err(|e| MLError::ModelError(format!("ep htod: {e}")))?;
        // SAFETY: episode_ids and ep_buf are valid device allocations. Reinterpret i32 as u32 (same size).
        unsafe {
            let ep_dst = &*(&self.episode_ids as *const CudaSlice<i32> as *const CudaSlice<u32>);
            let ep_src = &*(&ep_buf as *const CudaSlice<i32> as *const CudaSlice<u32>);
            self.stream.launch_builder(&self.kernels.scatter_insert_u32)
                .arg(ep_dst).arg(ep_src).arg(&ci).arg(&cpi).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc ep: {e}")))?;
        }

        self.write_cursor = (self.write_cursor + eff) % cap;
        self.size = (self.size + eff).min(cap);
        Ok(())
    }

    pub fn sample_proportional(&mut self, batch_size: usize) -> Result<GpuBatchSlices, MLError> {
        let _nvtx = NvtxRange::new("per_sample_proportional");
        if !self.can_sample(batch_size) {
            return Err(MLError::ModelError(format!("Cannot sample {batch_size} from {}", self.size)));
        }
        let mbs = if self.config.max_batch_size == 0 { 1024 } else { self.config.max_batch_size };
        if batch_size > mbs {
            return Err(MLError::ModelError(format!(
                "batch_size {batch_size} exceeds max_batch_size {mbs}"
            )));
        }
        let (n, sd) = (self.size, self.config.state_dim);
        let nb = -self.current_beta();
        let bsi = batch_size as i32;
        let cap_pow2_i = self.capacity_pow2 as i32;
        let buf_size_i = n as i32;

        // Step 1: Segment tree sample — parallel root-to-leaf traversal with Philox RNG.
        // O(log n) per sample instead of O(n) prefix sum + O(log n) binary search.
        // Replaces: pow_alpha + prefix_sum + threshold_gen + searchsorted.
        self.rng_step = self.rng_step.wrapping_add(1);
        let seed = self.rng_step;
        // SAFETY: seg_tree, sample_indices_i64 are valid device allocations.
        unsafe {
            self.stream.launch_builder(&self.kernels.seg_tree_sample)
                .arg(&self.seg_tree).arg(&mut self.sample_indices_i64)
                .arg(&seed).arg(&cap_pow2_i).arg(&bsi).arg(&buf_size_i)
                .launch(lcfg(batch_size))
                .map_err(|e| MLError::ModelError(format!("st sample: {e}")))?;
        }

        // Step 2: i64 -> u32 indices for output
        // SAFETY: sample_indices_u32 and sample_indices_i64 are valid device allocations.
        unsafe {
            self.stream.launch_builder(&self.kernels.i64_to_u32)
                .arg(&mut self.sample_indices_u32).arg(&self.sample_indices_i64).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("i2u: {e}")))?;
        }

        // Step 3: gather into pre-allocated buffers
        let sdi = sd as i32;
        // SAFETY: sample_states, states, sample_indices_i64 are valid device allocations.
        unsafe {
            self.stream.launch_builder(&self.kernels.gather_bf16_rows)
                .arg(&mut self.sample_states).arg(&self.states).arg(&self.sample_indices_i64).arg(&sdi).arg(&bsi)
                .launch(lcfg(batch_size * sd)).map_err(|e| MLError::ModelError(format!("g s: {e}")))?;
        }
        // SAFETY: same as above for next_states gather.
        unsafe {
            self.stream.launch_builder(&self.kernels.gather_bf16_rows)
                .arg(&mut self.sample_next_states).arg(&self.next_states).arg(&self.sample_indices_i64).arg(&sdi).arg(&bsi)
                .launch(lcfg(batch_size * sd)).map_err(|e| MLError::ModelError(format!("g n: {e}")))?;
        }
        // SAFETY: same context, actions buffer valid.
        unsafe {
            self.stream.launch_builder(&self.kernels.gather_u32)
                .arg(&mut self.sample_actions).arg(&self.actions).arg(&self.sample_indices_i64).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g a: {e}")))?;
        }
        // SAFETY: same context, rewards buffer valid.
        unsafe {
            self.stream.launch_builder(&self.kernels.gather_f32)
                .arg(&mut self.sample_rewards).arg(&self.rewards).arg(&self.sample_indices_i64).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g r: {e}")))?;
        }
        // SAFETY: same context, dones buffer valid.
        unsafe {
            self.stream.launch_builder(&self.kernels.gather_f32)
                .arg(&mut self.sample_dones).arg(&self.dones).arg(&self.sample_indices_i64).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g d: {e}")))?;
        }

        // Step 3b: gather episode_ids for HER strategies.
        // Reinterpret i32 buffers as u32 (same bit width) for the gather_u32 kernel.
        // SAFETY: episode_ids and sample_episode_ids are CudaSlice<i32>, same layout as u32.
        // The gather kernel copies raw bytes, so reinterpret is safe for same-size types.
        unsafe {
            let ep_src = &*(&self.episode_ids as *const CudaSlice<i32> as *const CudaSlice<u32>);
            let ep_dst = &mut *(&mut self.sample_episode_ids as *mut CudaSlice<i32> as *mut CudaSlice<u32>);
            self.stream.launch_builder(&self.kernels.gather_u32)
                .arg(ep_dst).arg(ep_src).arg(&self.sample_indices_i64).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g ep: {e}")))?;
        }

        // Step 4: gather sampled priority^alpha from segment tree leaves for IS weights.
        // SAFETY: sample_priorities, seg_tree, sample_indices_i64 are valid device allocations.
        unsafe {
            self.stream.launch_builder(&self.kernels.seg_tree_gather_prios)
                .arg(&mut self.sample_priorities).arg(&self.seg_tree)
                .arg(&self.sample_indices_i64).arg(&cap_pow2_i).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g sp: {e}")))?;
        }

        // Step 5: DtoD copy tree[1] (root = total sum) -> total_sum_buf for IS weights
        {
            let tree_root = self.seg_tree.slice(1..2);
            let num_bytes = std::mem::size_of::<f32>();
            let src_ptr = {
                let (ptr, guard) = tree_root.device_ptr(&self.stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };
            let dst_ptr = {
                let (ptr, guard) = self.total_sum_buf.device_ptr_mut(&self.stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };
            // SAFETY: src_ptr and dst_ptr are valid device pointers. num_bytes = sizeof(f32).
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("ts dtod: {e}")))?;
            }
        }

        // Step 6: IS weights via GPU-resident total_sum (zero CPU readback)
        // is_weights_f32 reads total_sum from GPU pointer.
        // SAFETY: sample_weights, sample_priorities, total_sum_buf are valid device allocations.
        unsafe {
            self.stream.launch_builder(&self.kernels.is_weights_f32)
                .arg(&mut self.sample_weights).arg(&self.sample_priorities).arg(&self.total_sum_buf)
                .arg(&nb).arg(&(n as i32)).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("isw: {e}")))?;
        }

        // Step 7: reduce max weight
        self.stream.memset_zeros(&mut self.sample_max_weight)
            .map_err(|e| MLError::ModelError(format!("mw zero: {e}")))?;
        let rt = 256_u32.min(batch_size as u32).max(1);
        let rb = (batch_size as u32).div_ceil(rt);
        // SAFETY: sample_weights, sample_max_weight are valid device allocations.
        unsafe {
            self.stream.launch_builder(&self.kernels.reduce_max_f32)
                .arg(&self.sample_weights).arg(&mut self.sample_max_weight).arg(&bsi)
                .launch(LaunchConfig { grid_dim: (rb.max(1),1,1), block_dim: (rt,1,1), shared_mem_bytes: rt*4 })
                .map_err(|e| MLError::ModelError(format!("rm: {e}")))?;
        }

        // Step 8: normalize weights by max
        // SAFETY: sample_weights, sample_max_weight are valid. Normalization divides element-wise.
        unsafe {
            self.stream.launch_builder(&self.kernels.normalize_weights_f32)
                .arg(&mut self.sample_weights).arg(&self.sample_max_weight).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("nw: {e}")))?;
        }

        // Return DtoD clones of pre-allocated slices sized to actual batch_size.
        // The caller owns the returned GpuBatchSlices (consumed by into_gpu_batch),
        // so we DtoD-clone the relevant portions. All copies are async on the stream.
        let ep_ids = dtod_clone_i32(&self.stream, &self.sample_episode_ids, batch_size, "o_ep")?;

        Ok(GpuBatchSlices {
            states: dtod_clone_u16(&self.stream, &self.sample_states, batch_size * sd, "o_s")?,
            next_states: dtod_clone_u16(&self.stream, &self.sample_next_states, batch_size * sd, "o_n")?,
            actions: dtod_clone_u32(&self.stream, &self.sample_actions, batch_size, "o_act")?,
            rewards: dtod_clone_f32(&self.stream, &self.sample_rewards, batch_size, "o_r")?,
            dones: dtod_clone_f32(&self.stream, &self.sample_dones, batch_size, "o_d")?,
            weights: dtod_clone_f32(&self.stream, &self.sample_weights, batch_size, "o_w")?,
            indices: dtod_clone_u32(&self.stream, &self.sample_indices_u32, batch_size, "o_i")?,
            episode_ids: Some(ep_ids),
            batch_size,
            state_dim: sd,
        })
    }

    /// Update priorities from GPU-resident index and `td_error` `CudaSlices`.
    ///
    /// Uses the segment tree update kernel: computes `(|td|^alpha + eps)`,
    /// writes to priorities buffer AND tree leaves, propagates sums to root.
    /// O(log n) per update instead of O(n) prefix sum rebuild.
    pub fn update_priorities_gpu(
        &mut self,
        indices: &CudaSlice<u32>,
        td_errors: &CudaSlice<f32>,
        bs: usize,
    ) -> Result<(), MLError> {
        let _nvtx = NvtxRange::new("per_update_priorities");
        if bs == 0 { return Ok(()); }
        let (al, ep, bsi) = (self.config.alpha, self.config.epsilon, bs as i32);
        let cap_pow2_i = self.capacity_pow2 as i32;
        let mut bm = a32f(&self.stream, 1, "bm")?;
        self.stream.memset_zeros(&mut bm).map_err(|e| MLError::ModelError(format!("{e}")))?;
        // seg_tree_update: writes priorities[], tree leaves, propagates, atomicMax.
        // SAFETY: seg_tree, priorities, bm, indices, td_errors are valid device allocations.
        unsafe {
            self.stream.launch_builder(&self.kernels.seg_tree_update)
                .arg(&self.seg_tree).arg(&self.priorities).arg(&mut bm)
                .arg(indices).arg(td_errors)
                .arg(&al).arg(&ep).arg(&cap_pow2_i).arg(&bsi)
                .launch(lcfg(bs)).map_err(|e| MLError::ModelError(format!("st update: {e}")))?;
        }
        self.pending_max_priority = Some(match self.pending_max_priority.take() {
            Some(prev) => {
                let mut r = a32f(&self.stream, 1, "mm")?;
                // SAFETY: r, prev, bm are valid 1-element device allocations on same context.
                unsafe {
                    self.stream.launch_builder(&self.kernels.max_of_two_f32)
                        .arg(&mut r).arg(&prev).arg(&bm)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
                        .map_err(|e| MLError::ModelError(format!("max_of_two: {e}")))?;
                }
                r
            }
            None => bm,
        });
        Ok(())
    }

    pub fn flush_max_priority(&mut self) -> Result<(), MLError> {
        if let Some(pend) = self.pending_max_priority.take() {
            // GPU-only max of pending vs current (zero CPU readback)
            let mut result = a32f(&self.stream, 1, "fm")?;
            // SAFETY: result, pend, max_priority are valid 1-element device allocations.
            unsafe {
                self.stream.launch_builder(&self.kernels.max_of_two_f32)
                    .arg(&mut result).arg(&pend).arg(&self.max_priority)
                    .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
                    .map_err(|e| MLError::ModelError(format!("flush max_of_two: {e}")))?;
            }
            // DtoD copy result into self.max_priority
            let num_bytes = std::mem::size_of::<f32>();
            let src_ptr = {
                let (ptr, guard) = result.device_ptr(&self.stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };
            let dst_ptr = {
                let (ptr, guard) = self.max_priority.device_ptr_mut(&self.stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };
            // SAFETY: dst_ptr and src_ptr are valid device pointers. num_bytes = sizeof(f32).
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("flush dtod: {e}")))?;
            }
        }
        Ok(())
    }

    pub fn apply_max_priority_scalar(&mut self, mp: f32) -> Result<(), MLError> {
        if mp > 0.0 {
            // Upload CPU scalar to GPU, then GPU-only max comparison (zero download)
            let mut mp_gpu = a32f(&self.stream, 1, "amp")?;
            self.stream.memcpy_htod(&[mp], &mut mp_gpu)
                .map_err(|e| MLError::ModelError(format!("amp htod: {e}")))?;
            let mut result = a32f(&self.stream, 1, "amr")?;
            // SAFETY: result, mp_gpu, max_priority are valid 1-element device allocations.
            unsafe {
                self.stream.launch_builder(&self.kernels.max_of_two_f32)
                    .arg(&mut result).arg(&mp_gpu).arg(&self.max_priority)
                    .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
                    .map_err(|e| MLError::ModelError(format!("amp max: {e}")))?;
            }
            // DtoD copy result into self.max_priority
            let num_bytes = std::mem::size_of::<f32>();
            let src_ptr = {
                let (ptr, guard) = result.device_ptr(&self.stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };
            let dst_ptr = {
                let (ptr, guard) = self.max_priority.device_ptr_mut(&self.stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };
            // SAFETY: dst_ptr and src_ptr are valid device pointers. num_bytes = sizeof(f32).
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("amp dtod: {e}")))?;
            }
        }
        Ok(())
    }

    /// Raw `CudaSlice` accessors for direct GPU consumption.
    pub const fn states_slice(&self) -> &CudaSlice<u16> { &self.states }
    pub const fn next_states_slice(&self) -> &CudaSlice<u16> { &self.next_states }
    pub const fn actions_slice(&self) -> &CudaSlice<u32> { &self.actions }
    pub const fn rewards_slice(&self) -> &CudaSlice<f32> { &self.rewards }
    pub const fn dones_slice(&self) -> &CudaSlice<f32> { &self.dones }
    pub const fn priorities_slice(&self) -> &CudaSlice<f32> { &self.priorities }

    /// Sample proportional indices and IS weights as GPU-resident `CudaSlices`.
    ///
    /// Returns `(indices: CudaSlice<u32>, weights: CudaSlice<f32>)` on GPU.
    /// Callers process indices on GPU -- zero CPU download.
    pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice<u32>, CudaSlice<f32>), MLError> {
        let b = self.sample_proportional(bs)?;
        Ok((b.indices, b.weights))
    }

}

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

fn a32f(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<f32>, MLError> {
    s.alloc_zeros::<f32>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
}
fn a32u(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<u32>, MLError> {
    s.alloc_zeros::<u32>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
}
fn a16(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<u16>, MLError> {
    s.alloc_zeros::<u16>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
}
fn a32i(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<i32>, MLError> {
    s.alloc_zeros::<i32>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
}

/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
fn dtod_clone_f32(s: &Arc<CudaStream>, src: &CudaSlice<f32>, n: usize, nm: &str) -> Result<CudaSlice<f32>, MLError> {
    let mut dst = a32f(s, n, nm)?;
    let sv = src.slice(..n);
    s.memcpy_dtod(&sv, &mut dst).map_err(|e| MLError::ModelError(format!("dtod {nm}: {e}")))?;
    Ok(dst)
}

/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
fn dtod_clone_u32(s: &Arc<CudaStream>, src: &CudaSlice<u32>, n: usize, nm: &str) -> Result<CudaSlice<u32>, MLError> {
    let mut dst = a32u(s, n, nm)?;
    let sv = src.slice(..n);
    s.memcpy_dtod(&sv, &mut dst).map_err(|e| MLError::ModelError(format!("dtod {nm}: {e}")))?;
    Ok(dst)
}

/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
fn dtod_clone_i32(s: &Arc<CudaStream>, src: &CudaSlice<i32>, n: usize, nm: &str) -> Result<CudaSlice<i32>, MLError> {
    let mut dst = a32i(s, n, nm)?;
    let sv = src.slice(..n);
    s.memcpy_dtod(&sv, &mut dst).map_err(|e| MLError::ModelError(format!("dtod {nm}: {e}")))?;
    Ok(dst)
}

/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
fn dtod_clone_u16(s: &Arc<CudaStream>, src: &CudaSlice<u16>, n: usize, nm: &str) -> Result<CudaSlice<u16>, MLError> {
    let mut dst = a16(s, n, nm)?;
    let sv = src.slice(..n);
    s.memcpy_dtod(&sv, &mut dst).map_err(|e| MLError::ModelError(format!("dtod {nm}: {e}")))?;
    Ok(dst)
}

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

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
            .expect("CUDA required")
            .new_stream()
            .expect("CUDA stream")
    }

    #[test]
    fn test_creation() {
        let c = GpuReplayBufferConfig { capacity: 1000, state_dim: 48, alpha: 0.6, beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 100_000, epsilon: 1e-6, max_memory_bytes: 4<<30, max_batch_size: 256 };
        let b = GpuReplayBuffer::new(c, &make_stream()).expect("buf");
        assert_eq!(b.len(), 0); assert_eq!(b.capacity(), 1000); assert!(b.is_empty());
    }
    #[test]
    fn test_beta() {
        let c = GpuReplayBufferConfig { capacity: 100, state_dim: 4, alpha: 0.6, beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6, max_memory_bytes: 4<<30, max_batch_size: 64 };
        let mut b = GpuReplayBuffer::new(c, &make_stream()).expect("buf");
        assert!((b.current_beta() - 0.4).abs() < 1e-6);
        for _ in 0..500 { b.step(); }
        assert!(b.current_beta() > 0.4 && b.current_beta() < 1.0);
        for _ in 0..600 { b.step(); }
        assert!((b.current_beta() - 1.0).abs() < 1e-6);
    }
    #[test]
    fn test_clear() {
        let c = GpuReplayBufferConfig { capacity: 100, state_dim: 4, alpha: 0.6, beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6, max_memory_bytes: 4<<30, max_batch_size: 64 };
        let mut b = GpuReplayBuffer::new(c, &make_stream()).expect("buf");
        b.step(); b.clear().expect("clear");
        assert_eq!(b.len(), 0); assert_eq!(b.current_step, 0);
    }
}
