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

//! GPU-accelerated DQN experience collection via timestep-level loop + cuBLAS.
//!
//! Replaces the monolithic `dqn_full_experience_kernel` (3,272 lines, 1.56%
//! occupancy from warp-cooperative matvec) with a host-driven timestep loop
//! that calls:
//!   1. `experience_state_gather`  — assemble batch state tensor (NVRTC kernel)
//!   2. `CublasForward::forward_online` — Q-network forward via cuBLAS SGEMM
//!   3. `compute_expected_q`       — C51 distributional → expected Q-values
//!   4. `experience_action_select` — epsilon-greedy on Q-values (NVRTC kernel)
//!   5. `experience_env_step`      — portfolio simulation + reward (NVRTC kernel)
//!
//! cuBLAS SGEMM processes the entire batch in a single GEMM call per layer,
//! raising GPU occupancy from ~1.56% to >60% on H100.

use std::ffi::c_void;
use std::mem::ManuallyDrop;
use std::sync::Arc;

use cudarc::driver::{
    CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg,
};
use ml_core::cuda_autograd::GpuVarStore;
use tracing::{debug, info};

use crate::MLError;
use super::batched_forward::{CublasForward, f32_weight_ptrs};
use super::gpu_curiosity_trainer::GpuCuriosityTrainer;
use super::gpu_dqn_trainer::{
    GpuDqnTrainConfig, compute_param_sizes, compute_total_params, dtod_copy, raw_device_ptr,
};
use super::gpu_weights::{
    BranchingWeightSet, CuriosityWeightSet, DuelingWeightSet,
    RmsNormWeightSet,
    extract_branching_weights, extract_curiosity_weights,
    extract_dueling_weights_branching, extract_rmsnorm_weights,
    sync_branching_weights, sync_curiosity_weights,
    sync_dueling_weights_branching, sync_rmsnorm_weights,
};

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

/// Absolute upper bound for validation — reject configs above this.
/// H100 (132 SMs) can run 32768 episodes; `optimal_n_episodes()` sizes dynamically.
const MAX_EPISODES_LIMIT: usize = 0x8000;
/// Absolute upper bound for validation — reject configs above this.
const MAX_TIMESTEPS_LIMIT: usize = 1000;
const PORTFOLIO_STATE_SIZE: usize = 8;
/// Portfolio stride for DQN experience kernels (20 floats per episode).
/// Matches PORTFOLIO_STRIDE in experience_kernels.cu.
/// Do NOT change PORTFOLIO_STATE_SIZE above — it's for the PPO/legacy path.
const PORTFOLIO_STRIDE: usize = 20;

/// Number of floats in the trade_stats_reduce output buffer.
/// Layout: [win_count, loss_count, sum_wins, sum_losses, sum_returns, sum_sq_returns]
const TRADE_STATS_FLOATS: usize = 6;

// ---------------------------------------------------------------------------
// TradeStats — GPU-sourced per-epoch trade statistics
// ---------------------------------------------------------------------------

/// Per-epoch trade statistics extracted from GPU portfolio_states[N, 14:19].
///
/// These fields are accumulated by `experience_env_step` kernel at each trade
/// completion event (flat→position or position→flat transition). The
/// `collect_trade_stats()` method runs a GPU reduction kernel to sum across
/// all N episodes, then downloads the 6-float result.
#[derive(Debug, Clone, Default)]
pub struct TradeStats {
    /// Total completed trades (entries that later exited back to flat)
    pub total_trades: usize,
    /// Trades with positive return
    pub winning_trades: usize,
    /// Trades with negative or zero return
    pub losing_trades: usize,
    /// Sum of returns from winning trades (positive)
    pub sum_wins: f64,
    /// Sum of |returns| from losing trades (positive)
    pub sum_losses: f64,
    /// Net sum of all trade returns
    pub sum_returns: f64,
    /// Sum of squared trade returns (for Sharpe/Sortino computation)
    pub sum_sq_returns: f64,
    /// Per-bar step returns from rewards_out (for MaxDD and Sortino)
    pub step_returns: Vec<f64>,
}

// ---------------------------------------------------------------------------
// Pinned host memory for DtoH transfers
// ---------------------------------------------------------------------------

/// Pinned (page-locked) host buffer for async DtoH memcpy.
///
/// CUDA DMA engine can transfer directly to pinned memory without staging
/// through a system bounce buffer, achieving ~2x PCIe bandwidth vs pageable.
/// All buffers are allocated once in the constructor and reused across launches.
struct PinnedHostBuf<T: Copy> {
    ptr: *mut T,
    len: usize,
}

// Safety: PinnedHostBuf is only accessed from the thread that created the
// GpuExperienceCollector, which holds a CUDA context. The ptr is valid
// for the lifetime of the allocation (freed in Drop).
unsafe impl<T: Copy> Send for PinnedHostBuf<T> {}
unsafe impl<T: Copy> Sync for PinnedHostBuf<T> {}

impl<T: Copy> PinnedHostBuf<T> {
    /// Allocate `len` elements of pinned host memory.
    ///
    /// # Safety
    /// Caller must ensure a CUDA context is active on the current thread.
    unsafe fn new(len: usize) -> Result<Self, MLError> {
        if len == 0 {
            return Ok(Self { ptr: std::ptr::null_mut(), len: 0 });
        }
        let num_bytes = len * std::mem::size_of::<T>();
        // CU_MEMHOSTALLOC_PORTABLE: accessible from any CUDA context
        let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;
        let host_ptr = cudarc::driver::result::malloc_host(num_bytes, flags)
            .map_err(|e| MLError::ModelError(format!(
                "pinned host alloc ({len} x {} bytes): {e}",
                std::mem::size_of::<T>(),
            )))?
            as *mut T;
        // Zero-init for safety
        std::ptr::write_bytes(host_ptr, 0, len);
        Ok(Self { ptr: host_ptr, len })
    }
}

