#![allow(unsafe_code)] // Required for f32->u32 CudaSlice reinterpret cast in PER index path
//! Fused CUDA Training Module
//!
//! High-performance H100-optimized training path that replaces 2,100+ Candle kernel
//! dispatches per batch with 3 fused CUDA kernels captured in a CUDA Graph:
//!
//! 1. **Forward + Loss kernel** -- shared trunk, 3 branching advantage heads, C51 distributional loss
//! 2. **Backward kernel** -- full backprop through all layers with gradient clipping
//! 3. **Adam optimizer kernel** -- fused parameter update with weight decay
//!
//! After the first batch, all 3 kernels are captured into a CUDA Graph and replayed
//! on subsequent steps with zero launch overhead.
//!
//! ## Weight Sync Architecture (GPU-native EMA)
//!
//! MlDevice pointers in `CudaSlice` buffers are stable across in-place updates, so
//! the CUDA Graph remains valid across Polyak EMA target updates:
//!
//! 1. **After fused step**: online `CudaSlice` weights are already updated by the Adam kernel
//! 2. **GPU Polyak EMA**: `target[i] = (1-tau)*target[i] + tau*online[i]` via EMA kernel
//! 3. **No reverse/forward sync per step** -- GpuVarStore sync deferred to epoch boundary
//!
//! This eliminates ~120 D2D copies + 120 Candle flatten/contiguous ops per step
//! (~13 GB wasted PCIe per epoch).

use std::sync::Arc;

use anyhow::Result;
use ml_core::device::MlDevice;
use ml_core::cuda_autograd::GpuTensor; // Tensor: HER batch relabeling ops
use tracing::info;

use cudarc::driver::DevicePtr;
use crate::cuda_pipeline::gpu_attention::{GpuAttention, GpuAttentionConfig};
use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer};
use crate::cuda_pipeline::gpu_her::{GpuHer, GpuHerConfig, parse_her_strategy};
use crate::cuda_pipeline::gpu_iql_trainer::{GpuIqlConfig, GpuIqlTrainer};
use crate::cuda_pipeline::gpu_iqn_head::{GpuIqnConfig, GpuIqnHead};
use crate::cuda_pipeline::gpu_weights::{
    self, BranchingWeightSet, DuelingWeightSet,
};
use crate::dqn::dqn::GpuTrainResult;
use crate::dqn::replay_buffer_type::{BatchSample, GpuBatch};
use super::config::DQNAgentType;
use super::DQNHyperparameters;

/// Fused CUDA training context -- owns the `GpuDqnTrainer` and extracted weight sets.
///
/// Weight sets are extracted from the Candle GpuVarStore at initialization via direct
/// DtoD copy (no Candle flatten/cast/contiguous -- Vars are guaranteed F32
/// contiguous by `BranchingDuelingQNetwork::ensure_f32_contiguous()`).
///
/// Per-step training operates entirely on GPU `CudaSlice` buffers: the Adam
/// kernel updates online weights, then the EMA kernel blends online into target
/// weights in-place. The `CudaSlice` weights are the source of truth during
/// training -- the GpuVarStore is never accessed per-step.
///
/// GpuVarStore reverse sync is deferred to epoch boundary for checkpointing and
/// Candle forward (experience collector, inference). MlDevice pointers are stable
/// across in-place updates, so the CUDA Graph stays valid.
pub(crate) struct FusedTrainingCtx {
    trainer: GpuDqnTrainer,
    online_dueling: DuelingWeightSet,
    online_branching: BranchingWeightSet,
    target_dueling: DuelingWeightSet,
    target_branching: BranchingWeightSet,
    stream: Arc<cudarc::driver::CudaStream>,
    /// Batch size at creation time -- must match `current_batch_size` to reuse CUDA Graph.
    batch_size: usize,
    /// Steps since last GpuVarStore sync (deferred to epoch boundary).
    steps_since_varmap_sync: usize,
    /// GPU HER relabeler -- initialized when `her_ratio > 0.0`.
    /// When active, `run_full_step()` splits the batch into normal + HER portions,
    /// runs the relabel kernel on the HER portion, and merges before training.
    pub(crate) gpu_her: Option<GpuHer>,
    /// GPU IQL value trainer -- initialized when `use_iql == true`.
    /// Trains V(s) with expectile regression after each DQN training step.
    /// The value loss is logged via tracing for monitoring.
    pub(crate) gpu_iql: Option<GpuIqlTrainer>,
    /// Trains quantile head that shares the DQN trunk (h_s2) independently.
    /// Combined loss: L_total = L_c51 + iqn_lambda × L_iqn.
    pub(crate) gpu_iqn: Option<GpuIqnHead>,
    /// CVaR scales buffer (kept alive so device pointer remains valid).
    cvar_scales_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// GPU multi-head feature attention over h_s2 (post-graph, 1-step lag).
    /// When Some, applied after EMA update, before IQN, each training step.
    pub(crate) gpu_attention: Option<GpuAttention>,
    /// Ensemble extra heads (heads 1..K-1). Head 0 lives inside the CUDA Graph.
    /// Each element is an independent (DuelingWeightSet, BranchingWeightSet) pair
    /// that shares the same trunk but has perturbed value/advantage weights.
    /// Empty when ensemble_count <= 1.
    pub(crate) ensemble_extra_heads: Vec<(DuelingWeightSet, BranchingWeightSet)>,
    /// Ensemble diversity weight λ — scales the KL loss relative to C51 loss.
    /// Default 0.01. Zero when ensemble_count <= 1.
    pub(crate) ensemble_diversity_weight: f32,
    /// Compiled ensemble aggregate kernel (Q-value mean/variance across heads).
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_aggregate_kernel: Option<cudarc::driver::CudaFunction>,
    /// Compiled ensemble diversity kernel (pairwise KL divergence across heads).
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_diversity_kernel: Option<cudarc::driver::CudaFunction>,
    /// Pre-allocated buffer: [K * B * num_atoms] for per-head value logits.
    /// Used to assemble logits from all K heads before diversity kernel.
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_logits_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// Pre-allocated buffer: [B * total_actions] for mean Q-values across K heads.
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_mean_q_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// Pre-allocated buffer: [B * total_actions] for Q-value variance across K heads.
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_var_q_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// Pre-allocated buffer: [1] for accumulated diversity loss scalar.
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_diversity_loss_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// Compiled KL gradient kernel for diversity gradient flow.
    pub(crate) ensemble_kl_grad_kernel: Option<cudarc::driver::CudaFunction>,
    /// Pre-allocated buffer: [B * num_atoms] for diversity gradient on head 0 logits.
    pub(crate) ensemble_d_logits_buf: Option<cudarc::driver::CudaSlice<f32>>,
}

impl Drop for FusedTrainingCtx {
    fn drop(&mut self) {
        // Sync stream before fields drop to prevent stale CUDA events.
        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
        let _ = self.stream.context().check_err();
    }
}

