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

//! GPU-accelerated PPO experience collection via the zero-roundtrip kernel.
//!
//! Wraps `ppo_full_experience_kernel` from `ppo_experience_kernel.cu`, managing
//! all GPU buffers, compiling the kernel at runtime via NVRTC, and providing a
//! `collect_experiences()` method that launches the kernel and downloads results.
//!
//! Each kernel launch processes N independent episodes of L timesteps entirely
//! on GPU -- actor forward (stochastic policy), critic forward (value estimation),
//! portfolio simulation, barrier tracking, diversity entropy, curiosity reward,
//! GAE advantage computation, and return calculation -- with zero CPU-GPU
//! roundtrips per timestep.

use std::sync::Arc;

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

use ml_core::nvtx::NvtxRange;
use crate::MLError;
use super::gpu_weights::{
    CuriosityWeightSet, PpoActorWeightSet, PpoCriticWeightSet,
    extract_ppo_actor_weights, extract_ppo_critic_weights, extract_curiosity_weights,
    sync_ppo_actor_weights, sync_ppo_critic_weights, sync_curiosity_weights,
};

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

const STATE_DIM: usize = 54;
const MAX_EPISODES: usize = 256;
const MAX_TIMESTEPS: usize = 1000;
const PORTFOLIO_STATE_SIZE: usize = 8;
const BARRIER_STATE_SIZE: usize = 5;
const DIVERSITY_WINDOW: usize = 100;

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

/// Configuration for the GPU PPO experience collection kernel.
#[derive(Debug, Clone)]
pub struct PpoCollectorConfig {
    /// 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,
    /// Discount factor for return computation
    pub gamma: f32,
    /// GAE lambda for advantage estimation
    pub gae_lambda: 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,
}

impl Default for PpoCollectorConfig {
    fn default() -> Self {
        Self {
            max_position: 1.0,
            episode_length: 500,
            total_bars: 10_000,
            gamma: 0.99,
            gae_lambda: 0.95,
            curiosity_max_reward: 0.1,
            n_episodes: 128,
            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,
        }
    }
}

impl PpoCollectorConfig {
    /// 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
// ---------------------------------------------------------------------------

/// A batch of PPO experiences that stay GPU-resident after kernel launch.
///
/// All tensor data lives in `CudaSlice` buffers on-device. The training loop
/// passes these directly to `PPO::update_gpu()` with zero GPU-CPU-GPU roundtrip.
/// Use the `download_*()` methods only for debugging or checkpoint serialization.
#[allow(missing_debug_implementations)]
pub struct PpoExperienceBatch {
    /// Flattened state vectors `[N * L * STATE_DIM]` -- GPU-resident
    pub states: CudaSlice<f32>,
    /// Selected action indices `[N * L]` -- GPU-resident
    pub actions: CudaSlice<i32>,
    /// Log-probabilities of selected actions `[N * L]` -- GPU-resident
    pub log_probs: CudaSlice<f32>,
    /// GAE advantages `[N * L]` -- GPU-resident
    pub advantages: CudaSlice<f32>,
    /// Discounted returns (value targets) `[N * L]` -- GPU-resident
    pub returns: CudaSlice<f32>,
    /// Episode-done flags `[N * L]` -- GPU-resident
    pub done_flags: CudaSlice<i32>,
    /// Number of episodes in this batch
    pub n_episodes: usize,
    /// Timesteps per episode
    pub timesteps: usize,
    /// State dimension (features per timestep)
    pub state_dim: usize,
    /// CUDA stream for any downstream operations on this batch
    pub stream: Arc<CudaStream>,
}

impl PpoExperienceBatch {
    /// Total number of experience tuples in this batch.
    pub fn total(&self) -> usize {
        self.n_episodes * self.timesteps
    }

    /// Download states to CPU — epoch boundary readback for hyperopt validation.
    pub fn download_states(&self) -> Result<Vec<f32>, MLError> {
        let count = self.total() * self.state_dim;
        let view = self.states.slice(..count);
        let mut host = vec![0.0_f32; count];
        self.stream
            .memcpy_dtoh(&view, &mut host) // gpu-exit: epoch-boundary hyperopt replay
            .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_states: {e}")))?;
        Ok(host)
    }