impl<T: Copy> Drop for PinnedHostBuf<T> {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            // Safety: ptr was returned by malloc_host and is non-null.
            unsafe {
                #[allow(clippy::let_underscore_must_use)]
                let _ = cudarc::driver::result::free_host(self.ptr as *mut c_void);
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

/// Configuration for the GPU experience collection kernel.
#[derive(Debug, Clone)]
pub struct ExperienceCollectorConfig {
    /// Maximum absolute position size
    pub max_position: f32,
    /// Episode length in bars (for wrap detection)
    pub episode_length: i32,
    /// Total bars in the market data buffer
    pub total_bars: i32,
    /// Epsilon for epsilon-greedy action selection
    pub epsilon: f32,
    /// Discount factor for TD target computation
    pub gamma: f32,
    /// Maximum curiosity reward (clamped)
    pub curiosity_max_reward: f32,
    /// Number of episodes to run
    pub n_episodes: i32,
    /// Timesteps per episode
    pub timesteps_per_episode: i32,
    /// Barrier profit multiplier (e.g. 1.02)
    pub barrier_profit_mult: f32,
    /// Barrier loss multiplier (e.g. 0.98)
    pub barrier_loss_mult: f32,
    /// Barrier max bars before forced exit (e.g. 500.0)
    pub barrier_max_bars: f32,
    /// Barrier reward scaling factor (e.g. 0.5)
    pub barrier_scale: f32,
    /// Diversity entropy reward scaling factor (e.g. 1.0)
    pub diversity_scale: f32,
    /// Curiosity intrinsic reward scaling factor (e.g. 1.0)
    pub curiosity_scale: f32,
    /// Risk penalty weight (e.g. 0.1)
    pub risk_weight: f32,
    /// Asymmetric loss scaling factor (prospect theory, default 1.5)
    pub loss_aversion: f32,
    /// Transaction cost multiplier (hyperopt-tunable, scales base 0.01% rate)
    pub tx_cost_multiplier: f32,
    /// UCB count-bonus coefficient for GPU action selection (0.0 = disabled)
    pub count_bonus_coefficient: f32,
    /// Minimum Q-value clamp (e.g. -200.0)
    pub q_clip_min: f32,
    /// Maximum Q-value clamp (e.g. 200.0)
    pub q_clip_max: f32,
    /// Huber loss kappa for TD error (robust priority). 0.0 = disabled (raw L1).
    pub huber_kappa: f32,
    /// D5: Whether to use NoisyNet exploration in the GPU kernel
    pub use_noisy_nets: bool,
    /// D5: Initial noise std dev for NoisyNet (Rainbow DQN default: 0.5)
    pub noisy_sigma_init: f32,
    /// D6: Whether to use C51 distributional RL
    pub use_distributional: bool,
    /// D6: Number of atoms in C51 distribution (default: 51)
    pub num_atoms: i32,
    /// D6: Minimum value support for C51 (default: -50.0)
    pub v_min: f32,
    /// D6: Maximum value support for C51 (default: 50.0)
    pub v_max: f32,
    /// EMA decay rate for per-thread reward normalization in the GPU kernel.
    pub reward_norm_alpha: f32,
    /// Fill simulation: median spread for order routing (0.0 = disabled)
    pub fill_median_spread: f32,
    /// Fill simulation: median volatility for urgency routing
    pub fill_median_vol: f32,
    /// Fill simulation: IoC base fill probability (default 0.85)
    pub fill_ioc_fill_prob: f32,
    /// Fill simulation: minimum limit fill probability (default 0.30)
    pub fill_limit_fill_min: f32,
    /// Fill simulation: maximum limit fill probability (default 0.80)
    pub fill_limit_fill_max: f32,
    /// Fill simulation: fraction of spread paid by market orders (default 0.50)
    pub fill_spread_cost_frac: f32,
    /// Fill simulation: fraction of spread captured by limit orders (default 0.50)
    pub fill_spread_capture_frac: f32,
    /// Whether fill simulation is enabled (requires median_spread > 0)
    pub fill_simulation_enabled: bool,
    /// Q-value gap threshold for trade conviction filter.
    /// When greedy Q(best_exposure) - Q(flat) < threshold, default to flat.
    /// 0.0 = disabled (trade on every bar). Typical range [0.0, 0.5].
    pub q_gap_threshold: f32,
    /// Kelly optimal fraction for position sizing (0.0-1.0).
    /// Computed per-epoch on CPU from trade statistics, uploaded as scalar.
    /// 1.0 = full Kelly (aggressive), 0.5 = half-Kelly (conservative).
    /// DSR EMA decay rate (e.g. 0.01 = ~100-step window). DSR is always enabled.
    pub dsr_eta: f32,
    /// N-step returns lookahead (1 = standard TD, 3-5 typical for Rainbow DQN)
    pub n_steps: i32,
    /// Enable action masking: filters invalid exposure actions in the GPU kernel.
    pub enable_action_masking: bool,
}

impl Default for ExperienceCollectorConfig {
    fn default() -> Self {
        Self {
            max_position: 1.0,
            episode_length: 500,
            total_bars: 10_000,
            epsilon: 0.1,
            gamma: 0.99,
            curiosity_max_reward: 0.1,
            n_episodes: 2048,
            timesteps_per_episode: 500,
            barrier_profit_mult: 1.02,
            barrier_loss_mult: 0.98,
            barrier_max_bars: 500.0,
            barrier_scale: 0.5,
            diversity_scale: 1.0,
            curiosity_scale: 1.0,
            risk_weight: 0.1,
            loss_aversion: 1.5,
            tx_cost_multiplier: 1.0,
            count_bonus_coefficient: 0.0,
            q_clip_min: -200.0,  // Reward v6: tighter than old -500 but covers v_range + safety margin
            q_clip_max: 200.0,
            huber_kappa: 0.0,
            use_noisy_nets: true,   // Rainbow DQN: NoisyNet exploration always enabled
            noisy_sigma_init: 0.5,
            use_distributional: true,  // Rainbow DQN: C51 distributional RL always enabled
            num_atoms: 51,          // C51 standard atom count
            // v_min/v_max computed from reward_scale=10, gamma=0.95:
            // Q* ~ 10 / (1-0.95) * 1.2 = 240, clamped to [20, 300] → ±240
            v_min: -(10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            v_max: (10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            reward_norm_alpha: 0.0, // Disabled for reward v4: per-bar returns are already naturally scaled
            fill_median_spread: 0.0,
            fill_median_vol: 0.0,
            fill_ioc_fill_prob: 0.85,
            fill_limit_fill_min: 0.30,
            fill_limit_fill_max: 0.80,
            fill_spread_cost_frac: 0.50,
            fill_spread_capture_frac: 0.50,
            fill_simulation_enabled: false,
            q_gap_threshold: 0.0,
            dsr_eta: 0.01,
            n_steps: 1,
            enable_action_masking: false,
        }
    }
}

impl ExperienceCollectorConfig {
    /// Total number of experience tuples produced by one kernel launch.
    pub fn total_experiences(&self) -> usize {
        self.n_episodes as usize * self.timesteps_per_episode as usize
    }
}

// ---------------------------------------------------------------------------
// Output batch
// ---------------------------------------------------------------------------

// ExperienceBatch (CPU Vec<f32>) DELETED — use GpuExperienceBatch only.
// Production training loop uses collect_experiences_gpu() which keeps data GPU-resident.

/// GPU-resident experience batch -- pure cudarc `CudaSlice`.
///
/// All buffers are raw `CudaSlice` on the collector's forked stream.
/// The caller syncs the stream before reading or converting on another stream.
#[allow(missing_debug_implementations)]
pub struct GpuExperienceBatch {
    /// States `[total * state_dim]` on GPU (f32, row-major)
    pub states: CudaSlice<f32>,
    /// Next-states `[total * state_dim]` on GPU (f32, episode-aware shift)
    pub next_states: CudaSlice<f32>,
    /// Action indices `[total]` on GPU (i32, values 0..4)
    pub actions: CudaSlice<i32>,
    /// Rewards `[total]` on GPU (f32)
    pub rewards: CudaSlice<f32>,
    /// Done flags `[total]` on GPU (i32, 0 or 1)
    pub dones: CudaSlice<f32>,
    /// Episode index per transition `[total]` on GPU (i32).
    ///
    /// `episode_ids[i] = i / timesteps_per_episode`.
    /// Episode `ep` occupies flat indices `[ep * timesteps, (ep+1) * timesteps)`.
    /// Required by HER Future and Final strategies for GPU-native donor selection.
    pub episode_ids: CudaSlice<i32>,
    /// State dimensionality (so the consumer can reshape)
    pub state_dim: usize,
    /// Number of episodes
    pub n_episodes: usize,
    /// Timesteps per episode
    pub timesteps: usize,
}

// ---------------------------------------------------------------------------
// GpuExperienceCollector
// ---------------------------------------------------------------------------

/// GPU experience collector using timestep-level loop + cuBLAS forward pass.
///
/// Replaces the monolithic warp-cooperative experience kernel with 3 focused
/// NVRTC kernels (state gather, action select, env step) interleaved with
/// cuBLAS SGEMM for the Q-network forward pass. This raises GPU occupancy
/// from ~1.56% to >60% on H100.
///
/// Manages all GPU buffers for Q-network weights (as flat F32 parameter
/// buffer), per-episode state, and output arrays.
#[allow(missing_debug_implementations)]
pub struct GpuExperienceCollector {
    stream: Arc<CudaStream>,

    /// Actual STATE_DIM injected into the kernel (matches network input_dim).
    state_dim: usize,
    /// Market feature dimension (e.g. 42).
    market_dim: usize,
    /// Network hidden dims — (shared_h1, shared_h2, value_h, adv_h).
    network_dims: (usize, usize, usize, usize),
    /// Number of C51 atoms (1 for non-distributional).
    num_atoms: usize,
    /// C51 support range.
    v_min: f32,
    v_max: f32,
    /// Branch sizes [9, 3, 3] for branching DQN.
    branch_sizes: [usize; 3],

    /// Number of episodes buffers were allocated for (from config, not MAX constant).
    alloc_episodes: usize,
    /// Number of timesteps buffers were allocated for (from config, not MAX constant).
    alloc_timesteps: usize,

    // ── cuBLAS forward pass context ─────────────────────────────────
    cublas_forward: CublasForward,

    // ── Flat F32 online weight buffer for cuBLAS ────────────────────
    /// Flattened online network parameters [total_params] — single DtoD copy for sync.
    online_params_flat: CudaSlice<f32>,
    /// Total number of F32 parameters in the flat buffer.
    total_params: usize,
    /// Per-tensor sizes for computing weight pointers.
    param_sizes: [usize; 20],

    // ── CVaR position scaling (IQN dual-head) ──────────────────────
    /// Device pointer to CVaR scales [N]. 0 = NULL = no scaling.
    /// Set by `set_cvar_scales()` before experience collection.
    cvar_scales_ptr: u64,

    // ── Compiled experience kernels ─────────────────────────────────
    state_gather_kernel: CudaFunction,
    action_select_kernel: CudaFunction,
    expert_override_kernel: CudaFunction,
    env_step_kernel: CudaFunction,
    expected_q_kernel: CudaFunction,
    nstep_kernel: CudaFunction,
    _reward_norm_kernel: CudaFunction, // Kept for kernel compilation but never launched (reward v4)
    fill_episode_ids_kernel: CudaFunction,
    trade_stats_kernel: CudaFunction,

    // ── Per-timestep batch buffers (reused each timestep) ───────────
    /// Assembled state batch for cuBLAS: [N, state_dim].
    batch_states: CudaSlice<f32>,
    /// Selected actions per episode: [N].
    batch_actions: CudaSlice<i32>,
    /// Q-gap conviction per episode [N]: output by action_select, input to env_step.
    q_gaps_buf: CudaSlice<f32>,
    /// Kelly optimal fraction (per-epoch scalar from CPU).
    /// Current timestep counter per episode: [N].
    current_timesteps: CudaSlice<i32>,
    /// Q-values output from expected Q or direct logits: [N, total_branch_actions].
    q_values: CudaSlice<f32>,

    // ── cuBLAS activation scratch (reused each timestep) ────────────
    exp_h_s1: CudaSlice<f32>,      // [N, shared_h1]
    exp_h_s2: CudaSlice<f32>,      // [N, shared_h2]
    exp_h_v: CudaSlice<f32>,       // [N, value_h]
    exp_h_b0: CudaSlice<f32>,      // [N, adv_h]
    exp_h_b1: CudaSlice<f32>,      // [N, adv_h]
    exp_h_b2: CudaSlice<f32>,      // [N, adv_h]
    exp_v_logits: CudaSlice<f32>,  // [N, num_atoms]
    exp_b_logits: CudaSlice<f32>,  // [N, total_branch_atoms]

    // ── Legacy weight sets (for sync_online_weights/sync_target_weights compat) ──
    online_weights: DuelingWeightSet,
    target_weights: DuelingWeightSet,
    curiosity_weights: CuriosityWeightSet,
    online_rmsnorm: RmsNormWeightSet,
    target_rmsnorm: RmsNormWeightSet,
    online_branching: BranchingWeightSet,
    target_branching: BranchingWeightSet,

    /// GPU buffer for trade_stats_reduce output: [TRADE_STATS_FLOATS] = 6 floats.
    trade_stats_buf: CudaSlice<f32>,

    // Per-episode state buffers [alloc_episodes, ...]
    portfolio_states: CudaSlice<f32>,  // [alloc_episodes * PORTFOLIO_STRIDE] (20 floats per episode)
    rng_states: CudaSlice<u32>,        // [alloc_episodes]
    episode_starts_buf: CudaSlice<i32>,// [alloc_episodes]

    // Output buffers [alloc_episodes * alloc_timesteps, ...]
    states_out: CudaSlice<f32>,        // [alloc_episodes * alloc_timesteps * STATE_DIM]
    actions_out: CudaSlice<i32>,       // [alloc_episodes * alloc_timesteps]
    rewards_out: CudaSlice<f32>,       // [alloc_episodes * alloc_timesteps]
    done_out: CudaSlice<f32>,          // [alloc_episodes * alloc_timesteps]
    /// Raw per-bar portfolio returns (unshaped) for accurate Sharpe/MaxDD/Sortino.
    /// True fractional return: (equity_t - equity_{t-1}) / equity_{t-1}.
    raw_returns_out: CudaSlice<f32>,   // [alloc_episodes * alloc_timesteps]

    /// Running reward statistics for cross-epoch normalization (Welford's online algorithm)

    // Persistent epoch state — survives across kernel launches.
    epoch_state: CudaSlice<f32>,  // [8]
    /// Bitfield: bit 0 = reset portfolio, bit 1 = reset DSR, bit 2 = reset vol EMA
    reset_flags: u32,

    /// Expert demonstration actions [total_bars] — pre-computed MA crossover signals.
    /// -1 = no expert opinion (use Q-network), >= 0 = expert exposure action index.
    /// Uploaded once at training start, read by action selection kernel.
    expert_actions_gpu: Option<CudaSlice<i32>>,
    /// Expert demo override probability (0.0 = disabled, decays over epochs).
    expert_ratio: f32,

    // OFI features pre-uploaded to GPU [total_bars * OFI_DIM] (None when OFI disabled)
    ofi_gpu: Option<CudaSlice<f32>>,
    ofi_dim: usize,

    // GPU-resident curiosity forward model trainer (None if curiosity disabled)
    curiosity_trainer: Option<GpuCuriosityTrainer>,

    // Pre-allocated pinned host buffers for DtoH transfers.
    pinned_states: PinnedHostBuf<f32>,
    pinned_actions: PinnedHostBuf<i32>,
    pinned_rewards: PinnedHostBuf<f32>,
    pinned_dones: PinnedHostBuf<i32>,
}

impl Drop for GpuExperienceCollector {
    fn drop(&mut self) {
        // Synchronize stream before CudaSlice fields drop.
        // Pending GPU work (cuBLAS forward, kernel launches) may reference
        // buffers that cudarc's CudaSlice::drop frees. Without sync, the
        // driver keeps allocations alive → 2.5GB VRAM leak per hyperopt trial.
        #[allow(unsafe_code)]
        unsafe {
            cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
        }
    }
}

impl GpuExperienceCollector {
    /// Create a new GPU experience collector with cuBLAS-based Q-network forward pass.
    ///
    /// Compiles 3 focused NVRTC kernels, creates a `CublasForward` context for
    /// the Q-network, extracts initial weights, and allocates all GPU buffers.
    ///
    /// # Arguments
    /// * `stream` - CUDA stream for all GPU operations
    /// * `online_vars` - GpuVarStore for the online Q-network (Dueling architecture)
    /// * `target_vars` - GpuVarStore for the target Q-network
    /// * `curiosity_vars` - GpuVarStore for the curiosity forward model
    /// * `initial_capital` - Starting capital for each episode
    /// * `avg_spread` - Average bid-ask spread
    /// * `cash_reserve_pct` - Cash reserve percentage
    /// * `network_dims` - Actual layer dimensions `(shared_h1, shared_h2, value_h, adv_h)`
    /// * `kernel_dims` - Data layout dimensions `(state_dim, market_dim, num_atoms_max)`.
    /// * `n_episodes` - Number of episodes to allocate buffers for (clamped to 32768 max).
    /// * `timesteps_per_episode` - Timesteps per episode to allocate for (clamped to 1000 max).
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        stream: Arc<CudaStream>,
        online_vars: &GpuVarStore,
        target_vars: &GpuVarStore,
        curiosity_vars: Option<&GpuVarStore>,
        initial_capital: f32,
        _avg_spread: f32,
        _cash_reserve_pct: f32,
        network_dims: (usize, usize, usize, usize),
        kernel_dims: (usize, usize, usize),
        n_episodes: usize,
        timesteps_per_episode: usize,
    ) -> Result<Self, MLError> {
        // Drain stale CUDA errors from previous test/trainer instances.
        let _ = stream.context().check_err();

        let alloc_episodes = n_episodes.min(MAX_EPISODES_LIMIT);
        let alloc_timesteps = timesteps_per_episode.min(MAX_TIMESTEPS_LIMIT);
        let (shared_h1, shared_h2, value_h, adv_h) = network_dims;
        let (state_dim, market_dim, num_atoms_max) = kernel_dims;
        let portfolio_dim: usize = 3;
        let ofi_dim: usize = if state_dim >= market_dim + portfolio_dim + 8 { 8 } else { 0 };

        // Branch sizes for branching DQN (always enabled).
        let branch_sizes: [usize; 3] = [9, 3, 3];
        let total_branch_actions = branch_sizes[0] + branch_sizes[1] + branch_sizes[2]; // 11

        // ── Step 1: Compile experience kernels ──────────────────────────
        let (state_gather_kernel, action_select_kernel, expert_override_kernel, env_step_kernel, expected_q_kernel) =
            compile_experience_kernels(&stream, state_dim, market_dim)?;
        let (nstep_kernel, reward_norm_kernel) = compile_nstep_kernel(&stream)?;
        let fill_episode_ids_kernel = compile_fill_episode_ids_kernel(&stream)?;
        let trade_stats_kernel = compile_trade_stats_kernel(&stream)?;

        info!(
            state_dim,
            market_dim,
            num_atoms_max,
            "GPU experience collector: experience kernels compiled (3 focused + expected_q + episode_ids)"
        );

        // ── Step 2: Create cuBLAS forward context ───────────────────────
        // The experience collector uses num_atoms from the config. For
        // non-distributional mode (num_atoms_max=1), the logit outputs
        // are plain Q-values. For C51 (num_atoms_max>1), we compute
        // expected Q via compute_expected_q kernel.
        let num_atoms = num_atoms_max.max(1);
        let cublas_forward = CublasForward::new(
            &stream,
            alloc_episodes, // batch_size = number of episodes
            state_dim,
            shared_h1,
            shared_h2,
            value_h,
            adv_h,
            num_atoms,
            branch_sizes[0],
            branch_sizes[1],
            branch_sizes[2],
        )?;
        info!(
            alloc_episodes,
            num_atoms,
            "GPU experience collector: cuBLAS forward context created"
        );

        // ── Step 3: Compute flat parameter layout ───────────────────────
        let train_cfg = GpuDqnTrainConfig {
            state_dim,
            shared_h1,
            shared_h2,
            value_h,
            adv_h,
            num_atoms,
            branch_0_size: branch_sizes[0],
            branch_1_size: branch_sizes[1],
            branch_2_size: branch_sizes[2],
            ..GpuDqnTrainConfig::default()
        };
        let param_sizes = compute_param_sizes(&train_cfg);
        let total_params = compute_total_params(&train_cfg);

        // Allocate flat online parameter buffer
        let online_params_flat = stream
            .alloc_zeros::<f32>(total_params)
            .map_err(|e| MLError::ModelError(format!("alloc online_params_flat: {e}")))?;

        info!(
            total_params,
            total_bytes = total_params * 4,
            "GPU experience collector: flat parameter buffer allocated"
        );

        // ── Step 4: Extract legacy weight sets ──────────────────────────
        let online_weights = extract_dueling_weights_branching(online_vars, &stream)?;
        let target_weights = extract_dueling_weights_branching(target_vars, &stream)?;
        let curiosity_weights = match curiosity_vars {
            Some(vars) => extract_curiosity_weights(vars, &stream)?,
            None => {
                info!("GPU experience collector: curiosity disabled, using zero weights");
                CuriosityWeightSet::zeros(&stream)?
            }
        };

        let online_rmsnorm = match extract_rmsnorm_weights(online_vars, &stream)? {
            Some(set) => {
                info!("GPU experience collector: distributional RMSNorm weights extracted (online)");
                set
            }
            None => RmsNormWeightSet::ones(&stream, network_dims)?,
        };
        let target_rmsnorm = match extract_rmsnorm_weights(target_vars, &stream)? {
            Some(set) => {
                info!("GPU experience collector: distributional RMSNorm weights extracted (target)");
                set
            }
            None => RmsNormWeightSet::ones(&stream, network_dims)?,
        };

        info!("GPU experience collector: extracting branching weights (order + urgency heads)");
        let (online_branching, target_branching) = (
            extract_branching_weights(online_vars, &stream)?,
            extract_branching_weights(target_vars, &stream)?,
        );

        // ── Step 5: L2 cache persistence hints (Hopper GPUs) ────────────
        {
            let caps = ml_core::gpu::capabilities::cached_capabilities();
            if ml_core::gpu::l2_cache::supports_l2_persistence(&caps.device_name) {
                let total_weight_bytes = total_params * std::mem::size_of::<f32>();
                match ml_core::gpu::l2_cache::pin_weights_in_l2(&caps.device_name, total_weight_bytes) {
                    Ok(reserved) if reserved > 0 => {
                        info!(
                            weight_bytes = total_weight_bytes,
                            reserved_bytes = reserved,
                            "L2 cache: pinned {} MB model weights ({} MB L2 reserved)",
                            total_weight_bytes / 0x0010_0000,
                            reserved / 0x0010_0000,
                        );
                    }
                    Ok(_) => {}
                    Err(e) => {
                        tracing::debug!("L2 cache persistence setup skipped (non-fatal): {e}");
                    }
                }
            }
        }

        // ── Step 6: Allocate per-episode buffers ────────────────────────
        // Portfolio states for experience kernels: [N, PORTFOLIO_STRIDE=20]
        let mut portfolio_states = stream
            .alloc_zeros::<f32>(alloc_episodes * PORTFOLIO_STRIDE)
            .map_err(|e| MLError::ModelError(format!("alloc portfolio_states: {e}")))?;

        // Initialize portfolio state (20 floats per episode)
        let mut portfolio_init = vec![0.0_f32; alloc_episodes * PORTFOLIO_STRIDE];
        for i in 0..alloc_episodes {
            let off = i * PORTFOLIO_STRIDE;
            portfolio_init[off] = 0.0;                // [0] position
            portfolio_init[off + 1] = initial_capital; // [1] cash
            portfolio_init[off + 2] = initial_capital; // [2] portfolio_value
            // [3] dsr_A = 0.0, [4] dsr_B = 0.0 (already zero)
            // [5] pnl_ema = 0.0, [6] pnl_var = 0.0 (already zero)
            portfolio_init[off + 7] = initial_capital; // [7] peak_equity
            // [8] flat_counter = 0.0 (already zero)
            portfolio_init[off + 9] = initial_capital; // [9] prev_equity
            // [10] hold_time = 0.0, [11] realized_pnl = 0.0 (already zero)
        }
        stream
            .memcpy_htod(&portfolio_init, &mut portfolio_states)
            .map_err(|e| MLError::ModelError(format!("upload portfolio init: {e}")))?;

        let mut rng_states = stream
            .alloc_zeros::<u32>(alloc_episodes)
            .map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?;

        // Initialize RNG seeds
        let rng_seeds: Vec<u32> = (0..alloc_episodes)
            .map(|_| fastrand::u32(..))
            .collect();
        stream
            .memcpy_htod(&rng_seeds, &mut rng_states)
            .map_err(|e| MLError::ModelError(format!("upload RNG seeds: {e}")))?;

        let episode_starts_buf = stream
            .alloc_zeros::<i32>(alloc_episodes)
            .map_err(|e| MLError::ModelError(format!("alloc episode_starts_buf: {e}")))?;

        // Persistent epoch state
        let epoch_state_init: Vec<f32> = vec![
            0.01,            // vol_ema
            0.01,            // median_vol
            initial_capital, // portfolio_value
            0.0,             // portfolio_position
            initial_capital, // portfolio_cash
            0.0,             // dsr_mean
            1.0,             // dsr_var
            0.0,             // step_count
        ];
        let epoch_state = stream.clone_htod(&epoch_state_init)
            .map_err(|e| MLError::ModelError(format!("epoch_state alloc: {e}")))?;

        // ── Step 7: Allocate output buffers ─────────────────────────────
        let total_output = alloc_episodes * alloc_timesteps;
        let states_out = stream
            .alloc_zeros::<f32>(total_output * state_dim)
            .map_err(|e| MLError::ModelError(format!("alloc states_out: {e}")))?;
        let actions_out = stream
            .alloc_zeros::<i32>(total_output)
            .map_err(|e| MLError::ModelError(format!("alloc actions_out: {e}")))?;
        let rewards_out = stream
            .alloc_zeros::<f32>(total_output)
            .map_err(|e| MLError::ModelError(format!("alloc rewards_out: {e}")))?;
        let done_out = stream
            .alloc_zeros::<f32>(total_output)
            .map_err(|e| MLError::ModelError(format!("alloc done_out: {e}")))?;
        let raw_returns_out = stream
            .alloc_zeros::<f32>(total_output)
            .map_err(|e| MLError::ModelError(format!("alloc raw_returns_out: {e}")))?;
        let trade_stats_buf = stream
            .alloc_zeros::<f32>(TRADE_STATS_FLOATS)
            .map_err(|e| MLError::ModelError(format!("alloc trade_stats_buf: {e}")))?;

        // ── Step 8: Allocate per-timestep batch buffers ─────────────────
        let n = alloc_episodes;
        let batch_states = stream
            .alloc_zeros::<f32>(n * state_dim)
            .map_err(|e| MLError::ModelError(format!("alloc batch_states: {e}")))?;
        let batch_actions = stream
            .alloc_zeros::<i32>(n)
            .map_err(|e| MLError::ModelError(format!("alloc batch_actions: {e}")))?;
        let q_gaps_buf = stream
            .alloc_zeros::<f32>(n)
            .map_err(|e| MLError::ModelError(format!("alloc q_gaps_buf: {e}")))?;
        let current_timesteps = stream
            .alloc_zeros::<i32>(n)
            .map_err(|e| MLError::ModelError(format!("alloc current_timesteps: {e}")))?;

        let total_branch_atoms = (branch_sizes[0] + branch_sizes[1] + branch_sizes[2]) * num_atoms;
        let q_values = stream
            .alloc_zeros::<f32>(n * total_branch_actions)
            .map_err(|e| MLError::ModelError(format!("alloc q_values: {e}")))?;

        // cuBLAS activation scratch buffers
        let exp_h_s1 = stream
            .alloc_zeros::<f32>(n * shared_h1)
            .map_err(|e| MLError::ModelError(format!("alloc exp_h_s1: {e}")))?;
        let exp_h_s2 = stream
            .alloc_zeros::<f32>(n * shared_h2)
            .map_err(|e| MLError::ModelError(format!("alloc exp_h_s2: {e}")))?;
        let exp_h_v = stream
            .alloc_zeros::<f32>(n * value_h)
            .map_err(|e| MLError::ModelError(format!("alloc exp_h_v: {e}")))?;
        let exp_h_b0 = stream
            .alloc_zeros::<f32>(n * adv_h)
            .map_err(|e| MLError::ModelError(format!("alloc exp_h_b0: {e}")))?;
        let exp_h_b1 = stream
            .alloc_zeros::<f32>(n * adv_h)
            .map_err(|e| MLError::ModelError(format!("alloc exp_h_b1: {e}")))?;
        let exp_h_b2 = stream
            .alloc_zeros::<f32>(n * adv_h)
            .map_err(|e| MLError::ModelError(format!("alloc exp_h_b2: {e}")))?;
        let exp_v_logits = stream
            .alloc_zeros::<f32>(n * num_atoms)
            .map_err(|e| MLError::ModelError(format!("alloc exp_v_logits: {e}")))?;
        let exp_b_logits = stream
            .alloc_zeros::<f32>(n * total_branch_atoms)
            .map_err(|e| MLError::ModelError(format!("alloc exp_b_logits: {e}")))?;

        // ── Step 9: Initialize curiosity trainer ────────────────────────
        let curiosity_trainer = if curiosity_vars.is_some() {
            let max_samples = alloc_episodes * alloc_timesteps;
            match GpuCuriosityTrainer::new(Arc::clone(&stream), state_dim, max_samples) {
                Ok(trainer) => {
                    info!("GPU curiosity trainer initialized");
                    Some(trainer)
                }
                Err(e) => {
                    return Err(MLError::ModelError(format!(
                        "GPU curiosity trainer init FAILED: {e}"
                    )));
                }
            }
        } else {
            None
        };

        // Pinned host buffers for DtoH transfers
        let pinned_states = unsafe { PinnedHostBuf::<f32>::new(total_output * state_dim)? };
        let pinned_actions = unsafe { PinnedHostBuf::<i32>::new(total_output)? };
        let pinned_rewards = unsafe { PinnedHostBuf::<f32>::new(total_output)? };
        let pinned_dones = unsafe { PinnedHostBuf::<i32>::new(total_output)? };

        // OFI placeholder
        let ofi_placeholder = stream.alloc_zeros::<f32>(1)
            .map_err(|e| MLError::ModelError(format!("alloc OFI placeholder: {e}")))
            .ok();

        // ── Step 10: Log buffer sizes ───────────────────────────────────
        let per_timestep_bytes = n * state_dim * 4     // batch_states
            + n * 4                                      // batch_actions
            + n * 4                                      // current_timesteps
            + n * total_branch_actions * 4               // q_values
            + n * shared_h1 * 4                          // h_s1
            + n * shared_h2 * 4                          // h_s2
            + n * value_h * 4                            // h_v
            + n * adv_h * 4 * 3                          // h_b0 + h_b1 + h_b2
            + n * num_atoms * 4                          // v_logits
            + n * total_branch_atoms * 4;                // b_logits
        let output_bytes = total_output * (state_dim * 4 + 4 + 4 + 4);
        info!(
            alloc_episodes,
            alloc_timesteps,
            per_timestep_scratch_kb = per_timestep_bytes / 1024,
            output_buffer_mb = output_bytes / (1024 * 1024),
            total_params,
            "GPU experience collector initialized (cuBLAS + 3 focused kernels)"
        );

        Ok(Self {
            stream,
            state_dim,
            market_dim,
            network_dims,
            num_atoms,
            // v_min/v_max computed from reward_scale=10, gamma=0.95:
            // Q* ~ 10 / (1-0.95) * 1.2 = 240, clamped to [20, 300] → ±240
            v_min: -(10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            v_max: (10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            branch_sizes,
            alloc_episodes,
            alloc_timesteps,
            cublas_forward,
            online_params_flat,
            total_params,
            param_sizes,
            cvar_scales_ptr: 0, // NULL = no CVaR scaling initially
            state_gather_kernel,
            action_select_kernel,
            expert_override_kernel,
            env_step_kernel,
            expected_q_kernel,
            nstep_kernel,
            _reward_norm_kernel: reward_norm_kernel,
            fill_episode_ids_kernel,
            trade_stats_kernel,
            trade_stats_buf,
            batch_states,
            batch_actions,
            q_gaps_buf,
            current_timesteps,
            q_values,
            exp_h_s1,
            exp_h_s2,
            exp_h_v,
            exp_h_b0,
            exp_h_b1,
            exp_h_b2,
            exp_v_logits,
            exp_b_logits,
            online_weights,
            target_weights,
            curiosity_weights,
            online_rmsnorm,
            target_rmsnorm,
            online_branching,
            target_branching,
            portfolio_states,
            rng_states,
            episode_starts_buf,
            states_out,
            actions_out,
            rewards_out,
            done_out,
            raw_returns_out,
            epoch_state,
            reset_flags: 0,
            expert_actions_gpu: None,
            expert_ratio: 0.0,
            ofi_gpu: ofi_placeholder,
            ofi_dim,
            curiosity_trainer,
            pinned_states,
            pinned_actions,
            pinned_rewards,
            pinned_dones,
        })
    }

    /// Collect experiences and return GPU-resident `CudaSlice` buffers.
    ///
    /// Runs the timestep-level loop: for each timestep, gathers states,
    /// runs cuBLAS Q-forward, selects actions, and steps the environment.
    /// Set CVaR position scaling from IQN dual-head.
    /// The device pointer will be passed to the env_step kernel.
    /// Call with 0 to disable (NULL pointer = no scaling).
    pub fn set_cvar_scales(&mut self, device_ptr: u64) {
        self.cvar_scales_ptr = device_ptr;
    }

    /// Upload pre-computed expert demonstration actions to GPU.
    ///
    /// `expert_actions[bar_index]` = expert exposure action index (-1 = no opinion).
    /// Generated once from MA crossover + ADX filter, stays GPU-resident for the
    /// entire training run. The action selection kernel reads this at each timestep
    /// and overrides Q-network actions with probability `expert_ratio`.
    pub fn upload_expert_actions(&mut self, actions: &[i32]) -> Result<(), MLError> {
        let mut buf = self.stream.alloc_zeros::<i32>(actions.len())
            .map_err(|e| MLError::ModelError(format!("alloc expert_actions: {e}")))?;
        self.stream.memcpy_htod(actions, &mut buf)
            .map_err(|e| MLError::ModelError(format!("upload expert_actions: {e}")))?;
        self.expert_actions_gpu = Some(buf);
        info!(n_actions = actions.len(), "Expert demo actions uploaded to GPU");
        Ok(())
    }

    /// Set the expert demo override ratio for the current epoch.
    pub fn set_expert_ratio(&mut self, ratio: f32) {
        self.expert_ratio = ratio;
    }

    /// Collect per-epoch trade statistics from GPU portfolio_states.
    ///
    /// Launches `trade_stats_reduce` kernel to sum ps[14:19] across all N episodes,
    /// then downloads the 6-float result (24 bytes) plus rewards_out for step returns.
    ///
    /// Called once per epoch at the boundary — cold path, not in hot training loop.
    pub fn collect_trade_stats(&mut self) -> Result<TradeStats, MLError> {
        let n = self.alloc_episodes;
        let n_i32 = n as i32;

        // Zero the output buffer before reduction
        let zeros = vec![0.0_f32; TRADE_STATS_FLOATS];
        self.stream.memcpy_htod(&zeros, &mut self.trade_stats_buf)
            .map_err(|e| MLError::ModelError(format!("zero trade_stats_buf: {e}")))?;

        // Launch reduction kernel: 1 block, 256 threads
        let blocks = 1_u32;
        let threads = 256_u32;

        // Safety: portfolio_states has N * PORTFOLIO_STRIDE allocated floats,
        // trade_stats_buf has TRADE_STATS_FLOATS allocated floats,
        // n_i32 matches the allocation size.
        unsafe {
            self.stream
                .launch_builder(&self.trade_stats_kernel)
                .arg(&self.portfolio_states)
                .arg(&mut self.trade_stats_buf)
                .arg(&n_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (threads, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("trade_stats_reduce launch: {e}")))?;
        }

        // Download 6-float result (24 bytes)
        let mut host_stats = vec![0.0_f32; TRADE_STATS_FLOATS];
        self.stream.memcpy_dtoh(&self.trade_stats_buf, &mut host_stats)
            .map_err(|e| MLError::ModelError(format!("trade_stats DtoH: {e}")))?;

        let win_count = host_stats[0];
        let loss_count = host_stats[1];
        let sum_wins = host_stats[2] as f64;
        let sum_losses = host_stats[3] as f64;
        let sum_returns = host_stats[4] as f64;
        let sum_sq_returns = host_stats[5] as f64;
        let total_trades = (win_count + loss_count) as usize;

        // Download raw_returns_out for step_returns (once per epoch, cold path).
        // raw_returns_out contains TRUE per-bar portfolio returns (unshaped):
        //   (equity_t - equity_{t-1}) / equity_{t-1}
        // NOT the shaped RL reward (which includes dense*0.1 + sparse*2.0 + loss_aversion).
        // Total = alloc_episodes * alloc_timesteps floats (~12KB for 3200 experiences).
        let total_output = self.alloc_episodes * self.alloc_timesteps;
        let mut host_raw_returns = vec![0.0_f32; total_output];
        self.stream.memcpy_dtoh(&self.raw_returns_out, &mut host_raw_returns)
            .map_err(|e| MLError::ModelError(format!("raw_returns_out DtoH: {e}")))?;

        let step_returns: Vec<f64> = host_raw_returns.iter().map(|&r| r as f64).collect();

        Ok(TradeStats {
            total_trades,
            winning_trades: win_count as usize,
            losing_trades: loss_count as usize,
            sum_wins,
            sum_losses,
            sum_returns,
            sum_sq_returns,
            step_returns,
        })
    }

    /// All outputs remain GPU-resident for zero-copy training.
    ///
    /// The caller **must** call `stream().synchronize()` before passing
    /// these buffers to ops that run on a different CUDA stream.
    pub fn collect_experiences_gpu(
        &mut self,
        market_features_buf: &CudaSlice<f32>,
        targets_buf: &CudaSlice<f32>,
        episode_starts: &[i32],
        config: &ExperienceCollectorConfig,
    ) -> Result<GpuExperienceBatch, MLError> {
        let (n_episodes, timesteps) =
            self.launch_timestep_loop(market_features_buf, targets_buf, episode_starts, config)?;
        let total = n_episodes * timesteps;
        let sd = self.state_dim;

        // Clone kernel output buffers so they can be consumed independently
        let states = dtod_clone_f32(&self.stream, &self.states_out, total * sd, "states")?;
        let rewards = dtod_clone_f32(&self.stream, &self.rewards_out, total, "rewards")?;
        let actions = dtod_clone_i32(&self.stream, &self.actions_out, total, "actions")?;
        let dones = dtod_clone_f32(&self.stream, &self.done_out, total, "dones")?;

        // ── N-step return accumulation ───────────────────────────────────
        // Converts 1-step rewards into R_n = sum(gamma^i * r_i) and
        // OR's done flags over n steps. Uses double-buffering to avoid races.
        let n_steps = config.n_steps.max(1);
        if n_steps > 1 {
            let raw_rewards = dtod_clone_f32(&self.stream, &rewards, total, "raw_rewards_nstep")?;
            let raw_dones = dtod_clone_f32(&self.stream, &dones, total, "raw_dones_nstep")?;

            let gamma_f32 = config.gamma;
            let n_steps_i32 = n_steps;
            let l_i32 = timesteps as i32;
            let n_i32 = n_episodes as i32;
            let blocks = ((total + 255) / 256) as u32;

            unsafe {
                self.stream
                    .launch_builder(&self.nstep_kernel)
                    .arg(&raw_rewards)
                    .arg(&raw_dones)
                    .arg(&rewards)
                    .arg(&dones)
                    .arg(&gamma_f32)
                    .arg(&n_steps_i32)
                    .arg(&l_i32)
                    .arg(&n_i32)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("nstep_accumulate_kernel: {e}")))?;
            }

            debug!(n_steps, gamma = config.gamma, "N-step return accumulation applied");
        }

        // Reward normalization: DISABLED for reward v4.
        // Per-bar percentage returns (~0.001) are naturally scaled.
        // Normalizing would scramble the C51 atom distribution.

        // Build next_states on GPU (episode-aware shift by n_steps)
        let next_states = build_next_states_dtod(
            &self.stream, &states, n_episodes, timesteps, sd, n_steps as usize,
        )?;

        // ── Episode ID fill (GPU-native, zero CPU) ────────────────────
        // episode_ids[i] = i / L — enables HER Future and Final strategies.
        // Direct integer arithmetic: no linear scan needed.
        let episode_ids = fill_episode_ids_gpu(
            &self.stream,
            &self.fill_episode_ids_kernel,
            total,
            timesteps,
        )?;

        debug!(
            n_episodes,
            timesteps,
            total_experiences = total,
            "GPU experience collection complete (cuBLAS timestep loop)"
        );

        Ok(GpuExperienceBatch {
            states,
            next_states,
            actions,
            rewards,
            dones,
            episode_ids,
            state_dim: sd,
            n_episodes,
            timesteps,
        })
    }

    /// Timestep-level loop replacing the monolithic kernel launch.
    ///
    /// For each timestep:
    ///   1. `experience_state_gather` — assemble [N, state_dim] batch
    ///   2. `CublasForward::forward_online` — Q-network via cuBLAS SGEMM
    ///   3. `compute_expected_q` — C51 logits → expected Q-values (skip if num_atoms==1)
    ///   4. `experience_action_select` — epsilon-greedy on Q-values
    ///   5. `experience_env_step` — portfolio simulation, reward, done, advance timestep
    fn launch_timestep_loop(
        &mut self,
        market_features_buf: &CudaSlice<f32>,
        targets_buf: &CudaSlice<f32>,
        episode_starts: &[i32],
        config: &ExperienceCollectorConfig,
    ) -> Result<(usize, usize), MLError> {
        // ── Validate buffer sizes ───────────────────────────────────────
        let n_episodes = episode_starts.len();
        if n_episodes > self.alloc_episodes {
            return Err(MLError::ModelError(format!(
                "Episode count {} exceeds allocated buffer size {}",
                n_episodes, self.alloc_episodes,
            )));
        }
        if n_episodes == 0 {
            return Err(MLError::ModelError(
                "episode_starts must not be empty".to_owned(),
            ));
        }
        let timesteps = config.timesteps_per_episode as usize;
        if timesteps > self.alloc_timesteps {
            return Err(MLError::ModelError(format!(
                "timesteps_per_episode {} exceeds allocated buffer size {}",
                timesteps, self.alloc_timesteps,
            )));
        }

        // ── Upload episode starts ───────────────────────────────────────
        self.stream
            .memcpy_htod(episode_starts, &mut self.episode_starts_buf)
            .map_err(|e| MLError::ModelError(format!("upload episode_starts: {e}")))?;

        // ── Initialize current_timesteps to zeros ───────────────────────
        self.stream
            .memset_zeros(&mut self.current_timesteps)
            .map_err(|e| MLError::ModelError(format!("memset current_timesteps: {e}")))?;

        // ── Compute weight pointers for cuBLAS ──────────────────────────
        let w_ptrs = f32_weight_ptrs(&self.online_params_flat, &self.param_sizes, &self.stream);

        let n = n_episodes;
        let blocks_256 = ((n + 255) / 256) as u32;
        let launch_cfg = LaunchConfig {
            grid_dim: (blocks_256, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let total_bars = config.total_bars;
        let sd = self.state_dim as i32;
        let md = self.market_dim as i32;
        let n_i32 = n as i32;

        let b0 = self.branch_sizes[0] as i32;
        let b1 = self.branch_sizes[1] as i32;
        let b2 = self.branch_sizes[2] as i32;

        debug!(
            n_episodes,
            timesteps,
            epsilon = config.epsilon,
            num_atoms = self.num_atoms,
            "Launching timestep loop (cuBLAS forward)"
        );

        for t in 0..timesteps {
            // ── 1. Gather states: [N, state_dim] ────────────────────────
            unsafe {
                self.stream
                    .launch_builder(&self.state_gather_kernel)
                    .arg(market_features_buf)
                    .arg(&self.episode_starts_buf)
                    .arg(&self.current_timesteps)
                    .arg(&self.portfolio_states)
                    .arg(&mut self.batch_states)
                    .arg(&n_i32)
                    .arg(&total_bars)
                    .arg(&sd)
                    .arg(&md)
                    .launch(launch_cfg)
                    .map_err(|e| MLError::ModelError(format!(
                        "experience_state_gather t={t}: {e}"
                    )))?;
            }

            // ── 2. cuBLAS Q-forward: batch_states → logits ──────────────
            self.cublas_forward.forward_online(
                &self.stream,
                &self.batch_states,
                &w_ptrs,
                &self.exp_h_s1,
                &self.exp_h_s2,
                &self.exp_h_v,
                &self.exp_h_b0,
                &self.exp_h_b1,
                &self.exp_h_b2,
                &self.exp_v_logits,
                &self.exp_b_logits,
            )?;

            // ── 3. Convert logits → expected Q-values ───────────────────
            // When num_atoms == 1, b_logits is already [N, 11] Q-values.
            // When num_atoms > 1, we need softmax + expectation over z-support.
            let q_source: &CudaSlice<f32> = if self.num_atoms > 1 {
                let na = self.num_atoms as i32;
                let v_min_f = self.v_min;
                let v_max_f = self.v_max;
                unsafe {
                    self.stream
                        .launch_builder(&self.expected_q_kernel)
                        .arg(&self.exp_v_logits)
                        .arg(&self.exp_b_logits)
                        .arg(&mut self.q_values)
                        .arg(&n_i32)
                        .arg(&na)
                        .arg(&b0)
                        .arg(&b1)
                        .arg(&b2)
                        .arg(&v_min_f)
                        .arg(&v_max_f)
                        .launch(launch_cfg)
                        .map_err(|e| MLError::ModelError(format!(
                            "compute_expected_q t={t}: {e}"
                        )))?;
                }
                &self.q_values
            } else {
                // For num_atoms==1, b_logits [N, 11] already contains Q-values.
                // The value head is a single scalar per sample — for action selection
                // in branching DQN, the advantage branches are sufficient (value is
                // shared across all actions and doesn't affect argmax).
                &self.exp_b_logits
            };

            // ── 4. Action selection: epsilon-greedy + Q-gap conviction filter
            let epsilon = config.epsilon;
            let q_gap = config.q_gap_threshold;
            unsafe {
                self.stream
                    .launch_builder(&self.action_select_kernel)
                    .arg(q_source)
                    .arg(&mut self.batch_actions)
                    .arg(&mut self.rng_states)
                    .arg(&mut self.q_gaps_buf)
                    .arg(&epsilon)
                    .arg(&n_i32)
                    .arg(&b0)
                    .arg(&b1)
                    .arg(&b2)
                    .arg(&q_gap)
                    .launch(launch_cfg)
                    .map_err(|e| MLError::ModelError(format!(
                        "experience_action_select t={t}: {e}"
                    )))?;
            }

            // ── 4b. Expert action override (GPU-native MA crossover + ADX)
            // When expert_ratio > 0, overrides Q-network exposure actions with
            // expert signals. Reads prices and ADX directly from GPU buffers.
            if self.expert_ratio > 0.0 {
                let er = self.expert_ratio;
                let tb = total_bars;
                unsafe {
                    self.stream
                        .launch_builder(&self.expert_override_kernel)
                        .arg(&mut self.batch_actions)
                        .arg(targets_buf)
                        .arg(market_features_buf)
                        .arg(&self.portfolio_states)
                        .arg(&self.episode_starts_buf)
                        .arg(&self.current_timesteps)
                        .arg(&mut self.rng_states)
                        .arg(&er)
                        .arg(&n_i32)
                        .arg(&tb)
                        .arg(&md)
                        .arg(&b0)
                        .arg(&b1)
                        .arg(&b2)
                        .launch(launch_cfg)
                        .map_err(|e| MLError::ModelError(format!(
                            "expert_action_override t={t}: {e}"
                        )))?;
                }
            }

            // ── 5. Environment step (reward v5: trade-aware hybrid) ──────
            let max_pos = config.max_position;
            let tx_cost = config.tx_cost_multiplier;
            let rw_loss_av = config.loss_aversion;
            let l_i32 = timesteps as i32;
            let t_i32 = t as i32;
            unsafe {
                self.stream
                    .launch_builder(&self.env_step_kernel)
                    .arg(targets_buf)
                    .arg(&self.episode_starts_buf)
                    .arg(&mut self.current_timesteps)
                    .arg(&self.batch_actions)
                    .arg(&mut self.portfolio_states)
                    .arg(&mut self.states_out)
                    .arg(&mut self.actions_out)
                    .arg(&mut self.rewards_out)
                    .arg(&mut self.done_out)
                    .arg(&self.batch_states)
                    .arg(&max_pos)
                    .arg(&tx_cost)
                    .arg(&rw_loss_av)
                    .arg(market_features_buf)
                    .arg(&md)
                    .arg(&l_i32)
                    .arg(&n_i32)
                    .arg(&total_bars)
                    .arg(&sd)
                    .arg(&b0)
                    .arg(&b1)
                    .arg(&b2)
                    .arg(&t_i32)
                    .arg(&self.cvar_scales_ptr) // CVaR position scaling (0 = NULL = no scaling)
                    .arg(&self.q_gaps_buf)     // Q-gap conviction scaling
                    .arg(&mut self.raw_returns_out) // Raw portfolio returns (unshaped) for Sharpe/MaxDD
                    .launch(launch_cfg)
                    .map_err(|e| MLError::ModelError(format!(
                        "experience_env_step t={t}: {e}"
                    )))?;
            }
        }

        // Auto-clear reset flags after launch
        self.reset_flags = 0;

        Ok((n_episodes, timesteps))
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Weight synchronization
    // ═══════════════════════════════════════════════════════════════════════

    /// Sync online Q-network weights from a `GpuVarStore` (after training step).
    ///
    /// Extracts per-tensor weights from VarStore into the legacy weight sets,
    /// then flattens them into the contiguous `online_params_flat` buffer for cuBLAS.
    pub fn sync_online_weights(&mut self, vars: &GpuVarStore) -> Result<(), MLError> {
        sync_dueling_weights_branching(vars, &mut self.online_weights, &self.stream)?;
        self.flatten_online_weights()
    }

    /// Sync target Q-network weights from a `GpuVarStore` (after soft update).
    ///
    /// Updates the legacy target weight sets. The target network is not used
    /// by the cuBLAS timestep loop (only online network is used for action
    /// selection during experience collection).
    pub fn sync_target_weights(&mut self, vars: &GpuVarStore) -> Result<(), MLError> {
        sync_dueling_weights_branching(vars, &mut self.target_weights, &self.stream)
    }

    /// Sync online branching DQN extra head weights (branches 1+2).
    pub fn sync_online_branching(&mut self, vars: &GpuVarStore) -> Result<(), MLError> {
        sync_branching_weights(vars, &mut self.online_branching, &self.stream)?;
        self.flatten_online_weights()
    }

    /// Sync target branching DQN extra head weights (branches 1+2).
    pub fn sync_target_branching(&mut self, vars: &GpuVarStore) -> Result<(), MLError> {
        sync_branching_weights(vars, &mut self.target_branching, &self.stream)
    }

    /// Sync curiosity forward model weights from a `GpuVarStore`.
    pub fn sync_curiosity_weights_from(&mut self, vars: &GpuVarStore) -> Result<(), MLError> {
        sync_curiosity_weights(vars, &mut self.curiosity_weights, &self.stream)
    }

    /// Sync online RMSNorm gamma weights (distributional dueling).
    pub fn sync_online_rmsnorm(&mut self, vars: &GpuVarStore) -> Result<(), MLError> {
        sync_rmsnorm_weights(vars, &mut self.online_rmsnorm, &self.stream)
    }

    /// Sync target RMSNorm gamma weights (distributional dueling).
    pub fn sync_target_rmsnorm(&mut self, vars: &GpuVarStore) -> Result<(), MLError> {
        sync_rmsnorm_weights(vars, &mut self.target_rmsnorm, &self.stream)
    }

    /// Sync from a flat parameter buffer (single DtoD copy).
    ///
    /// Use this when the caller already has a contiguous F32 parameter buffer
    /// (e.g. from `GpuDqnTrainer`). Avoids the per-tensor extraction overhead.
    pub fn sync_weights_flat(&mut self, params_buf: &CudaSlice<f32>) -> Result<(), MLError> {
        let num_bytes = self.total_params * std::mem::size_of::<f32>();
        let src = raw_device_ptr(params_buf, &self.stream);
        let dst = raw_device_ptr(&self.online_params_flat, &self.stream);
        dtod_copy(dst, src, num_bytes, &self.stream, 0, "weight_sync_flat")
    }

    /// Flatten the legacy per-tensor weight sets into the contiguous `online_params_flat`.
    ///
    /// Copies each of the 20 weight tensors from the DuelingWeightSet and
    /// BranchingWeightSet into their respective positions in the flat buffer.
    #[allow(unused_assignments)]
    fn flatten_online_weights(&mut self) -> Result<(), MLError> {
        let mut byte_offset: u64 = 0;
        let dst_base = raw_device_ptr(&self.online_params_flat, &self.stream);

        // Helper: copy a single weight tensor into the flat buffer at current offset.
        macro_rules! copy_weight {
            ($slice:expr, $idx:expr) => {{
                let n_elems = self.param_sizes[$idx];
                let n_bytes = n_elems * std::mem::size_of::<f32>();
                let src = raw_device_ptr($slice, &self.stream);
                dtod_copy(dst_base + byte_offset, src, n_bytes, &self.stream, $idx, "flatten")?;
                byte_offset += n_bytes as u64;
            }};
        }

        // Shared trunk: w_s1, b_s1, w_s2, b_s2
        copy_weight!(&self.online_weights.w_s1, 0);
        copy_weight!(&self.online_weights.b_s1, 1);
        copy_weight!(&self.online_weights.w_s2, 2);
        copy_weight!(&self.online_weights.b_s2, 3);

        // Value head: w_v1, b_v1, w_v2, b_v2
        copy_weight!(&self.online_weights.w_v1, 4);
        copy_weight!(&self.online_weights.b_v1, 5);
        copy_weight!(&self.online_weights.w_v2, 6);
        copy_weight!(&self.online_weights.b_v2, 7);

        // Branch 0 (exposure — from DuelingWeightSet): w_b0fc, b_b0fc, w_b0out, b_b0out
        copy_weight!(&self.online_weights.w_a1, 8);
        copy_weight!(&self.online_weights.b_a1, 9);
        copy_weight!(&self.online_weights.w_a2, 10);
        copy_weight!(&self.online_weights.b_a2, 11);

        // Branch 1 (order — from BranchingWeightSet): w_b1fc, b_b1fc, w_b1out, b_b1out
        copy_weight!(&self.online_branching.w_bo1, 12);
        copy_weight!(&self.online_branching.b_bo1, 13);
        copy_weight!(&self.online_branching.w_bo2, 14);
        copy_weight!(&self.online_branching.b_bo2, 15);

        // Branch 2 (urgency — from BranchingWeightSet): w_b2fc, b_b2fc, w_b2out, b_b2out
        copy_weight!(&self.online_branching.w_bu1, 16);
        copy_weight!(&self.online_branching.b_bu1, 17);
        copy_weight!(&self.online_branching.w_bu2, 18);
        copy_weight!(&self.online_branching.b_bu2, 19);

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Episode management
    // ═══════════════════════════════════════════════════════════════════════

    /// Reset all per-episode state for a new training epoch.
    pub fn reset_episodes(
        &mut self,
        initial_capital: f32,
        _avg_spread: f32,
        _cash_reserve_pct: f32,
    ) -> Result<(), MLError> {
        // Reset portfolio states [N, PORTFOLIO_STRIDE=20]
        let mut portfolio_init = vec![0.0_f32; self.alloc_episodes * PORTFOLIO_STRIDE];
        for i in 0..self.alloc_episodes {
            let off = i * PORTFOLIO_STRIDE;
            portfolio_init[off] = 0.0;                // [0] position
            portfolio_init[off + 1] = initial_capital; // [1] cash
            portfolio_init[off + 2] = initial_capital; // [2] portfolio_value
            // [3..6] dsr_A, dsr_B, pnl_ema, pnl_var = 0.0 (already zero)
            portfolio_init[off + 7] = initial_capital; // [7] peak_equity
            // [8] flat_counter = 0.0 (already zero)
            portfolio_init[off + 9] = initial_capital; // [9] prev_equity
            // [10] hold_time = 0.0, [11] realized_pnl = 0.0 (already zero)
        }
        self.stream
            .memcpy_htod(&portfolio_init, &mut self.portfolio_states)
            .map_err(|e| MLError::ModelError(format!("reset portfolio_states: {e}")))?;

        // Fresh RNG seeds
        let rng_seeds: Vec<u32> = (0..self.alloc_episodes)
            .map(|_| fastrand::u32(..))
            .collect();
        self.stream
            .memcpy_htod(&rng_seeds, &mut self.rng_states)
            .map_err(|e| MLError::ModelError(format!("reset rng_states: {e}")))?;

        debug!(
            initial_capital,
            "Episode state reset complete"
        );

        Ok(())
    }

    /// Set epoch-boundary reset flags for next kernel launch.
    /// Bit 0: reset portfolio. Bit 1: reset DSR. Bit 2: reset vol EMA.
    pub fn set_reset_flags(&mut self, flags: u32) {
        self.reset_flags = flags;
    }

    /// Clear reset flags (called automatically after kernel launch).
    pub fn clear_reset_flags(&mut self) {
        self.reset_flags = 0;
    }

    /// Get a reference to the epoch state GPU buffer (for monitoring kernel).
    pub fn epoch_state_gpu(&self) -> &CudaSlice<f32> {
        &self.epoch_state
    }

    /// Read DSR EMA state from GPU after experience collection.
    pub fn read_epoch_dsr_state(&self) -> Result<(f32, f32), MLError> {
        let mut host = vec![0.0_f32; 8];
        self.stream
            .memcpy_dtoh(&self.epoch_state, &mut host)
            .map_err(|e| MLError::ModelError(format!("epoch_state DtoH readback: {e}")))?;
        Ok((host[5], host[6]))
    }

    /// Get a reference to the CUDA stream used by this collector.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Get a reference to the rewards output GPU buffer.
    pub fn rewards_gpu(&self) -> &CudaSlice<f32> {
        &self.rewards_out
    }

    /// Get a reference to the actions output GPU buffer.
    pub fn actions_gpu(&self) -> &CudaSlice<i32> {
        &self.actions_out
    }

    /// Train curiosity forward model directly on GPU using experience data.
    pub fn train_curiosity_gpu(
        &mut self,
        n_episodes: usize,
        timesteps: usize,
    ) -> Result<(), MLError> {
        if let Some(ref mut trainer) = self.curiosity_trainer {
            let total = n_episodes.saturating_mul(timesteps);
            if total < 2 {
                return Ok(());
            }
            trainer.train_on_collector_buffers(
                &mut self.curiosity_weights,
                &self.states_out,
                &self.actions_out,
                total,
            )?;
        }
        Ok(())
    }

    /// Whether the GPU curiosity trainer is active.
    pub fn has_curiosity_trainer(&self) -> bool {
        self.curiosity_trainer.is_some()
    }

    /// Disable curiosity training (e.g. after an async kernel crash).
    pub fn disable_curiosity(&mut self) {
        self.curiosity_trainer = None;
    }

    /// Upload OFI features to GPU for kernel-side state construction.
    pub fn upload_ofi_features(&mut self, ofi_flat: &[f32]) -> Result<(), MLError> {
        if self.ofi_dim == 0 {
            tracing::warn!("upload_ofi_features called but OFI_DIM=0 (state_dim too small)");
            return Ok(());
        }
        if ofi_flat.is_empty() {
            return Ok(());
        }
        let mut gpu_buf = self.stream.alloc_zeros::<f32>(ofi_flat.len())
            .map_err(|e| MLError::ModelError(format!("alloc OFI GPU buffer: {e}")))?;
        self.stream.memcpy_htod(ofi_flat, &mut gpu_buf)
            .map_err(|e| MLError::ModelError(format!("upload OFI features to GPU: {e}")))?;
        info!(
            "OFI features uploaded to GPU: {} bars x {} dims ({:.1} KB)",
            ofi_flat.len() / self.ofi_dim,
            self.ofi_dim,
            (ofi_flat.len() * 4) as f64 / 1024.0,
        );
        self.ofi_gpu = Some(gpu_buf);
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Experience kernel compilation
// ---------------------------------------------------------------------------

/// Compile the 3 focused experience kernels + compute_expected_q from NVRTC source.
///
/// Injects STATE_DIM and MARKET_DIM as #defines before the kernel source
/// (which uses #ifndef guards for fallback defaults).
fn compile_experience_kernels(
    stream: &Arc<CudaStream>,
    state_dim: usize,
    market_dim: usize,
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
    let kernel_src = include_str!("experience_kernels.cu");

    // Inject compile-time constants before the source
    let dim_overrides = format!(
        "#define STATE_DIM {state_dim}\n\
         #define MARKET_DIM {market_dim}\n\
         #define PORTFOLIO_DIM 8\n"
    );
    let full_source = format!("{dim_overrides}\n{kernel_src}");

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        .map_err(|e| MLError::ModelError(format!("experience kernel compilation: {e}")))?;
    let module = context
        .load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("experience kernel module load: {e}")))?;

    let state_gather = module
        .load_function("experience_state_gather")
        .map_err(|e| MLError::ModelError(format!("load experience_state_gather: {e}")))?;
    let action_select = module
        .load_function("experience_action_select")
        .map_err(|e| MLError::ModelError(format!("load experience_action_select: {e}")))?;
    let expert_override = module
        .load_function("expert_action_override")
        .map_err(|e| MLError::ModelError(format!("load expert_action_override: {e}")))?;
    let env_step = module
        .load_function("experience_env_step")
        .map_err(|e| MLError::ModelError(format!("load experience_env_step: {e}")))?;
    let expected_q = module
        .load_function("compute_expected_q")
        .map_err(|e| MLError::ModelError(format!("load compute_expected_q: {e}")))?;

    Ok((state_gather, action_select, expert_override, env_step, expected_q))
}

/// Compile the n-step return accumulation kernel.
///
/// Converts 1-step (s,a,r,s',d) transitions into n-step (s,a,R_n,s_{t+n},d_n)
/// by accumulating discounted rewards and OR-ing done flags over n steps.
fn compile_nstep_kernel(
    stream: &Arc<CudaStream>,
) -> Result<(CudaFunction, CudaFunction), MLError> {
    let kernel_src = include_str!("nstep_kernel.cu");
    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, &context)
        .map_err(|e| MLError::ModelError(format!("nstep_kernel compilation: {e}")))?;
    let module = context
        .load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("nstep_kernel module load: {e}")))?;
    let nstep = module
        .load_function("nstep_accumulate_kernel")
        .map_err(|e| MLError::ModelError(format!("load nstep_accumulate_kernel: {e}")))?;
    let reward_norm = module
        .load_function("reward_normalize_kernel")
        .map_err(|e| MLError::ModelError(format!("load reward_normalize_kernel: {e}")))?;
    // regime_scale_td_errors — loaded from dqn_utility_kernels.cu (trainer module),
    // not from nstep_kernel.cu. The trainer's version uses STATE_DIM compile-time constant.
    Ok((nstep, reward_norm))
}

/// Compile the HER episode boundary kernels (fill_episode_ids, future, final).
fn compile_fill_episode_ids_kernel(
    stream: &Arc<CudaStream>,
) -> Result<CudaFunction, MLError> {
    let kernel_src = include_str!("her_episode_kernel.cu");
    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, &context)
        .map_err(|e| MLError::ModelError(format!("her_episode_kernel compilation: {e}")))?;
    let module = context
        .load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("her_episode_kernel module load: {e}")))?;
    module
        .load_function("fill_episode_ids_kernel")
        .map_err(|e| MLError::ModelError(format!("load fill_episode_ids_kernel: {e}")))
}

/// Compile the trade_stats_reduce kernel that sums portfolio_states[N, 14:19]
/// across all N episodes into a 6-float output.
///
/// Uses the same `compile_ptx_for_device` (nvcc) path as other experience kernels.
fn compile_trade_stats_kernel(
    stream: &Arc<CudaStream>,
) -> Result<CudaFunction, MLError> {
    let kernel_src = r#"
// Trade stats reduction kernel.
// Reads portfolio_states[N, PORTFOLIO_STRIDE=20] and sums fields [14:19]
// (win_count, loss_count, sum_wins, sum_losses, sum_returns, sum_sq_returns)
// across all N episodes into a 6-float output buffer.
// Uses shared-memory warp reduction (same pattern as monitoring_reduce).

#define PORTFOLIO_STRIDE 20

extern "C" __global__ void trade_stats_reduce(
    const float* __restrict__ portfolio_states, // [N * PORTFOLIO_STRIDE]
    float*       out,                           // [6]: sums of ps[14..19]
    int          N
) {
    __shared__ float s_sums[6];

    int tid = threadIdx.x;
    int stride = blockDim.x;

    // Init shared memory
    if (tid == 0) {
        for (int k = 0; k < 6; k++) s_sums[k] = 0.0f;
    }
    __syncthreads();

    // Thread-local accumulators
    float local_sums[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};

    for (int i = tid; i < N; i += stride) {
        int base = i * PORTFOLIO_STRIDE;
        for (int k = 0; k < 6; k++) {
            local_sums[k] += portfolio_states[base + 14 + k];
        }
    }

    // Warp-level reduction before atomic to shared (reduces contention 32x)
    for (int mask = 16; mask >= 1; mask >>= 1) {
        for (int k = 0; k < 6; k++) {
            local_sums[k] += __shfl_xor_sync(0xFFFFFFFF, local_sums[k], mask);
        }
    }
    // Only lane 0 of each warp atomicAdd
    if ((tid & 31) == 0) {
        for (int k = 0; k < 6; k++) {
            atomicAdd(&s_sums[k], local_sums[k]);
        }
    }
    __syncthreads();

    // Thread 0 writes output
    if (tid == 0) {
        for (int k = 0; k < 6; k++) {
            out[k] = s_sums[k];
        }
    }
}
"#;
    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, &context)
        .map_err(|e| MLError::ModelError(format!("trade_stats_reduce compilation: {e}")))?;
    let module = context
        .load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("trade_stats_reduce module load: {e}")))?;
    module
        .load_function("trade_stats_reduce")
        .map_err(|e| MLError::ModelError(format!("load trade_stats_reduce: {e}")))
}