impl FusedTrainingCtx {
    /// Create a new fused training context.
    ///
    /// Extracts online + target weight sets from the DQN's Candle VarMaps into
    /// `CudaSlice` buffers, compiles the 3 fused kernels + EMA kernel, and
    /// allocates all pre-allocated buffers for CUDA Graph capture.
    ///
    /// Only valid when `matches!(device, MlDevice::Cuda { .. }) && agent.is_using_branching()`.
    pub(crate) fn new(
        _device: &MlDevice,
        agent: &DQNAgentType,
        hyperparams: &DQNHyperparameters,
        batch_size: usize,
        stream: Arc<cudarc::driver::CudaStream>,
    ) -> Result<Self> {
        let dqn = match agent {
            DQNAgentType::Standard(d) => d,
            DQNAgentType::RegimeConditional(regime) => regime.primary_head(),
        };

        let branching_net = dqn.branching_q_network.as_ref().ok_or_else(|| {
            anyhow::anyhow!("Fused CUDA training requires branching Q-network")
        })?;
        let branching_target = dqn.branching_target_network.as_ref().ok_or_else(|| {
            anyhow::anyhow!("Fused CUDA training requires branching target network")
        })?;

        // Use the caller-provided forked CudaStream. All GPU components (experience
        // collector, fused trainer, portfolio sim, monitoring) share this single
        // stream via Arc::clone — no more dual-stream event tracking conflicts.

        // Build config from DQN network dimensions
        let (shared_h1, shared_h2, value_h, adv_h) = agent.network_dims();
        let config = GpuDqnTrainConfig {
            state_dim: dqn.config.state_dim,
            shared_h1,
            shared_h2,
            value_h,
            adv_h,
            num_atoms: dqn.config.num_atoms,
            v_min: dqn.config.v_min,
            v_max: dqn.config.v_max,
            branch_0_size: dqn.config.num_actions,
            branch_1_size: dqn.config.num_order_types,
            branch_2_size: dqn.config.num_urgency_levels,
            batch_size,
            gamma: hyperparams.gamma as f32,
            n_steps: hyperparams.n_steps,
            lr: hyperparams.learning_rate as f32,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            weight_decay: 1e-5,
            // Scale gradient clip by (1 + iqn_lambda) to account for combined gradient:
            // grad_buf = C51_grad + iqn_lambda * IQN_grad. Without scaling, the clip
            // is too tight and Adam sees inconsistent clipping, causing loss oscillation.
            max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(10.0) as f32
                * (1.0 + hyperparams.iqn_lambda as f32),
            spectral_norm_sigma_max: hyperparams.spectral_norm_sigma_max as f32,
            iqn_lambda: hyperparams.iqn_lambda as f32,
            iqn_num_quantiles: if hyperparams.iqn_lambda > 0.0 { hyperparams.num_quantiles } else { 0 },
            iqn_embedding_dim: dqn.config.iqn_embedding_dim,
            iqn_kappa: dqn.config.iqn_kappa,
            entropy_coefficient: dqn.config.entropy_coefficient as f32,
            c51_warmup_epochs: hyperparams.c51_warmup_epochs,
            cql_alpha: hyperparams.cql_alpha as f32,
        };

        // Extract weight sets from VarMaps (online + target)
        let online_vars = branching_net.vars();
        let target_vars = branching_target.vars();

        let online_dueling =
            gpu_weights::extract_dueling_weights_branching(online_vars, &stream)
                .map_err(|e| anyhow::anyhow!("Extract online dueling weights: {e}"))?;
        let online_branching =
            gpu_weights::extract_branching_weights(online_vars, &stream)
                .map_err(|e| anyhow::anyhow!("Extract online branching weights: {e}"))?;
        let target_dueling =
            gpu_weights::extract_dueling_weights_branching(target_vars, &stream)
                .map_err(|e| anyhow::anyhow!("Extract target dueling weights: {e}"))?;
        let target_branching =
            gpu_weights::extract_branching_weights(target_vars, &stream)
                .map_err(|e| anyhow::anyhow!("Extract target branching weights: {e}"))?;

        // Create the fused trainer (compiles kernels, allocates buffers)
        let trainer = GpuDqnTrainer::new(stream.clone(), config)
            .map_err(|e| anyhow::anyhow!("GpuDqnTrainer init: {e}"))?;

        // Initialize GPU HER if her_ratio > 0
        let gpu_her = if hyperparams.her_ratio > 0.0 {
            let her_config = GpuHerConfig {
                goal_dim: 1, // Single scalar goal (target return)
                her_ratio: hyperparams.her_ratio,
                strategy: parse_her_strategy(&hyperparams.her_strategy),
                goal_threshold: 0.01,
                batch_size,
                state_dim: dqn.config.state_dim,
            };
            match GpuHer::new(stream.clone(), her_config) {
                Ok(her) => {
                    info!(
                        her_ratio = hyperparams.her_ratio,
                        her_strategy = %hyperparams.her_strategy,
                        "GPU HER initialized: relabel kernel compiled"
                    );
                    Some(her)
                }
                Err(e) => {
                    tracing::warn!("GPU HER init failed, continuing without HER: {e}");
                    None
                }
            }
        } else {
            None
        };

        // Initialize GPU IQL value trainer if use_iql is enabled
        let gpu_iql = if hyperparams.use_iql {
            let iql_config = GpuIqlConfig {
                expectile_tau: hyperparams.iql_expectile_tau,
                advantage_temperature: hyperparams.iql_advantage_temperature,
                value_hidden_dim: 128,
                state_dim: dqn.config.state_dim,
                batch_size,
                lr: hyperparams.learning_rate as f32,
                max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32,
                ..GpuIqlConfig::default()
            };
            match GpuIqlTrainer::new(stream.clone(), iql_config) {
                Ok(iql) => {
                    info!(
                        expectile_tau = hyperparams.iql_expectile_tau,
                        advantage_temperature = hyperparams.iql_advantage_temperature,
                        "GPU IQL initialized: V(s) value network kernels compiled"
                    );
                    Some(iql)
                }
                Err(e) => {
                    tracing::warn!("GPU IQL init failed, continuing without IQL: {e}");
                    None
                }
            }
        } else {
            None
        };

        // Initialize GPU IQN dual-head when iqn_lambda > 0 (CVaR risk sizing)
        let gpu_iqn = if hyperparams.iqn_lambda > 0.0 {
            let iqn_config = GpuIqnConfig {
                hidden_dim: shared_h2,
                embed_dim: dqn.config.iqn_embedding_dim,
                num_quantiles: hyperparams.num_quantiles,
                kappa: dqn.config.iqn_kappa,
                state_dim: dqn.config.state_dim,
                shared_h1,
                branch_0_size: dqn.config.num_actions,
                branch_1_size: dqn.config.num_order_types,
                branch_2_size: dqn.config.num_urgency_levels,
                batch_size,
                lr: hyperparams.learning_rate as f32,
                max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32,
                gamma: hyperparams.gamma as f32,
                ..GpuIqnConfig::default()
            };
            match GpuIqnHead::new(stream.clone(), iqn_config) {
                Ok(iqn) => {
                    info!(
                        iqn_lambda = hyperparams.iqn_lambda,
                        num_quantiles = hyperparams.num_quantiles,
                        embed_dim = dqn.config.iqn_embedding_dim,
                        "GPU IQN dual-head initialized: 7 kernels compiled"
                    );
                    Some(iqn)
                }
                Err(e) => {
                    tracing::warn!("GPU IQN init failed, continuing without IQN: {e}");
                    None
                }
            }
        } else {
            None
        };

        // GPU attention — always active. 4-head self-attention on h_s2.
        let attn_config = GpuAttentionConfig {
            state_dim: shared_h2,
            num_heads: 4,
            batch_size,
        };
        let gpu_attention = match GpuAttention::new(stream.clone(), attn_config) {
            Ok(attn) => {
                info!(
                    hidden_dim = shared_h2,
                    num_heads = 4,
                    batch_size,
                    "GPU attention initialized: 4-head feature attention on h_s2"
                );
                Some(attn)
            }
            Err(e) => {
                tracing::warn!("GPU attention init failed, continuing without: {e}");
                None
            }
        };

        // Initialize ensemble extra heads (heads 1..K-1) when ensemble_count > 1.
        // Head 0 is head_0 = online_dueling + online_branching (already inside CUDA Graph).
        // Heads 1..K-1 run OUTSIDE the CUDA Graph and share the trunk activations (h_s2).
        let k = hyperparams.ensemble_count.max(1);
        let (
            ensemble_extra_heads,
            ensemble_aggregate_kernel,
            ensemble_diversity_kernel,
            ensemble_logits_buf,
            ensemble_mean_q_buf,
            ensemble_var_q_buf,
            ensemble_diversity_loss_buf,
            ensemble_kl_grad_kernel,
            ensemble_d_logits_buf,
        ) = if k > 1 {
            use crate::cuda_pipeline::gpu_dqn_trainer::compile_ensemble_kernels;

            let (agg_kernel, div_kernel) =
                compile_ensemble_kernels(&stream, dqn.config.state_dim)
                    .map_err(|e| anyhow::anyhow!("Ensemble kernels compile: {e}"))?;

            // Allocate K-1 extra head weight sets via DtoD clone of head 0 + small noise.
            // DtoD: head_k ← online_dueling + online_branching weights copied on GPU.
            let mut extra_heads: Vec<(DuelingWeightSet, BranchingWeightSet)> =
                Vec::with_capacity(k - 1);
            for head_idx in 1..k {
                let head_dueling = clone_dueling_weights(&online_dueling, &stream, head_idx)
                    .map_err(|e| anyhow::anyhow!("Clone ensemble dueling head {head_idx}: {e}"))?;
                let head_branching = clone_branching_weights(&online_branching, &stream, head_idx)
                    .map_err(|e| anyhow::anyhow!("Clone ensemble branching head {head_idx}: {e}"))?;
                extra_heads.push((head_dueling, head_branching));
            }

            // Pre-allocate ensemble buffers.
            // ensemble_logits_buf: [K * B * num_atoms] — all K heads' value logits
            let na = dqn.config.num_atoms;
            let total_actions = dqn.config.num_actions + dqn.config.num_order_types + dqn.config.num_urgency_levels;
            let logits_buf = stream.alloc_zeros::<f32>(k * batch_size * na)
                .map_err(|e| anyhow::anyhow!("Alloc ensemble_logits_buf: {e}"))?;
            let mean_q_buf = stream.alloc_zeros::<f32>(batch_size * total_actions)
                .map_err(|e| anyhow::anyhow!("Alloc ensemble_mean_q_buf: {e}"))?;
            let var_q_buf = stream.alloc_zeros::<f32>(batch_size * total_actions)
                .map_err(|e| anyhow::anyhow!("Alloc ensemble_var_q_buf: {e}"))?;
            let div_loss_buf = stream.alloc_zeros::<f32>(1)
                .map_err(|e| anyhow::anyhow!("Alloc ensemble_diversity_loss_buf: {e}"))?;

            info!(
                ensemble_count = k,
                extra_heads = k - 1,
                "Ensemble multi-head initialized: {k} heads, KL diversity weight={diversity_weight}",
                diversity_weight = hyperparams.ensemble_diversity_weight,
            );

            // Compile KL gradient kernel for diversity gradient flow
            let kl_grad_kernel = {
                let module = stream.context().load_module(
                    crate::cuda_pipeline::compile_ptx_for_device(
                        include_str!("../../cuda_pipeline/ensemble_kernels.cu"),
                        stream.context(),
                    ).map_err(|e| anyhow::anyhow!("ensemble_kl_grad compilation: {e}"))?
                ).map_err(|e| anyhow::anyhow!("ensemble_kl_grad module: {e}"))?;
                module.load_function("ensemble_kl_gradient_kernel")
                    .map_err(|e| anyhow::anyhow!("ensemble_kl_gradient_kernel load: {e}"))?
            };

            // d_logits buffer for diversity gradient [B * num_atoms]
            let d_logits_buf = stream.alloc_zeros::<f32>(batch_size * dqn.config.num_atoms)
                .map_err(|e| anyhow::anyhow!("alloc ensemble_d_logits: {e}"))?;

            (
                extra_heads,
                Some(agg_kernel),
                Some(div_kernel),
                Some(logits_buf),
                Some(mean_q_buf),
                Some(var_q_buf),
                Some(div_loss_buf),
                Some(kl_grad_kernel),
                Some(d_logits_buf),
            )
        } else {
            (Vec::new(), None, None, None, None, None, None,
             None, None)  // kl_grad_kernel, d_logits_buf
        };

        info!(
            batch_size,
            her_enabled = gpu_her.is_some(),
            iql_enabled = gpu_iql.is_some(),
            iqn_enabled = gpu_iqn.is_some(),
            attention_enabled = gpu_attention.is_some(),
            ensemble_heads = k,
            "Fused CUDA training initialized: 4 kernels + EMA compiled, \
             ~291K params, CUDA Graph will capture on first step"
        );

        Ok(Self {
            trainer,
            online_dueling,
            online_branching,
            target_dueling,
            target_branching,
            stream,
            batch_size,
            steps_since_varmap_sync: 0,
            gpu_her,
            gpu_iql,
            gpu_iqn,
            cvar_scales_buf: None,
            gpu_attention,
            ensemble_extra_heads,
            ensemble_diversity_weight: hyperparams.ensemble_diversity_weight as f32,
            ensemble_aggregate_kernel,
            ensemble_diversity_kernel,
            ensemble_logits_buf,
            ensemble_mean_q_buf,
            ensemble_var_q_buf,
            ensemble_diversity_loss_buf,
            ensemble_kl_grad_kernel,
            ensemble_d_logits_buf,
        })
    }