    /// Download actions to CPU — epoch boundary readback for metrics.
    pub fn download_actions(&self) -> Result<Vec<i32>, MLError> {
        let count = self.total();
        let view = self.actions.slice(..count);
        let mut host = vec![0_i32; count];
        self.stream
            .memcpy_dtoh(&view, &mut host) // gpu-exit: epoch-boundary entropy metrics
            .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_actions: {e}")))?;
        Ok(host)
    }

    /// Download log_probs to CPU — epoch boundary readback for hyperopt replay.
    pub fn download_log_probs(&self) -> Result<Vec<f32>, MLError> {
        let count = self.total();
        let view = self.log_probs.slice(..count);
        let mut host = vec![0.0_f32; count];
        self.stream
            .memcpy_dtoh(&view, &mut host) // gpu-exit: epoch-boundary hyperopt replay
            .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_log_probs: {e}")))?;
        Ok(host)
    }

    /// Download advantages to CPU — epoch boundary readback for metrics.
    pub fn download_advantages(&self) -> Result<Vec<f32>, MLError> {
        let count = self.total();
        let view = self.advantages.slice(..count);
        let mut host = vec![0.0_f32; count];
        self.stream
            .memcpy_dtoh(&view, &mut host) // gpu-exit: epoch-boundary explained variance
            .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_advantages: {e}")))?;
        Ok(host)
    }

    /// Download returns to CPU — epoch boundary readback for metrics.
    pub fn download_returns(&self) -> Result<Vec<f32>, MLError> {
        let count = self.total();
        let view = self.returns.slice(..count);
        let mut host = vec![0.0_f32; count];
        self.stream
            .memcpy_dtoh(&view, &mut host) // gpu-exit: epoch-boundary explained variance
            .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_returns: {e}")))?;
        Ok(host)
    }

    /// Download done_flags to CPU — epoch boundary readback for hyperopt replay.
    pub fn download_done_flags(&self) -> Result<Vec<i32>, MLError> {
        let count = self.total();
        let view = self.done_flags.slice(..count);
        let mut host = vec![0_i32; count];
        self.stream
            .memcpy_dtoh(&view, &mut host) // gpu-exit: epoch-boundary hyperopt replay
            .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_done_flags: {e}")))?;
        Ok(host)
    }
}

// ---------------------------------------------------------------------------
// GpuPpoExperienceCollector
// ---------------------------------------------------------------------------

/// GPU experience collector that launches the full PPO experience kernel.
///
/// Manages all GPU buffers for actor/critic weights, per-episode state, and
/// output arrays. The `collect_experiences()` method runs the entire
/// experience pipeline on GPU with zero CPU-GPU roundtrips per timestep.
#[allow(missing_debug_implementations)]
pub struct GpuPpoExperienceCollector {
    stream: Arc<CudaStream>,
    kernel_func: CudaFunction,

    // Actor + Critic + Curiosity weights on GPU
    actor_weights: PpoActorWeightSet,
    critic_weights: PpoCriticWeightSet,
    curiosity_weights: CuriosityWeightSet,

    // Per-episode state buffers [N, ...]
    portfolio_states: CudaSlice<f32>,  // [MAX_EPISODES * PORTFOLIO_STATE_SIZE]
    barrier_states: CudaSlice<f32>,    // [MAX_EPISODES * BARRIER_STATE_SIZE]
    diversity_windows: CudaSlice<i32>, // [MAX_EPISODES * DIVERSITY_WINDOW]
    diversity_metas: CudaSlice<i32>,   // [MAX_EPISODES * 2]
    barrier_config: CudaSlice<f32>,    // [3]
    rng_states: CudaSlice<u32>,        // [MAX_EPISODES]
    episode_starts_buf: CudaSlice<i32>,// [MAX_EPISODES]

    // Output buffers [MAX_EPISODES * MAX_TIMESTEPS, ...]
    states_out: CudaSlice<f32>,        // [MAX_EPISODES * MAX_TIMESTEPS * STATE_DIM]
    actions_out: CudaSlice<i32>,       // [MAX_EPISODES * MAX_TIMESTEPS]
    log_probs_out: CudaSlice<f32>,     // [MAX_EPISODES * MAX_TIMESTEPS]
    advantages_out: CudaSlice<f32>,    // [MAX_EPISODES * MAX_TIMESTEPS]
    returns_out: CudaSlice<f32>,       // [MAX_EPISODES * MAX_TIMESTEPS]
    dones_out: CudaSlice<i32>,         // [MAX_EPISODES * MAX_TIMESTEPS]