/// Fill episode_ids[i] = i / L on GPU. Zero CPU in hot path.
///
/// Returns a fresh `CudaSlice<i32>` of length `total` with episode IDs.
fn fill_episode_ids_gpu(
    stream: &Arc<CudaStream>,
    kernel: &CudaFunction,
    total: usize,
    timesteps_per_episode: usize,
) -> Result<CudaSlice<i32>, MLError> {
    let mut episode_ids = stream.alloc_zeros::<i32>(total)
        .map_err(|e| MLError::ModelError(format!("alloc episode_ids[{total}]: {e}")))?;

    let l_i32 = timesteps_per_episode as i32;
    let total_i32 = total as i32;
    let blocks = ((total + 255) / 256) as u32;

    // Safety: episode_ids is a valid GPU allocation of `total` i32 elements.
    // l_i32 and total_i32 are non-negative and consistent with the allocation.
    unsafe {
        stream
            .launch_builder(kernel)
            .arg(&mut episode_ids)
            .arg(&l_i32)
            .arg(&total_i32)
            .launch(LaunchConfig {
                grid_dim: (blocks, 1, 1),
                block_dim: (256, 1, 1),
                shared_mem_bytes: 0,
            })
            .map_err(|e| MLError::ModelError(format!("fill_episode_ids_kernel: {e}")))?;
    }

    Ok(episode_ids)
}