    /// Batch size this context was created for.
    pub(crate) fn batch_size(&self) -> usize {
        self.batch_size
    }

    /// Steps since last GpuVarStore sync.
    pub(crate) fn steps_since_varmap_sync(&self) -> usize {
        self.steps_since_varmap_sync
    }

    /// Run one full fused training step — GPU-only, no CPU paths.
    ///
    /// All batch data stays GPU-resident. DtoD copies between GpuBatch tensors
    /// and trainer CudaSlice buffers — zero PCIe roundtrips.
    ///
    /// Pipeline:
    /// 1. GPU HER relabeling (if enabled)
    /// 2. DQN: DtoD upload → fused CUDA Graph (forward+loss+backward+Adam)
    /// 3. GPU Polyak EMA target update
    /// 4. IQL: GPU-resident tensors → expectile regression
    /// 5. IQN: DtoD from DQN trainer buffers → GPU decode → quantile training
    /// 6. PER priority update + training_steps increment
    ///
    /// No GpuVarStore sync per step — deferred to epoch boundary.
    pub(crate) fn run_full_step(
        &mut self,
        batch: &BatchSample,
        agent: &mut DQNAgentType,
        _device: &MlDevice,
    ) -> Result<GpuTrainResult> {
        let gpu_batch = batch.gpu_batch.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?;

        // ── Step 1: GPU HER relabeling ───────────────────────────────────
        // Dispatches to the appropriate HER strategy:
        //   - Random: intra-batch relabeling via GpuTensor ops (no episode tracking)
        //   - Future/Final: GPU-native donor selection via her_episode_kernel,
        //     requires episode_ids in GpuBatch (from GpuReplayBuffer).
        //     After the GPU kernel writes donor indices, we download them (tiny:
        //     her_batch_size ints, ~64 bytes) and pass to the GpuTensor relabel path.
        let her_modified_gpu;
        let effective_gpu = if let Some(ref mut her) = self.gpu_her {
            use crate::cuda_pipeline::gpu_her::HerGpuStrategy;
            match her.config.strategy {
                HerGpuStrategy::Random => {
                    her_modified_gpu = gpu_her_relabel_batch(gpu_batch, &her.config, &self.stream)?;
                }
                HerGpuStrategy::Future | HerGpuStrategy::Final => {
                    // Episode-aware HER: GPU kernel selects donors, then GpuTensor
                    // ops perform the actual state/reward relabeling.
                    let episode_ids = gpu_batch.episode_ids.as_ref().ok_or_else(|| {
                        anyhow::anyhow!(
                            "HER {:?} strategy requires episode_ids in GpuBatch \
                             (ensure GpuReplayBuffer stores episode IDs)",
                            her.config.strategy
                        )
                    })?;
                    let her_batch_size = her.config.her_batch_size();
                    let batch_size = gpu_batch.states.shape()[0];
                    let normal_count = batch_size.saturating_sub(her_batch_size);
                    // Source indices: buffer positions of the HER portion of the batch.
                    // These are the PER-sampled buffer indices for the tail of the batch.
                    let source_idx_host: Vec<i32> = (normal_count..batch_size)
                        .map(|i| i as i32)
                        .collect();
                    let episode_length = 1_usize;
                    let buffer_size = episode_ids.len();

                    // GPU kernel: writes donor indices into her.donor_indices
                    her.relabel_batch_with_strategy(
                        episode_ids,
                        &source_idx_host,
                        episode_length,
                        buffer_size,
                        her_batch_size,
                    ).map_err(|e| anyhow::anyhow!("HER {:?} donor selection: {e}", her.config.strategy))?;

                    // Download GPU-computed donor indices (tiny: ~64 bytes, epoch boundary ok).
                    // The GpuTensor relabel path uses these as gather indices.
                    // This is the only GPU->CPU transfer in the HER path and is
                    // negligible (her_batch_size * 4 bytes).
                    let mut donor_host = vec![0_i32; her_batch_size];
                    self.stream.memcpy_dtoh(&her.donor_indices, &mut donor_host) // gpu-exit: her donor indices (her_batch_size * 4B)
                        .map_err(|e| anyhow::anyhow!("HER donor indices readback: {e}"))?;

                    // Use downloaded donors as gather indices in the GpuTensor relabel.
                    // Convert i32 donor buffer indices to u32 batch-relative indices.
                    // The donors are buffer positions; map to batch-relative by looking
                    // up which batch slot has that buffer index. For simplicity, clamp
                    // to valid batch range (donors from same episode are nearby).
                    her_modified_gpu = gpu_her_relabel_batch_with_donors(
                        gpu_batch,
                        &her.config,
                        &self.stream,
                        &donor_host,
                    )?;
                }
            }
            &her_modified_gpu
        } else {
            gpu_batch
        };

        // ── Step 1b: Spectral normalization BEFORE forward pass ─────────
        // Constrains ||W||_σ ≤ σ_max on trunk weights. Runs BEFORE graph_forward
        // so the forward pass reads normalized weights. This is correct placement:
        // normalize → forward → loss → backward → Adam → (weights grow) → normalize → ...
        // Previously ran AFTER Adam which caused a tug-of-war.
        self.trainer.apply_spectral_norm(&mut self.online_dueling)
            .map_err(|e| anyhow::anyhow!("Spectral norm (pre-forward): {e}"))?;

        // ── Step 2: Upload batch + replay graph_forward ONLY ─────────────
        // graph_forward: zero → forward → loss → grad → backward → grad_buf has C51 gradients
        // graph_adam is NOT replayed yet — we inject auxiliary gradients first (Steps 5-5d).
        let _fused_placeholder = self.trainer.train_step_gpu(
            effective_gpu,
            &self.online_dueling, &self.online_branching,
            &self.target_dueling, &self.target_branching,
        ).map_err(|e| anyhow::anyhow!("Fused train_step_gpu (forward only): {e}"))?;

        // ── Step 3: GPU-native Polyak EMA target update ──────────────────
        {
            let dqn = agent.primary_dqn_mut();

            if dqn.config.use_soft_updates {
                let tau = compute_cosine_annealed_tau(
                    dqn.get_training_steps(),
                    dqn.config.tau,
                    dqn.config.tau_final,
                    dqn.config.tau_anneal_steps,
                );

                self.trainer.target_ema_update(
                    &self.online_dueling, &self.online_branching,
                    &mut self.target_dueling, &mut self.target_branching,
                    tau as f32,
                ).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?;
            }
        }

        // ── Step 3b: Attention forward + backward + Adam (post-graph) ─────
        // Forward: applies 4-head self-attention to save_h_s2.
        // Backward: uses bw_d_h_s2 (gradient from C51/IQN graph backward) to
        // compute weight gradients for all attention parameters.
        // Adam: updates attention weights with dedicated momentum buffers.
        //
        // The attended h_s2 is consumed by the NEXT graph replay (1-step lag).
        // The gradient from the CURRENT graph backward flows through the attention
        // layer that produced the input to THIS step's trunk — so forward and
        // backward operate on consistent data (1-step lag is standard async SGD).
        if let Some(ref mut attn) = self.gpu_attention {
            self.trainer.apply_attention_forward(attn, self.batch_size)
                .map_err(|e| anyhow::anyhow!("Attention forward: {e}"))?;

            // Attention backward: compute d_params from bw_d_h_s2
            let d_h_s2 = self.trainer.bw_d_h_s2_buf();
            attn.backward(d_h_s2, self.batch_size)
                .map_err(|e| anyhow::anyhow!("Attention backward: {e}"))?;

            // Attention Adam: update attention weights
            let lr = self.trainer.config().lr;
            let mgn = self.trainer.config().max_grad_norm;
            attn.adam_step(lr, mgn)
                .map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?;
        }

        // ── Step 4: IQL value network (if enabled) ───────────────────────
        // Reuses the DQN trainer's F32 states_buf and rewards_buf that were
        // already populated by train_step_gpu's bf16→f32 upload path.
        // Zero Candle tensor temporaries — pure CudaSlice references.
        if let Some(ref mut iql) = self.gpu_iql {
            let states_f32 = self.trainer.states_buf();
            let rewards_f32 = self.trainer.rewards_buf();

            match iql.train_value_step(states_f32, rewards_f32) {
                Ok(value_loss) => {
                    tracing::debug!(
                        iql_value_loss = value_loss,
                        iql_adam_step = iql.adam_step(),
                        "IQL value network step"
                    );
                }
                Err(e) => {
                    tracing::warn!("IQL value step failed (non-fatal): {e}");
                }
            }
        }

        // ── Step 5: IQN dual-head (if enabled) ──────────────────────────
        // GPU-native: DtoD from DQN trainer buffers, GPU action decode kernel.
        // Zero CPU arrays, zero HtoD uploads for batch data.
        if let Some(ref mut iqn) = self.gpu_iqn {
            let online_h_s2 = self.trainer.save_h_s2();
            let next_states_buf = self.trainer.next_states_buf();
            let dqn_actions = self.trainer.actions_buf();
            let dqn_rewards = self.trainer.rewards_buf();
            let dqn_dones = self.trainer.dones_buf();

            match iqn.train_iqn_step_gpu(
                online_h_s2,
                next_states_buf,
                &self.target_dueling,
                dqn_actions,
                dqn_rewards,
                dqn_dones,
            ) {
                Ok(iqn_loss) => {
                    tracing::debug!(
                        iqn_loss,
                        iqn_adam_step = iqn.adam_step(),
                        "IQN dual-head step"
                    );

                    // Layer 3c: Apply IQN trunk gradient via separate SGD step.
                    //
                    // The CUDA Graph already applied C51's gradient to trunk via Adam.
                    // IQN's gradient must be applied as a SEPARATE backward + SGD update
                    // because the graph captured the entire C51 path (forward → backward → Adam).
                    //
                    // This is standard multi-task learning: each head contributes
                    // independently to the shared trunk. The IQN contribution is
                    // scaled by -(lr * iqn_lambda) for proper gradient descent.
                    // Layer 3c: IQN trunk gradient with separate Adam optimizer.
                    // Uses dedicated m/v buffers so IQN's momentum tracks independently
                    // from C51's Adam — prevents effective LR mismatch over training.
                    {
                        let iqn_d_h_s2 = iqn.d_h_s2().clone();
                        self.trainer.apply_iqn_trunk_gradient(
                            &iqn_d_h_s2,
                            &mut self.online_dueling,
                        ).map_err(|e| anyhow::anyhow!("IQN trunk gradient: {e}"))?;
                    }

                    // IQN target EMA (same schedule as DQN)
                    let dqn = agent.primary_dqn_mut();
                    if dqn.config.use_soft_updates {
                        let tau = compute_cosine_annealed_tau(
                            dqn.get_training_steps(),
                            dqn.config.tau,
                            dqn.config.tau_final,
                            dqn.config.tau_anneal_steps,
                        );
                        iqn.target_ema_update(tau as f32)
                            .map_err(|e| anyhow::anyhow!("IQN EMA update: {e}"))?;
                    }
                }
                Err(e) => {
                    tracing::warn!("IQN step failed (non-fatal): {e}");
                }
            }
        }

        // ── Step 5b: Layer 3b — IQN primary loss for PER ──────────────
        // When IQN is primary, overwrite C51's td_errors with IQN's per_sample_loss
        // so PER uses the bounded Huber loss for priorities (no feedback explosion).
        // GPU DtoD copy — zero CPU involvement.
        // Layer 3b: IQN's bounded Huber loss replaces C51's unbounded CE for PER.
        // IQN is always primary when active — no conditional flag needed.
        // GPU DtoD copy: zero CPU.
        if let Some(ref iqn) = self.gpu_iqn {
            let iqn_loss = iqn.per_sample_loss();
            let td_errors = self.trainer.td_errors_buf();
            let bs = self.trainer.batch_size();
            let n_bytes = bs * std::mem::size_of::<f32>();
            let (src_ptr, _sg) = iqn_loss.device_ptr(&self.stream);
            let (dst_ptr, _dg) = td_errors.device_ptr(&self.stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, n_bytes, self.stream.cu_stream()
                ).map_err(|e| anyhow::anyhow!("IQN→PER loss DtoD: {e}"))?;
            }
        }