    // Consolidated readback staging: D2D gather then single DtoH per dtype
    readback_f32_staging: CudaSlice<f32>,
    readback_f32_host: Vec<f32>,
    readback_i32_staging: CudaSlice<i32>,
    readback_i32_host: Vec<i32>,

    // Pre-allocated host-side staging for reset_episodes() to avoid per-call Vec allocations.
    // portfolio_init_staging is re-filled with structured data each reset; rng_seed_staging
    // is re-filled with random seeds. Both avoid hitting the allocator on every epoch boundary.
    portfolio_init_staging: Vec<f32>,
    rng_seed_staging: Vec<u32>,
}

impl GpuPpoExperienceCollector {
    /// Create a new GPU PPO experience collector.
    ///
    /// Compiles the CUDA kernel via NVRTC, extracts neural network weights
    /// from the provided `GpuVarStore` objects, and allocates all GPU buffers.
    ///
    /// # Arguments
    /// * `stream` - CUDA stream for all GPU operations
    /// * `actor_vars` - GpuVarStore for the PPO actor (PolicyNetwork)
    /// * `critic_vars` - GpuVarStore for the PPO critic (ValueNetwork)
    /// * `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
    pub fn new(
        stream: Arc<CudaStream>,
        actor_vars: &GpuVarStore,
        critic_vars: &GpuVarStore,
        curiosity_vars: &GpuVarStore,
        initial_capital: f32,
        avg_spread: f32,
        cash_reserve_pct: f32,
    ) -> Result<Self, MLError> {
        // ---- Step 1: Compile and load kernel ----
        let defines = "\
            #define STATE_DIM 48\n\
            #define MARKET_DIM 42\n\
            #define PORTFOLIO_DIM 8\n";
        let common_src = include_str!("common_device_functions.cuh");
        let kernel_src = include_str!("ppo_experience_kernel.cu");
        let full_source = format!("{defines}{common_src}\n{kernel_src}");
        let context = stream.context();
        // Use arch-aware compilation so __CUDA_ARCH__ is defined correctly.
        // On Hopper (sm_90+) this enables TMA bulk async tile loads.
        let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
            .map_err(|e| {
                MLError::ModelError(format!(
                    "CUDA ppo_experience_kernel compilation failed (is NVRTC available?): {e}"
                ))
            })?;
        let module = context.load_module(ptx).map_err(|e| {
            MLError::ModelError(format!("Failed to load ppo_experience_kernel module: {e}"))
        })?;
        let kernel_func = module
            .load_function("ppo_full_experience_kernel")
            .map_err(|e| {
                MLError::ModelError(format!(
                    "Failed to load ppo_full_experience_kernel function: {e}"
                ))
            })?;
        info!("GPU PPO experience collector: CUDA kernel compiled and loaded");

        // ---- Step 2: Extract weights ----
        let actor_weights = extract_ppo_actor_weights(actor_vars, &stream)?;
        let critic_weights = extract_ppo_critic_weights(critic_vars, &stream)?;
        let curiosity_weights = extract_curiosity_weights(curiosity_vars, &stream)?;

        // ---- Step 3: Allocate per-episode buffers ----
        let portfolio_states = stream
            .alloc_zeros::<f32>(MAX_EPISODES * PORTFOLIO_STATE_SIZE)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc portfolio_states: {e}"))
            })?;
        let barrier_states = stream
            .alloc_zeros::<f32>(MAX_EPISODES * BARRIER_STATE_SIZE)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc barrier_states: {e}"))
            })?;
        let diversity_windows = stream
            .alloc_zeros::<i32>(MAX_EPISODES * DIVERSITY_WINDOW)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc diversity_windows: {e}"))
            })?;
        let diversity_metas = stream
            .alloc_zeros::<i32>(MAX_EPISODES * 2)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc diversity_metas: {e}"))
            })?;
        let barrier_config = stream
            .alloc_zeros::<f32>(3)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc barrier_config: {e}"))
            })?;
        let rng_states = stream
            .alloc_zeros::<u32>(MAX_EPISODES)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc rng_states: {e}"))
            })?;
        let episode_starts_buf = stream
            .alloc_zeros::<i32>(MAX_EPISODES)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc episode_starts_buf: {e}"))
            })?;

        // ---- Step 4: Initialize portfolio states ----
        let mut portfolio_init = vec![0.0_f32; MAX_EPISODES * PORTFOLIO_STATE_SIZE];
        for i in 0..MAX_EPISODES {
            let off = i * PORTFOLIO_STATE_SIZE;
            portfolio_init[off] = initial_capital;      // cash
            // portfolio_init[off + 1] = 0.0;           // position (already zero)
            // portfolio_init[off + 2] = 0.0;           // entry_price
            portfolio_init[off + 3] = initial_capital;  // initial_cap
            portfolio_init[off + 4] = avg_spread;       // spread
            // portfolio_init[off + 5] = 0.0;           // last_price
            portfolio_init[off + 6] = cash_reserve_pct; // reserve_pct
            // portfolio_init[off + 7] = 0.0;           // cum_costs
        }
        let mut portfolio_states = portfolio_states;
        stream
            .memcpy_htod(&portfolio_init, &mut portfolio_states)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to upload portfolio init: {e}"))
            })?;

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

        // ---- Step 6: Allocate output buffers ----
        let total_output = MAX_EPISODES * MAX_TIMESTEPS;
        let states_out = stream
            .alloc_zeros::<f32>(total_output * STATE_DIM)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc states_out: {e}"))
            })?;
        let actions_out = stream
            .alloc_zeros::<i32>(total_output)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc actions_out: {e}"))
            })?;
        let log_probs_out = stream
            .alloc_zeros::<f32>(total_output)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc log_probs_out: {e}"))
            })?;
        let advantages_out = stream
            .alloc_zeros::<f32>(total_output)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc advantages_out: {e}"))
            })?;
        let returns_out = stream
            .alloc_zeros::<f32>(total_output)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc returns_out: {e}"))
            })?;
        let dones_out = stream
            .alloc_zeros::<i32>(total_output)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc dones_out: {e}"))
            })?;

        // ---- Step 6b: Allocate consolidated readback staging buffers ----
        // f32 staging: states (N*L*STATE_DIM) + log_probs (N*L) + advantages (N*L) + returns (N*L)
        let f32_staging_len = total_output * STATE_DIM + 3 * total_output;
        let readback_f32_staging = stream
            .alloc_zeros::<f32>(f32_staging_len)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc readback_f32_staging: {e}"))
            })?;
        let readback_f32_host = vec![0.0_f32; f32_staging_len];

        // i32 staging: actions (N*L) + done_flags (N*L)
        let i32_staging_len = 2 * total_output;
        let readback_i32_staging = stream
            .alloc_zeros::<i32>(i32_staging_len)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to alloc readback_i32_staging: {e}"))
            })?;
        let readback_i32_host = vec![0_i32; i32_staging_len];

        // ---- Step 6c: Pre-allocate host staging for reset_episodes() ----
        // Avoids per-epoch Vec allocations for portfolio init and RNG seeds.
        let portfolio_init_staging = vec![0.0_f32; MAX_EPISODES * PORTFOLIO_STATE_SIZE];
        let rng_seed_staging = vec![0_u32; MAX_EPISODES];

        // ---- Step 7: Log buffer sizes ----
        let per_episode_bytes = (PORTFOLIO_STATE_SIZE + BARRIER_STATE_SIZE) * 4
            + DIVERSITY_WINDOW * 4
            + 2 * 4  // diversity_metas
            + 4;     // rng_state
        let output_bytes = total_output * (STATE_DIM * 4 + 4 + 4 + 4 + 4 + 4);
        info!(
            max_episodes = MAX_EPISODES,
            max_timesteps = MAX_TIMESTEPS,
            per_episode_bytes,
            output_buffer_mb = output_bytes / (1024 * 1024),
            "GPU PPO experience collector buffers allocated"
        );

        Ok(Self {
            stream,
            kernel_func,
            actor_weights,
            critic_weights,
            curiosity_weights,
            portfolio_states,
            barrier_states,
            diversity_windows,
            diversity_metas,
            barrier_config,
            rng_states,
            episode_starts_buf,
            states_out,
            actions_out,
            log_probs_out,
            advantages_out,
            returns_out,
            dones_out,
            readback_f32_staging,
            readback_f32_host,
            readback_i32_staging,
            readback_i32_host,
            portfolio_init_staging,
            rng_seed_staging,
        })
    }

    /// Launch the PPO experience collection kernel and download results.
    ///
    /// Runs N independent episodes of L timesteps entirely on GPU. Each thread
    /// handles one episode: actor forward pass (stochastic softmax sampling),
    /// critic forward pass (value estimation), portfolio simulation, barrier
    /// tracking, diversity entropy, curiosity reward, combined reward shaping,
    /// and backward GAE advantage/return computation.
    ///
    /// # Arguments
    /// * `market_features_buf` - Pre-uploaded market features `[total_bars, 51]`
    /// * `targets_buf` - Pre-uploaded targets `[total_bars, 4]`
    /// * `episode_starts` - Global bar index where each episode begins `[N]`
    /// * `config` - Kernel configuration parameters
    pub fn collect_experiences(
        &mut self,
        market_features_buf: &CudaSlice<f32>,
        targets_buf: &CudaSlice<f32>,
        episode_starts: &[i32],
        config: &PpoCollectorConfig,
    ) -> Result<PpoExperienceBatch, MLError> {
        let _nvtx = NvtxRange::new("ppo_experience_collect");
        // ---- Step 1: Validate ----
        let n_episodes = episode_starts.len();
        if n_episodes > MAX_EPISODES {
            return Err(MLError::ModelError(format!(
                "Episode count {} exceeds MAX_EPISODES {}",
                n_episodes, MAX_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 > MAX_TIMESTEPS {
            return Err(MLError::ModelError(format!(
                "timesteps_per_episode {} exceeds MAX_TIMESTEPS {}",
                timesteps, MAX_TIMESTEPS,
            )));
        }

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

        // ---- Step 3: Upload barrier config ----
        let barrier_cfg = [
            config.barrier_profit_mult,
            config.barrier_loss_mult,
            config.barrier_max_bars,
        ];
        self.stream
            .memcpy_htod(&barrier_cfg, &mut self.barrier_config)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to upload barrier_config: {e}"))
            })?;

        // ---- Step 4: Launch config ----
        let n = n_episodes as u32;
        let (grid_x, block_x) = super::optimal_launch_dims(n, 0);
        let grid_dim = (grid_x, 1, 1);
        let block_dim = (block_x, 1, 1);
        let launch_config = LaunchConfig {
            grid_dim,
            block_dim,
            shared_mem_bytes: 0,
        };

        debug!(
            n_episodes,
            timesteps,
            grid_x = grid_dim.0,
            block_x = block_dim.0,
            gamma = config.gamma,
            gae_lambda = config.gae_lambda,
            "Launching ppo_full_experience_kernel"
        );

        // ---- Step 5: Launch kernel ----
        // Scalar args passed as references to stack variables.
        let n_i32 = n_episodes as i32;

        // Safety: kernel parameter order matches ppo_experience_kernel.cu lines 237-303 exactly.
        // All buffer sizes are pre-validated (MAX_EPISODES, MAX_TIMESTEPS).
        // All CudaSlice lifetimes are valid (owned by self).
        unsafe {
            self.stream
                .launch_builder(&self.kernel_func)
                // Market data (3)
                .arg(market_features_buf)
                .arg(targets_buf)
                .arg(&self.episode_starts_buf)
                // Actor weights (6)
                .arg(&self.actor_weights.pw1)
                .arg(&self.actor_weights.pb1)
                .arg(&self.actor_weights.pw2)
                .arg(&self.actor_weights.pb2)
                .arg(&self.actor_weights.pw3)
                .arg(&self.actor_weights.pb3)
                // Critic weights (12)
                .arg(&self.critic_weights.vw1)
                .arg(&self.critic_weights.vb1)
                .arg(&self.critic_weights.vw2)
                .arg(&self.critic_weights.vb2)
                .arg(&self.critic_weights.vw3)
                .arg(&self.critic_weights.vb3)
                .arg(&self.critic_weights.vw4)
                .arg(&self.critic_weights.vb4)
                .arg(&self.critic_weights.vw5)
                .arg(&self.critic_weights.vb5)
                .arg(&self.critic_weights.vw6)
                .arg(&self.critic_weights.vb6)
                // Curiosity weights (4)
                .arg(&self.curiosity_weights.w1)
                .arg(&self.curiosity_weights.b1)
                .arg(&self.curiosity_weights.w2)
                .arg(&self.curiosity_weights.b2)
                // Per-episode mutable state (4)
                .arg(&mut self.portfolio_states)
                .arg(&mut self.barrier_states)
                .arg(&mut self.diversity_windows)
                .arg(&mut self.diversity_metas)
                // Barrier config
                .arg(&self.barrier_config)
                // Scalars (12)
                .arg(&config.max_position)
                .arg(&config.episode_length)
                .arg(&config.total_bars)
                .arg(&config.timesteps_per_episode)
                .arg(&config.gamma)
                .arg(&config.gae_lambda)
                .arg(&config.curiosity_max_reward)
                .arg(&n_i32)
                .arg(&config.barrier_scale)
                .arg(&config.diversity_scale)
                .arg(&config.curiosity_scale)
                .arg(&config.risk_weight)
                // RNG
                .arg(&mut self.rng_states)
                // Outputs (6)
                .arg(&mut self.states_out)
                .arg(&mut self.actions_out)
                .arg(&mut self.log_probs_out)
                .arg(&mut self.advantages_out)
                .arg(&mut self.returns_out)
                .arg(&mut self.dones_out)
                .launch(launch_config)
                .map_err(|e| {
                    MLError::ModelError(format!(
                        "ppo_full_experience_kernel launch failed: {e}"
                    ))
                })?;
        }

        // ---- Step 6: GPU-resident batch (D2D copy into owned CudaSlices) ----
        // ZERO DtoH transfers. Data stays on GPU for direct consumption by update_gpu().
        let total = n_episodes * timesteps;
        let states_count = total * STATE_DIM;

        // Allocate owned CudaSlices and D2D copy the active sub-range from output buffers.
        // Output buffers are MAX_EPISODES * MAX_TIMESTEPS sized; we only need n * L.
        let batch_states = self.d2d_copy_f32(&self.states_out, states_count, "states")?;
        let batch_log_probs = self.d2d_copy_f32(&self.log_probs_out, total, "log_probs")?;
        let batch_advantages = self.d2d_copy_f32(&self.advantages_out, total, "advantages")?;
        let batch_returns = self.d2d_copy_f32(&self.returns_out, total, "returns")?;
        let batch_actions = self.d2d_copy_i32(&self.actions_out, total, "actions")?;
        let batch_done_flags = self.d2d_copy_i32(&self.dones_out, total, "done_flags")?;

        debug!(
            n_episodes,
            timesteps,
            total_experiences = total,
            "PPO experience collection complete (GPU-resident, 0 DtoH)"
        );

        // ---- Step 7: Return GPU-resident batch ----
        Ok(PpoExperienceBatch {
            states: batch_states,
            actions: batch_actions,
            log_probs: batch_log_probs,
            advantages: batch_advantages,
            returns: batch_returns,
            done_flags: batch_done_flags,
            n_episodes,
            timesteps,
            state_dim: STATE_DIM,
            stream: self.stream.clone(),
        })
    }

    /// D2D copy `count` f32 elements from `src` into a fresh CudaSlice (GPU-to-GPU, no CPU).
    fn d2d_copy_f32(
        &self,
        src: &CudaSlice<f32>,
        count: usize,
        label: &str,
    ) -> Result<CudaSlice<f32>, MLError> {
        let mut dst = self
            .stream
            .alloc_zeros::<f32>(count)
            .map_err(|e| MLError::ModelError(format!("alloc {label}: {e}")))?;
        let src_view = src.slice(..count);
        let nbytes = count * std::mem::size_of::<f32>();
        {
            let (src_ptr, _s) = src_view.device_ptr(&self.stream);
            let (dst_ptr, _d) = dst.device_ptr_mut(&self.stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, self.stream.cu_stream(),
                )
                .map_err(|e| MLError::ModelError(format!("D2D {label}: {e}")))?;
            }
        }
        Ok(dst)
    }

    /// D2D copy `count` i32 elements from `src` into a fresh CudaSlice (GPU-to-GPU, no CPU).
    fn d2d_copy_i32(
        &self,
        src: &CudaSlice<i32>,
        count: usize,
        label: &str,
    ) -> Result<CudaSlice<i32>, MLError> {
        let mut dst = self
            .stream
            .alloc_zeros::<i32>(count)
            .map_err(|e| MLError::ModelError(format!("alloc {label}: {e}")))?;
        let src_view = src.slice(..count);
        let nbytes = count * std::mem::size_of::<i32>();
        {
            let (src_ptr, _s) = src_view.device_ptr(&self.stream);
            let (dst_ptr, _d) = dst.device_ptr_mut(&self.stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, self.stream.cu_stream(),
                )
                .map_err(|e| MLError::ModelError(format!("D2D {label}: {e}")))?;
            }
        }
        Ok(dst)
    }

    /// Re-upload actor and critic weights from `GpuVarStore` objects (after training step).
    pub fn sync_weights(
        &mut self,
        actor_vars: &GpuVarStore,
        critic_vars: &GpuVarStore,
    ) -> Result<(), MLError> {
        sync_ppo_actor_weights(actor_vars, &mut self.actor_weights, &self.stream)?;
        sync_ppo_critic_weights(critic_vars, &mut self.critic_weights, &self.stream)?;
        Ok(())
    }

    /// Re-upload actor, critic, and curiosity weights from `GpuVarStore` objects.
    pub fn sync_all_weights(
        &mut self,
        actor_vars: &GpuVarStore,
        critic_vars: &GpuVarStore,
        curiosity_vars: &GpuVarStore,
    ) -> Result<(), MLError> {
        sync_ppo_actor_weights(actor_vars, &mut self.actor_weights, &self.stream)?;
        sync_ppo_critic_weights(critic_vars, &mut self.critic_weights, &self.stream)?;
        sync_curiosity_weights(curiosity_vars, &mut self.curiosity_weights, &self.stream)?;
        Ok(())
    }

    /// Reset all per-episode state for a new training epoch.
    ///
    /// Re-initializes portfolio states to starting capital, zeros barrier
    /// states and diversity windows, and regenerates RNG seeds.
    ///
    /// Zero-fill buffers use `memset_zeros` (async GPU memset) instead of
    /// allocating CPU-side `Vec`s and uploading via HtoD. Structured data
    /// (portfolio init, RNG seeds) reuses pre-allocated host staging vectors
    /// to avoid per-epoch allocations.
    pub fn reset_episodes(
        &mut self,
        initial_capital: f32,
        avg_spread: f32,
        cash_reserve_pct: f32,
    ) -> Result<(), MLError> {
        // Reset portfolio states using pre-allocated staging buffer
        // (structured init -- cannot use memset_zeros)
        for i in 0..MAX_EPISODES {
            let off = i * PORTFOLIO_STATE_SIZE;
            // Zero all fields first
            for j in 0..PORTFOLIO_STATE_SIZE {
                if let Some(slot) = self.portfolio_init_staging.get_mut(off + j) {
                    *slot = 0.0;
                }
            }
            // Set non-zero fields
            if let Some(slot) = self.portfolio_init_staging.get_mut(off) {
                *slot = initial_capital;       // cash
            }
            if let Some(slot) = self.portfolio_init_staging.get_mut(off + 3) {
                *slot = initial_capital;       // initial_cap
            }
            if let Some(slot) = self.portfolio_init_staging.get_mut(off + 4) {
                *slot = avg_spread;            // spread
            }
            if let Some(slot) = self.portfolio_init_staging.get_mut(off + 6) {
                *slot = cash_reserve_pct;      // reserve_pct
            }
        }
        self.stream
            .memcpy_htod(&self.portfolio_init_staging, &mut self.portfolio_states)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to reset portfolio_states: {e}"))
            })?;

        // Zero barrier states via async GPU memset (no CPU allocation)
        self.stream
            .memset_zeros(&mut self.barrier_states)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to reset barrier_states: {e}"))
            })?;

        // Zero diversity windows via async GPU memset (no CPU allocation)
        self.stream
            .memset_zeros(&mut self.diversity_windows)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to reset diversity_windows: {e}"))
            })?;

        // Zero diversity metas via async GPU memset (no CPU allocation)
        self.stream
            .memset_zeros(&mut self.diversity_metas)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to reset diversity_metas: {e}"))
            })?;

        // Fresh RNG seeds using pre-allocated staging buffer
        for slot in &mut self.rng_seed_staging {
            *slot = fastrand::u32(..);
        }
        self.stream
            .memcpy_htod(&self.rng_seed_staging, &mut self.rng_states)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to reset rng_states: {e}"))
            })?;

        debug!(
            initial_capital,
            avg_spread,
            cash_reserve_pct,
            "PPO episode state reset complete (memset_zeros path)"
        );

        Ok(())
    }
}