// ---------------------------------------------------------------------------
// Pure cudarc DtoD helpers
// ---------------------------------------------------------------------------

/// Allocate a fresh `CudaSlice<f32>` and DtoD-copy `n_elems` floats from `src`.
fn dtod_clone_f32(
    stream: &Arc<CudaStream>,
    src: &CudaSlice<f32>,
    n_elems: usize,
    label: &str,
) -> Result<CudaSlice<f32>, MLError> {
    let mut dst = stream.alloc_zeros::<f32>(n_elems)
        .map_err(|e| MLError::ModelError(format!("alloc {label} f32[{n_elems}]: {e}")))?;
    let src_view = src.slice(..n_elems);
    stream.memcpy_dtod(&src_view, &mut dst)
        .map_err(|e| MLError::ModelError(format!("DtoD {label}: {e}")))?;
    Ok(dst)
}

/// Allocate a fresh `CudaSlice<i32>` and DtoD-copy `n_elems` ints from `src`.
fn dtod_clone_i32(
    stream: &Arc<CudaStream>,
    src: &CudaSlice<i32>,
    n_elems: usize,
    label: &str,
) -> Result<CudaSlice<i32>, MLError> {
    let mut dst = stream.alloc_zeros::<i32>(n_elems)
        .map_err(|e| MLError::ModelError(format!("alloc {label} i32[{n_elems}]: {e}")))?;
    let src_view = src.slice(..n_elems);
    stream.memcpy_dtod(&src_view, &mut dst)
        .map_err(|e| MLError::ModelError(format!("DtoD {label}: {e}")))?;
    Ok(dst)
}