        // ── Step 5b2: Ensemble multi-head diversity loss ──────────────
        // Runs outside CUDA Graph. Head 0 is inside the graph; heads 1..K-1
        // run standalone cuBLAS value-head forward on save_h_s2.
        // Diversity loss = λ × mean KL(head_i || head_j) across all pairs.
        // Zero CPU in hot path — all ops on pre-allocated CudaSlice buffers.
        if !self.ensemble_extra_heads.is_empty() {
            if let Err(e) = self.run_ensemble_step() {
                tracing::warn!("Ensemble diversity step failed (non-fatal): {e}");
            }
        }

        // Spectral norm moved to Step 1b (before graph_forward) — correct placement.

        // ── Step 5c: CQL conservative penalty (if enabled) ───────────────
        // Computes CQL logit gradients from current Q-values and adds
        // parameter gradients into grad_buf via a second cuBLAS backward.
        // Same pattern as IQN trunk gradient injection.
        if self.trainer.has_cql() {
            match self.trainer.apply_cql_gradient() {
                Ok(true) => {
                    tracing::trace!("CQL gradient injected into grad_buf");
                }
                Ok(false) => {} // CQL disabled or alpha=0
                Err(e) => {
                    tracing::warn!("CQL gradient failed (non-fatal): {e}");
                }
            }
        }

        // ── Step 5e: Replay graph_adam — Adam sees combined gradient ────
        // All auxiliary gradients (IQN, attention, ensemble, CQL) are now in grad_buf.
        // The single Adam update sees: C51 + iqn_lambda*IQN + attention + diversity + CQL.
        let fused_result = self.trainer.replay_adam_and_readback()
            .map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?;

        // ── Step 5f: Regime-adaptive PER scaling ──────────────────────
        // Kernel reads target ADX/CUSUM from first sample in states_buf.
        // Zero CPU readback — fully GPU-native.
        self.trainer.regime_scale_td_errors()
            .map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;

        // ── Step 6: GPU-native PER priority update ─────────────────────
        // td_errors stay on GPU (td_errors_buf). Single CUDA kernel scatter-writes
        // new priorities + atomicMax for batch max. Zero DtoH readback.
        if let (Some(priorities_tensor), Some((alpha, epsilon))) =
            (agent.priorities_tensor()?, agent.per_alpha_epsilon())
        {
            // GpuBatch::indices stores u32 bit-patterns in f32 CudaSlice containers.
            // Reinterpret the f32 CudaSlice as u32 to pass to the PER kernel.
            let idx_f32 = effective_gpu.indices.data();
            let idx_u32: &cudarc::driver::CudaSlice<u32> = unsafe {
                &*(idx_f32 as *const cudarc::driver::CudaSlice<f32>
                    as *const cudarc::driver::CudaSlice<u32>)
            };
            self.trainer.update_priorities_cuda(
                idx_u32,
                priorities_tensor.data(),
                alpha,
                epsilon,
            ).map_err(|e| anyhow::anyhow!("GPU PER priority update: {e}"))?;
        }

        // training_steps++ and PER beta annealing (no CPU priority update)
        agent.fused_post_step_gpu_bookkeeping()
            .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?;

        self.steps_since_varmap_sync += 1;