/// Build next_states from states via episode-aware DtoD shift.
///
/// For n-step returns, `shift` is n (not 1): `next_states[t] = states[t+n]`.
/// When `t + shift >= timesteps`, duplicates the last available state.
fn build_next_states_dtod(
    stream: &Arc<CudaStream>,
    states: &CudaSlice<f32>,
    n_episodes: usize,
    timesteps: usize,
    sd: usize,
    shift: usize,
) -> Result<CudaSlice<f32>, MLError> {
    let total = n_episodes * timesteps;
    let total_elems = total * sd;
    let shift = shift.max(1).min(timesteps); // clamp shift to valid range

    if timesteps <= shift {
        return dtod_clone_f32(stream, states, total_elems, "next_states_trivial");
    }

    let mut dst = stream.alloc_zeros::<f32>(total_elems)
        .map_err(|e| MLError::ModelError(format!("alloc next_states f32[{total_elems}]: {e}")))?;

    let shifted_count = timesteps - shift;

    for ep in 0..n_episodes {
        let ep_offset = ep * timesteps * sd;

        // Copy states[ep][shift..timesteps] -> next_states[ep][0..timesteps-shift]
        let src_shifted = states.slice((ep_offset + shift * sd)..(ep_offset + timesteps * sd));
        let mut dst_shifted = dst.slice_mut(ep_offset..(ep_offset + shifted_count * sd));
        stream.memcpy_dtod(&src_shifted, &mut dst_shifted)
            .map_err(|e| MLError::ModelError(format!("DtoD next_states shift ep{ep}: {e}")))?;

        // Fill remaining positions with last available state
        let last_src_offset = ep_offset + (timesteps - 1) * sd;
        let src_last = states.slice(last_src_offset..(last_src_offset + sd));
        for t in shifted_count..timesteps {
            let dst_offset = ep_offset + t * sd;
            let mut dst_t = dst.slice_mut(dst_offset..(dst_offset + sd));
            stream.memcpy_dtod(&src_last, &mut dst_t)
                .map_err(|e| MLError::ModelError(format!("DtoD next_states fill ep{ep} t{t}: {e}")))?;
        }
    }

    Ok(dst)
}

/// Extract raw F32 device pointer (read-only).
fn raw_f32_ptr(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}