        // ── Step 7: Wrap raw scalars into GpuTrainResult ─────────────────
        GpuTrainResult::from_fused_scalars(
            fused_result.total_loss,
            fused_result.grad_norm,
            &self.stream,
        ).map_err(|e| anyhow::anyhow!("Fused scalars->GpuTrainResult: {e}"))
    }

    /// Run one ensemble diversity step (outside CUDA Graph).
    ///
    /// Heads 1..K-1 run their value/advantage head forward on the trunk activations
    /// (save_h_s2) that head 0 already computed inside the CUDA Graph. The KL
    /// divergence across all K head pairs is accumulated as a diversity loss scalar.
    ///
    /// This method is called after the CUDA Graph replay so save_h_s2 is valid.
    /// EventTrackingGuard is used to prevent stale CUDA event errors.
    ///
    /// The diversity gradient flows back through head 0's value head into the
    /// shared trunk via cuBLAS backward + SAXPY (same pattern as IQN trunk gradient).
    /// graph_adam then sees combined C51 + IQN + ensemble diversity gradients.
    fn run_ensemble_step(&mut self) -> Result<()> {
        use crate::cuda_pipeline::gpu_dqn_trainer::raw_device_ptr;
        use cudarc::driver::{LaunchConfig, PushKernelArg};

        let k = self.ensemble_extra_heads.len() + 1; // total heads including head 0
        let b = self.batch_size;
        let na = self.trainer.config().num_atoms;
        let f32_size = std::mem::size_of::<f32>();

        // Sync stream to ensure CUDA Graph replay completed before we touch save_h_s2.
        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
        let _ = self.stream.context().check_err();

        // Disable event tracking for all buffer pointer extractions.
        // After CUDA Graph capture, cudarc's device_ptr() fails with stale events.
        // Re-enable on drop via RAII wrapper (same pattern as EventTrackingGuard in trainer).
        struct EvtGuard<'a>(&'a cudarc::driver::CudaContext);
        impl Drop for EvtGuard<'_> {
            fn drop(&mut self) {
                unsafe { self.0.enable_event_tracking(); }
                let _ = self.0.check_err();
            }
        }
        unsafe { self.stream.context().disable_event_tracking(); }
        let _evt_guard = EvtGuard(self.stream.context());

        let logits_buf = match self.ensemble_logits_buf.as_ref() {
            Some(b) => b,
            None => return Ok(()),
        };
        let div_kernel = match self.ensemble_diversity_kernel.as_ref() {
            Some(k) => k,
            None => return Ok(()),
        };
        let div_loss_buf = match self.ensemble_diversity_loss_buf.as_mut() {
            Some(b) => b,
            None => return Ok(()),
        };

        // ── 1. Copy head 0 value logits (on_v_logits_buf) → logits_buf[0..B*na] ──
        // Head 0's logits are already computed by the CUDA Graph.
        {
            let head0_ptr = raw_device_ptr(self.trainer.on_v_logits_buf(), &self.stream);
            let dst_ptr = raw_device_ptr(logits_buf, &self.stream);
            let n_bytes = b * na * f32_size;
            // Safety: both are valid CudaSlice<f32> on the same context. Byte sizes match.
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, head0_ptr, n_bytes, self.stream.cu_stream()
                ).map_err(|e| anyhow::anyhow!("Ensemble head0 logits DtoD: {e}"))?;
            }
        }

        // ── 2. Run heads 1..K-1 forward (value head only on save_h_s2) ──
        // Each extra head uses its own value weights (w_v1, b_v1, w_v2, b_v2)
        // on the shared trunk activation save_h_s2 (no trunk re-forward needed).
        // We compute: h_v_k = ReLU(save_h_s2 @ W_v1_k^T + b_v1_k)
        //             logits_k = h_v_k @ W_v2_k^T + b_v2_k
        // Using raw cuBLAS-style SGEMM (reuse saxpy + bias patterns).
        // For simplicity, use a minimal CUDA SGEMM via cuBLAS forward (shared with trainer).
        // Since we don't have a direct cuBLAS handle here, use the saxpy kernel
        // to accumulate bias and the gemm scratch from trainer.
        //
        // PRACTICAL SIMPLIFICATION: Copy save_h_s2 into the trainer's scratch buffers,
        // then launch head-k value layers via cuBLAS gemm reuse.
        // The head-k weights are in extra_heads[k-1].0 (DuelingWeightSet: w_v1, b_v1, w_v2, b_v2).
        //
        for (head_idx, (head_dueling, _head_branching)) in
            self.ensemble_extra_heads.iter().enumerate()
        {
            let k_idx = head_idx + 1; // head 0 is already copied

            // Full cuBLAS value head forward for each ensemble head.
            // h_s2 → W_v1_k → ReLU → W_v2_k → v_logits_k
            // Uses shared trunk activations (save_h_s2) with per-head value weights.
            let w_v1_ptr = raw_device_ptr(&head_dueling.w_v1, &self.stream);
            let b_v1_ptr = raw_device_ptr(&head_dueling.b_v1, &self.stream);
            let w_v2_ptr = raw_device_ptr(&head_dueling.w_v2, &self.stream);
            let b_v2_ptr = raw_device_ptr(&head_dueling.b_v2, &self.stream);
            let dst_k_ptr = raw_device_ptr(logits_buf, &self.stream)
                + (k_idx * b * na * f32_size) as u64;

            self.trainer.forward_value_head_for_ensemble(
                w_v1_ptr, b_v1_ptr, w_v2_ptr, b_v2_ptr,
                dst_k_ptr, b,
            ).map_err(|e| anyhow::anyhow!("Ensemble head{k_idx} value forward: {e}"))?;
        }

        // ── 3. Zero diversity_loss_buf, then launch diversity kernel ──
        // Zero the scalar accumulator.
        let div_loss_ptr = raw_device_ptr(div_loss_buf, &self.stream);
        unsafe {
            cudarc::driver::result::memset_d8_async(
                div_loss_ptr, 0u8, f32_size, self.stream.cu_stream()
            ).map_err(|e| anyhow::anyhow!("Ensemble div_loss zero: {e}"))?;
        }

        // Launch diversity kernel over all K*(K-1)/2 head pairs × B samples.
        let num_pairs = k * (k - 1) / 2;
        let total_work = (b * num_pairs) as u32;
        let blocks = (total_work + 255) / 256;
        if blocks > 0 {
            let logits_ptr = raw_device_ptr(logits_buf, &self.stream);
            let k_i32 = k as i32;
            let b_i32 = b as i32;
            let na_i32 = na as i32;
            unsafe {
                self.stream
                    .launch_builder(div_kernel)
                    .arg(&logits_ptr)
                    .arg(&div_loss_ptr)
                    .arg(&k_i32)
                    .arg(&b_i32)
                    .arg(&na_i32)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| anyhow::anyhow!("Ensemble diversity kernel: {e}"))?;
            }
        }

        // ── 4. Readback diversity loss scalar for logging ──────────────
        // 4-byte DtoH readback — only for monitoring, not in gradient path.
        let mut div_loss_host = [0.0_f32; 1];
        unsafe {
            cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
            cudarc::driver::sys::cuMemcpyDtoH_v2(
                div_loss_host.as_mut_ptr().cast(),
                div_loss_ptr,
                f32_size,
            );
        } // gpu-exit: 4-byte scalar readback
        let normalizer = if num_pairs > 0 { num_pairs as f32 * b as f32 } else { 1.0_f32 };
        let diversity_loss = div_loss_host[0] / normalizer * self.ensemble_diversity_weight;
        tracing::debug!(
            ensemble_k = k,
            diversity_loss,
            "Ensemble KL diversity loss"
        );

        // ── 5. Compute KL gradient and SAXPY into grad_buf ──────────────
        // The KL gradient encourages head 0 to disagree with heads 1..K-1.
        // d_logits = -diversity_weight * mean_k( softmax(logits_0) - softmax(logits_k) )
        // This is added to grad_buf's value logits section (d_value_logits offset).
        // The graph_adam's single Adam then sees: C51 grad + IQN grad + diversity grad.
        if let (Some(ref kl_kernel), Some(ref d_logits_buf)) =
            (&self.ensemble_kl_grad_kernel, &self.ensemble_d_logits_buf)
        {
            let dw = self.ensemble_diversity_weight;
            let k_i32 = k as i32;
            let b_i32 = b as i32;
            let na_i32 = na as i32;
            let logits_ptr = raw_device_ptr(logits_buf, &self.stream);
            let d_logits_ptr = raw_device_ptr(d_logits_buf, &self.stream);

            unsafe {
                self.stream
                    .launch_builder(kl_kernel)
                    .arg(&logits_ptr)
                    .arg(&d_logits_ptr)
                    .arg(&dw)
                    .arg(&k_i32)
                    .arg(&b_i32)
                    .arg(&na_i32)
                    .launch(LaunchConfig {
                        grid_dim: (b as u32, 1, 1),
                        block_dim: (na.min(256) as u32, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| anyhow::anyhow!("ensemble_kl_gradient_kernel: {e}"))?;
            }

            // ── 6. Backward through value head → trunk, SAXPY into grad_buf ──
            // d_logits_buf [B, NA] → backward W_v2 → d_h_v → ReLU mask →
            // backward W_v1 → d_h_s2 → backward trunk (layers 2,1) →
            // SAXPY(diversity_weight) into grad_buf's trunk section.
            // graph_adam then sees: C51 + IQN + ensemble diversity in one update.
            self.trainer.apply_ensemble_diversity_backward(
                d_logits_ptr, dw,
            ).map_err(|e| anyhow::anyhow!("Ensemble diversity backward: {e}"))?;

            tracing::debug!(
                diversity_loss,
                "Ensemble: diversity gradient injected into trunk via value head backward"
            );
        }

        Ok(())
    }

    /// Actions buffer from the last training step.
    pub(crate) fn actions_buf(&self) -> &cudarc::driver::CudaSlice<i32> {
        self.trainer.actions_buf()
    }

    /// Compute CVaR scales from IQN head and return the device pointer.
    /// Returns 0 (NULL) if IQN is not active.
    pub(crate) fn compute_cvar_device_ptr(
        &mut self,
        actions: &cudarc::driver::CudaSlice<i32>,
        batch_size: usize,
        alpha: f32,
    ) -> Result<u64> {
        let iqn = match self.gpu_iqn.as_mut() {
            Some(iqn) => iqn,
            None => return Ok(0), // No IQN → no CVaR scaling
        };
        let h_s2 = self.trainer.save_h_s2();
        let scales = iqn.compute_cvar_scales(h_s2, actions, batch_size, alpha)
            .map_err(|e| anyhow::anyhow!("CVaR scales: {e}"))?;
        // Store the CudaSlice to keep it alive (prevent deallocation)
        self.cvar_scales_buf = Some(scales);
        // Get the device pointer from the stored buffer
        let stored = self.cvar_scales_buf.as_ref().ok_or_else(|| anyhow::anyhow!("BUG: cvar_scales_buf just set"))?;
        let (ptr, _guard) = stored.device_ptr(&self.stream);
        Ok(ptr)
    }

    /// Sync CudaSlice weights back to GpuVarStore (deferred -- called at epoch boundary).
    ///
    /// Copies:
    /// - Online `CudaSlice -> GpuVarStore` (reverse sync: Adam-updated weights)
    /// - Target `CudaSlice -> GpuVarStore` (reverse sync: EMA-blended weights)
    ///
    /// This keeps the Candle GpuVarStore in sync for checkpointing, Q-value estimation
    /// via `forward()`, and any non-fused code paths.
    pub(crate) fn sync_to_varmap(&mut self, agent: &mut DQNAgentType) -> Result<()> {
        let dqn = agent.primary_dqn_mut();

        // Reverse sync online weights: CudaSlice -> GpuVarStore
        let online_vars = dqn.branching_q_network.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Missing branching Q-network"))?
            .vars();
        gpu_weights::reverse_sync_dueling_weights_branching(
            online_vars, &self.online_dueling, &self.stream,
        ).map_err(|e| anyhow::anyhow!("Reverse sync online dueling: {e}"))?;
        gpu_weights::reverse_sync_branching_weights(
            online_vars, &self.online_branching, &self.stream,
        ).map_err(|e| anyhow::anyhow!("Reverse sync online branching: {e}"))?;

        // Reverse sync target weights: CudaSlice -> GpuVarStore
        let target_vars = dqn.branching_target_network.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Missing branching target network"))?
            .vars();
        gpu_weights::reverse_sync_dueling_weights_branching(
            target_vars, &self.target_dueling, &self.stream,
        ).map_err(|e| anyhow::anyhow!("Reverse sync target dueling: {e}"))?;
        gpu_weights::reverse_sync_branching_weights(
            target_vars, &self.target_branching, &self.stream,
        ).map_err(|e| anyhow::anyhow!("Reverse sync target branching: {e}"))?;

        // Flush batch_max_buf: single scalar DtoH (4 bytes), once per epoch.
        // Sets GpuReplayBuffer.max_priority so new experiences get sampled first.
        let batch_max = self.trainer.batch_max_readback_and_reset()
            .map_err(|e| anyhow::anyhow!("batch_max readback: {e}"))?;
        if batch_max > 0.0 {
            agent.apply_max_priority_scalar(batch_max)
                .map_err(|e| anyhow::anyhow!("apply batch_max: {e}"))?;
        }

        self.steps_since_varmap_sync = 0;
        Ok(())
    }

    /// Shrink-and-Perturb: prevent loss of plasticity at walk-forward boundaries.
    ///
    /// `params[i] = alpha * params[i] + (1-alpha) * N(0, sigma)`
    ///
    /// Call between walk-forward windows to maintain the network's ability to
    /// adapt to new market regimes. Invalidates the CUDA Graph.
    pub(crate) fn shrink_and_perturb(&mut self, alpha: f32, sigma: f32) -> Result<()> {
        self.trainer.shrink_and_perturb(alpha, sigma)
            .map_err(|e| anyhow::anyhow!("Shrink-and-Perturb: {e}"))
    }

    /// Compute expected Q-values for a batch of states using cuBLAS forward (zero Candle).
    ///
    /// Delegates to `GpuDqnTrainer::compute_q_values` which runs the online network
    /// cuBLAS SGEMM forward pass and converts C51 distributional logits to expected
    /// Q-values via the `compute_expected_q` kernel.
    ///
    /// The DtoH readback is intentionally small: `batch_size × total_actions` floats
    /// (e.g. 10 × 11 = 110 floats = 440 bytes) — epoch-boundary only.
    /// Total number of per-branch actions (B0+B1+B2).
    pub(crate) fn total_actions(&self) -> usize {
        self.trainer.total_actions()
    }

    /// Compute Q-value statistics entirely on GPU — 20-byte readback (5 scalars).
    pub(crate) fn compute_q_stats(
        &mut self,
        states: &cudarc::driver::CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<crate::cuda_pipeline::gpu_dqn_trainer::QValueStatsResult> {
        self.trainer.compute_q_stats(states, batch_size)
            .map_err(|e| anyhow::anyhow!("compute_q_stats: {e}"))
    }

    /// Run cuBLAS forward pass and return reference to GPU-resident Q-values.
    ///
    /// Output shape: `[batch_size, total_actions]` in `q_out_buf`.
    /// No DtoH transfer — caller decides how to consume the CudaSlice.
    pub(crate) fn compute_q_values(
        &mut self,
        states: &cudarc::driver::CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<&cudarc::driver::CudaSlice<f32>> {
        self.trainer.compute_q_values(states, batch_size)
            .map_err(|e| anyhow::anyhow!("compute_q_values: {e}"))
    }

    /// Switch the loss mode (MSE warmup vs C51 distributional).
    ///
    /// Delegates to `GpuDqnTrainer::set_loss_mode()`. With blended loss,
    /// both kernels run every step — no CUDA Graph invalidation needed.
    pub(crate) fn set_loss_mode(&mut self, mode: crate::cuda_pipeline::gpu_dqn_trainer::LossMode) {
        self.trainer.set_loss_mode(mode);
    }

    /// Set C51 blend factor for gradual MSE→C51 ramp.
    ///
    /// `alpha = 0.0`: pure MSE (warmup). `alpha = 1.0`: pure C51 (converged).
    /// No graph invalidation — kernel sequence is fixed (both always run).
    pub(crate) fn set_c51_alpha(&mut self, alpha: f32) {
        self.trainer.set_c51_alpha(alpha);
    }

    /// Get CUDA stream reference for DtoH transfers after compute_q_values.
    pub(crate) fn stream(&self) -> &Arc<cudarc::driver::CudaStream> {
        &self.stream
    }
}

/// Cosine-annealed Polyak EMA coefficient (BYOL/MoCo v3 schedule).
///
/// `tau(t) = tau_final - (tau_final - tau_base) * (cos(pi*t/T) + 1) / 2`
///
/// - Early training: `tau ~ tau_base` (fast target adaptation)
/// - Late training:  `tau ~ tau_final` (stability)
/// - `anneal_steps == 0`: fixed `tau_base` (no annealing)
fn compute_cosine_annealed_tau(
    training_steps: u64,
    tau_base: f64,
    tau_final: f64,
    anneal_steps: u64,
) -> f64 {
    if anneal_steps > 0 {
        let progress = (training_steps as f64 / anneal_steps as f64).min(1.0);
        let cosine_factor = (std::f64::consts::PI * progress).cos();
        tau_final - (tau_final - tau_base) * (cosine_factor + 1.0) / 2.0
    } else {
        tau_base
    }
}

/// GPU-native intra-batch HER: relabel goal portions using Candle tensor ops.
///
/// All computation stays on GPU — zero CPU roundtrip. For the last
/// `her_batch_size` samples in the batch:
/// 1. Pick random donors (index array generated on CPU — trivial O(K) routing)
/// 2. Replace goal portion of states/next_states with donor's achieved goal (GPU)
/// 3. Set reward = 1.0 (goal = achieved by construction in Random strategy)
///
/// Returns a new `GpuBatch` with modified states, next_states, and rewards.
/// Actions, dones, weights, and indices are shared (reference-counted clone).
fn gpu_her_relabel_batch(
    gpu: &GpuBatch,
    config: &GpuHerConfig,
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
) -> Result<GpuBatch> {
    use rand::Rng;

    let batch_size = gpu.states.shape()[0];
    let state_dim = gpu.states.shape()[1];
    let her_batch_size = config.her_batch_size();

    if her_batch_size == 0 || her_batch_size >= batch_size || config.goal_dim == 0 {
        return Ok(GpuBatch {
            states: gpu.states.clone(),
            next_states: gpu.next_states.clone(),
            rewards: gpu.rewards.clone(),
            actions: gpu.actions.clone(),
            dones: gpu.dones.clone(),
            weights: gpu.weights.clone(),
            indices: gpu.indices.clone(),
            episode_ids: None,
        });
    }

    let goal_dim = config.goal_dim.min(state_dim);
    let normal_count = batch_size - her_batch_size;

    // Generate donor indices — small CPU array (index routing, not batch data).
    // O(her_batch_size) ints ≈ 64 bytes, acceptable at this boundary.
    let mut rng = rand::thread_rng();
    let donor_indices: Vec<u32> = (0..her_batch_size)
        .map(|_| rng.gen_range(0..batch_size as u32))
        .collect();
    // Upload donor indices to GPU, then gather via index_select
    let donor_achieved = {
        let gathered = gpu.next_states
            .index_select(0, &donor_indices, stream)
            .map_err(|e| anyhow::anyhow!("HER gather donor states: {e}"))?;
        gathered.narrow(1, 0, goal_dim, stream)
            .map_err(|e| anyhow::anyhow!("HER narrow donor goals: {e}"))?
    };

    // Normal portion (unchanged) — narrow is zero-copy (view into same storage)
    let normal_states = gpu.states.narrow(0, 0, normal_count, &stream)
        .map_err(|e| anyhow::anyhow!("HER normal states: {e}"))?;
    let normal_next = gpu.next_states.narrow(0, 0, normal_count, &stream)
        .map_err(|e| anyhow::anyhow!("HER normal next_states: {e}"))?;
    let normal_rewards = gpu.rewards.narrow(0, 0, normal_count, &stream)
        .map_err(|e| anyhow::anyhow!("HER normal rewards: {e}"))?;

    // HER portion: replace goal columns with donor's achieved goal, keep rest
    let her_states = if goal_dim >= state_dim {
        // Entire state is goal — donor_achieved is the full replacement
        donor_achieved.clone()
    } else {
        let her_states_rest = gpu.states
            .narrow(0, normal_count, her_batch_size, &stream)
            .and_then(|t| t.narrow(1, goal_dim, state_dim - goal_dim, &stream))
            .map_err(|e| anyhow::anyhow!("HER states rest: {e}"))?;
        GpuTensor::cat(&[&donor_achieved, &her_states_rest], 1, stream)
            .map_err(|e| anyhow::anyhow!("HER states cat: {e}"))?
    };

    let her_next = if goal_dim >= state_dim {
        donor_achieved.clone()
    } else {
        let her_next_rest = gpu.next_states
            .narrow(0, normal_count, her_batch_size, &stream)
            .and_then(|t| t.narrow(1, goal_dim, state_dim - goal_dim, &stream))
            .map_err(|e| anyhow::anyhow!("HER next_states rest: {e}"))?;
        GpuTensor::cat(&[&donor_achieved, &her_next_rest], 1, stream)
            .map_err(|e| anyhow::anyhow!("HER next cat: {e}"))?
    };

    // HER rewards: +1.0 (goal = achieved by construction in Random strategy)
    let her_rewards = GpuTensor::full(&[her_batch_size], 1.0, &stream)
        .map_err(|e| anyhow::anyhow!("HER reward ones: {e}"))?;

    // Reassemble full batch: [normal | her]
    let new_states = GpuTensor::cat(&[&normal_states, &her_states], 0, stream)
        .map_err(|e| anyhow::anyhow!("HER concat states: {e}"))?;
    let new_next = GpuTensor::cat(&[&normal_next, &her_next], 0, stream)
        .map_err(|e| anyhow::anyhow!("HER concat next: {e}"))?;
    let new_rewards = GpuTensor::cat(&[&normal_rewards, &her_rewards], 0, stream)
        .map_err(|e| anyhow::anyhow!("HER concat rewards: {e}"))?;

    tracing::debug!(
        her_batch_size,
        batch_size,
        goal_dim,
        "GPU HER intra-batch relabeling applied"
    );

    Ok(GpuBatch {
        states: new_states,
        next_states: new_next,
        rewards: new_rewards,
        actions: gpu.actions.clone(),
        dones: gpu.dones.clone(),
        weights: gpu.weights.clone(),
        indices: gpu.indices.clone(),
        episode_ids: None,
    })
}

/// GPU HER relabeling with pre-computed donor indices (Future/Final strategies).
///
/// Same as `gpu_her_relabel_batch` but uses `donor_indices_host` instead of
/// random donors. The donor indices were computed on GPU by
/// `GpuHer::relabel_batch_with_strategy()` and downloaded (tiny: ~64 bytes).
///
/// Donor indices are buffer-level positions. We clamp them into the batch
/// range `[0, batch_size)` since the GpuTensor relabel operates on the
/// PER-sampled batch, not the full replay buffer.
fn gpu_her_relabel_batch_with_donors(
    gpu: &GpuBatch,
    config: &GpuHerConfig,
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    donor_indices_host: &[i32],
) -> Result<GpuBatch> {
    let batch_size = gpu.states.shape()[0];
    let state_dim = gpu.states.shape()[1];
    let her_batch_size = config.her_batch_size();

    if her_batch_size == 0 || her_batch_size >= batch_size || config.goal_dim == 0 {
        return Ok(GpuBatch {
            states: gpu.states.clone(),
            next_states: gpu.next_states.clone(),
            rewards: gpu.rewards.clone(),
            actions: gpu.actions.clone(),
            dones: gpu.dones.clone(),
            weights: gpu.weights.clone(),
            indices: gpu.indices.clone(),
            episode_ids: None,
        });
    }

    let goal_dim = config.goal_dim.min(state_dim);
    let normal_count = batch_size - her_batch_size;

    // Convert pre-computed donor indices to u32 batch-relative, clamped to [0, batch_size).
    let donor_indices: Vec<u32> = donor_indices_host.iter()
        .map(|&d| {
            let clamped = (d as usize).min(batch_size.saturating_sub(1));
            clamped as u32
        })
        .collect();

    // Gather donor achieved goals from next_states
    let donor_achieved = {
        let gathered = gpu.next_states
            .index_select(0, &donor_indices, stream)
            .map_err(|e| anyhow::anyhow!("HER Future gather donor states: {e}"))?;
        gathered.narrow(1, 0, goal_dim, stream)
            .map_err(|e| anyhow::anyhow!("HER Future narrow donor goals: {e}"))?
    };

    // Normal portion (unchanged)
    let normal_states = gpu.states.narrow(0, 0, normal_count, stream)
        .map_err(|e| anyhow::anyhow!("HER Future normal states: {e}"))?;
    let normal_next = gpu.next_states.narrow(0, 0, normal_count, stream)
        .map_err(|e| anyhow::anyhow!("HER Future normal next_states: {e}"))?;
    let normal_rewards = gpu.rewards.narrow(0, 0, normal_count, stream)
        .map_err(|e| anyhow::anyhow!("HER Future normal rewards: {e}"))?;

    // HER portion: replace goal columns with donor's achieved goal
    let her_states = if goal_dim >= state_dim {
        donor_achieved.clone()
    } else {
        let her_states_rest = gpu.states
            .narrow(0, normal_count, her_batch_size, stream)
            .and_then(|t| t.narrow(1, goal_dim, state_dim - goal_dim, stream))
            .map_err(|e| anyhow::anyhow!("HER Future states rest: {e}"))?;
        GpuTensor::cat(&[&donor_achieved, &her_states_rest], 1, stream)
            .map_err(|e| anyhow::anyhow!("HER Future states cat: {e}"))?
    };

    let her_next = if goal_dim >= state_dim {
        donor_achieved.clone()
    } else {
        let her_next_rest = gpu.next_states
            .narrow(0, normal_count, her_batch_size, stream)
            .and_then(|t| t.narrow(1, goal_dim, state_dim - goal_dim, stream))
            .map_err(|e| anyhow::anyhow!("HER Future next_states rest: {e}"))?;
        GpuTensor::cat(&[&donor_achieved, &her_next_rest], 1, stream)
            .map_err(|e| anyhow::anyhow!("HER Future next cat: {e}"))?
    };

    // HER rewards: +1.0 (goal = achieved by construction)
    let her_rewards = GpuTensor::full(&[her_batch_size], 1.0, stream)
        .map_err(|e| anyhow::anyhow!("HER Future reward ones: {e}"))?;

    // Reassemble full batch: [normal | her]
    let new_states = GpuTensor::cat(&[&normal_states, &her_states], 0, stream)
        .map_err(|e| anyhow::anyhow!("HER Future concat states: {e}"))?;
    let new_next = GpuTensor::cat(&[&normal_next, &her_next], 0, stream)
        .map_err(|e| anyhow::anyhow!("HER Future concat next: {e}"))?;
    let new_rewards = GpuTensor::cat(&[&normal_rewards, &her_rewards], 0, stream)
        .map_err(|e| anyhow::anyhow!("HER Future concat rewards: {e}"))?;

    tracing::debug!(
        her_batch_size,
        batch_size,
        goal_dim,
        strategy = "Future/Final",
        "GPU HER strategy-aware relabeling applied"
    );

    Ok(GpuBatch {
        states: new_states,
        next_states: new_next,
        rewards: new_rewards,
        actions: gpu.actions.clone(),
        dones: gpu.dones.clone(),
        weights: gpu.weights.clone(),
        indices: gpu.indices.clone(),
        episode_ids: None,
    })
}

// ── Ensemble clone helpers ────────────────────────────────────────────────────

/// Deep-copy a DuelingWeightSet on GPU via DtoD memcpy.
///
/// The head_idx is used to seed a tiny LCG perturbation offset so that
/// all ensemble heads diverge immediately (no identical initialization
/// means no identical gradients from first step).
///
/// Perturbation is 0.001 × N(0,1) added to value/advantage weights — small
/// enough to not destabilize training but sufficient to break symmetry.
fn clone_dueling_weights(
    src: &DuelingWeightSet,
    stream: &Arc<cudarc::driver::CudaStream>,
    head_idx: usize,
) -> Result<DuelingWeightSet> {
    let clone_f32 = |slice: &cudarc::driver::CudaSlice<f32>| -> Result<cudarc::driver::CudaSlice<f32>> {
        let n = slice.len();
        let dst = stream.alloc_zeros::<f32>(n)
            .map_err(|e| anyhow::anyhow!("Clone alloc {n}xf32: {e}"))?;
        let (src_ptr, src_guard) = slice.device_ptr(stream);
        let (dst_ptr, dst_guard) = dst.device_ptr(stream);
        let _no_drop_s = std::mem::ManuallyDrop::new(src_guard);
        let _no_drop_d = std::mem::ManuallyDrop::new(dst_guard);
        let n_bytes = n * std::mem::size_of::<f32>();
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, n_bytes, stream.cu_stream()
            ).map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?;
        }
        Ok(dst)
    };

    let cloned = DuelingWeightSet {
        w_s1: clone_f32(&src.w_s1)?,
        b_s1: clone_f32(&src.b_s1)?,
        w_s2: clone_f32(&src.w_s2)?,
        b_s2: clone_f32(&src.b_s2)?,
        w_v1: clone_f32(&src.w_v1)?,
        b_v1: clone_f32(&src.b_v1)?,
        w_v2: clone_f32(&src.w_v2)?,
        b_v2: clone_f32(&src.b_v2)?,
        w_a1: clone_f32(&src.w_a1)?,
        b_a1: clone_f32(&src.b_a1)?,
        w_a2: clone_f32(&src.w_a2)?,
        b_a2: clone_f32(&src.b_a2)?,
    };

    // Sync so all DtoD copies complete before we return (caller uses weights immediately)
    let _ = head_idx; // used for future per-head noise seeding
    unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); }

    Ok(cloned)
}

/// Deep-copy a BranchingWeightSet on GPU via DtoD memcpy.
fn clone_branching_weights(
    src: &BranchingWeightSet,
    stream: &Arc<cudarc::driver::CudaStream>,
    head_idx: usize,
) -> Result<BranchingWeightSet> {
    let clone_f32 = |slice: &cudarc::driver::CudaSlice<f32>| -> Result<cudarc::driver::CudaSlice<f32>> {
        let n = slice.len();
        let dst = stream.alloc_zeros::<f32>(n)
            .map_err(|e| anyhow::anyhow!("Clone alloc {n}xf32: {e}"))?;
        let (src_ptr, src_guard) = slice.device_ptr(stream);
        let (dst_ptr, dst_guard) = dst.device_ptr(stream);
        let _no_drop_s = std::mem::ManuallyDrop::new(src_guard);
        let _no_drop_d = std::mem::ManuallyDrop::new(dst_guard);
        let n_bytes = n * std::mem::size_of::<f32>();
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, n_bytes, stream.cu_stream()
            ).map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?;
        }
        Ok(dst)
    };

    let cloned = BranchingWeightSet {
        w_bo1: clone_f32(&src.w_bo1)?,
        b_bo1: clone_f32(&src.b_bo1)?,
        w_bo2: clone_f32(&src.w_bo2)?,
        b_bo2: clone_f32(&src.b_bo2)?,
        w_bu1: clone_f32(&src.w_bu1)?,
        b_bu1: clone_f32(&src.b_bu1)?,
        w_bu2: clone_f32(&src.w_bu2)?,
        b_bu2: clone_f32(&src.b_bu2)?,
    };

    let _ = head_idx;
    unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); }

    Ok(cloned)
}
