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

//! GPU-accelerated IQN (Implicit Quantile Network) dual-head trainer.
//!
//! Auxiliary quantile head that shares the DQN trunk (h_s2) but trains
//! its own weights independently. Provides quantile-based Q-values for
//! risk-sensitive action selection (CVaR) alongside the primary C51 head.
//!
//! ## Architecture
//!
//! For each τ ∈ (0,1):
//!   cos_feat  = [cos(π·1·τ), ..., cos(π·D·τ)]
//!   embed     = ReLU(W_embed × cos_feat + b_embed)
//!   combined  = h_s2 ⊙ embed
//!   q_branch  = W_bd × combined + b_bd    (per branch)
//!
//! ## Kernel pipeline
//!
//! 1. **Trunk forward** (`iqn_trunk_forward_kernel`): compute target h_s2 from next_states
//! 2. **Forward + loss** (`iqn_forward_loss_kernel`): quantile forward + Huber loss
//! 3. **Backward** (`iqn_backward_kernel`): gradients through IQN weights only
//! 4. **Grad norm** (`iqn_grad_norm_kernel`): L2 gradient norm
//! 5. **Adam** (`iqn_adam_kernel`): AdamW with gradient clipping
//! 6. **EMA** (`iqn_ema_kernel`): Polyak target weight update
//! 7. **Forward only** (`iqn_forward_kernel`): inference, returns expected Q per action
//!
//! ## Integration
//!
//! Called after the main DQN training step in `FusedTrainingCtx::run_full_step()`.
//! Reads `save_h_s2` (online trunk activation) from `GpuDqnTrainer` and computes
//! target h_s2 from `next_states_buf` using the target DQN trunk weights.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use tracing::info;

use crate::MLError;
use super::gpu_weights::DuelingWeightSet;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// Configuration for the GPU IQN dual-head trainer.
#[derive(Debug, Clone)]
pub struct GpuIqnConfig {
    /// Hidden dimension = SHARED_H2 from DQN trunk (default 256).
    pub hidden_dim: usize,
    /// Cosine embedding dimension (default 64).
    pub embed_dim: usize,
    /// Number of quantile samples per forward pass (default 32).
    pub num_quantiles: usize,
    /// Huber threshold for quantile loss (default 1.0).
    pub kappa: f32,
    /// Input state dimension (tensor-core-aligned, e.g. 48 or 56).
    pub state_dim: usize,
    /// First shared hidden layer width from DQN trunk (default 256).
    pub shared_h1: usize,
    /// Branch 0 (exposure) action count.
    pub branch_0_size: usize,
    /// Branch 1 (order) action count.
    pub branch_1_size: usize,
    /// Branch 2 (urgency) action count.
    pub branch_2_size: usize,
    /// Training batch size (must match DQN trainer).
    pub batch_size: usize,
    /// Learning rate for IQN Adam optimizer.
    pub lr: f32,
    /// Adam beta1 (first moment decay).
    pub beta1: f32,
    /// Adam beta2 (second moment decay).
    pub beta2: f32,
    /// Adam epsilon (numerical stability).
    pub epsilon: f32,
    /// Decoupled weight decay (AdamW).
    pub weight_decay: f32,
    /// Maximum gradient L2 norm for clipping.
    pub max_grad_norm: f32,
    /// Discount factor for Bellman projection.
    pub gamma: f32,
}

impl Default for GpuIqnConfig {
    fn default() -> Self {
        Self {
            hidden_dim: 256,
            embed_dim: 64,
            num_quantiles: 32,
            kappa: 1.0,
            state_dim: 72,
            shared_h1: 256,
            branch_0_size: 9,
            branch_1_size: 3,
            branch_2_size: 3,
            batch_size: 256,
            lr: 3e-4,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            weight_decay: 1e-5,
            max_grad_norm: 1.0,
            gamma: 0.99,
        }
    }
}

impl GpuIqnConfig {
    /// Total branch actions (sum of all branches).
    fn total_branch_actions(&self) -> usize {
        self.branch_0_size + self.branch_1_size + self.branch_2_size
    }

    /// Total IQN parameters (embedding + per-branch output layers).
    fn total_params(&self) -> usize {
        let h = self.hidden_dim;
        let d = self.embed_dim;
        let b0 = self.branch_0_size;
        let b1 = self.branch_1_size;
        let b2 = self.branch_2_size;
        // W_embed[H, D] + b_embed[H] + W_b0[B0, H] + b_b0[B0]
        // + W_b1[B1, H] + b_b1[B1] + W_b2[B2, H] + b_b2[B2]
        h * d + h + b0 * h + b0 + b1 * h + b1 + b2 * h + b2
    }
}

// ---------------------------------------------------------------------------
// GpuIqnHead
// ---------------------------------------------------------------------------

/// GPU-accelerated IQN dual-head trainer.
///
/// Owns online + target IQN weight buffers, Adam state, compiled CUDA kernels,
/// and all intermediate activation/loss buffers. Operates entirely on GPU.
#[allow(missing_debug_implementations)]
pub struct GpuIqnHead {
    config: GpuIqnConfig,
    stream: Arc<CudaStream>,

    // ── Compiled kernels ────────────────────────────────────────────
    trunk_forward_kernel: CudaFunction,
    forward_loss_kernel: CudaFunction,
    backward_kernel: CudaFunction,
    grad_norm_kernel: CudaFunction,
    adam_kernel: CudaFunction,
    forward_kernel: CudaFunction,
    ema_kernel: CudaFunction,
    decode_actions_kernel: CudaFunction,
    sample_taus_kernel: CudaFunction,

    // ── IQN online parameters (flat f32 on GPU) ──────────────────────
    online_params: CudaSlice<f32>,
    // ── IQN target parameters (for Bellman target) ───────────────────
    target_params: CudaSlice<f32>,

    // ── Adam optimizer state ─────────────────────────────────────────
    m_buf: CudaSlice<f32>,
    v_buf: CudaSlice<f32>,
    grad_buf: CudaSlice<f32>,
    grad_norm_buf: CudaSlice<f32>,
    /// IQN gradient w.r.t. shared trunk h_s2 [B, hidden_dim].
    /// Accumulated during backward, added to C51's trunk gradient.
    d_h_s2_buf: CudaSlice<f32>,

    // ── Per-step buffers ─────────────────────────────────────────────
    /// Pre-sampled τ values for online network [B, N]
    online_taus: CudaSlice<f32>,
    /// Pre-sampled τ values for target network [B, N]
    target_taus: CudaSlice<f32>,
    /// Branch actions decoded from flat actions [B, 3]
    branch_actions: CudaSlice<i32>,
    /// Target h_s2 computed from next_states + target trunk weights [B, H]
    target_h_s2: CudaSlice<f32>,
    /// Persistent rewards buffer [B] — avoids per-step temp allocs
    rewards_buf: CudaSlice<f32>,
    /// Persistent dones buffer [B] — avoids per-step temp allocs
    dones_buf: CudaSlice<f32>,

    // ── Activation save buffers (forward → backward) ─────────────────
    save_embed: CudaSlice<f32>,     // [B, N, H]
    save_combined: CudaSlice<f32>,  // [B, N, H]
    save_q_online: CudaSlice<f32>,  // [B, N, TOTAL_BRANCH_ACTIONS]
    save_q_target: CudaSlice<f32>,  // [B, N, TOTAL_BRANCH_ACTIONS]

    // ── Loss buffers ─────────────────────────────────────────────────
    per_sample_loss: CudaSlice<f32>, // [B]
    total_loss: CudaSlice<f32>,      // [1]

    // ── Training state ───────────────────────────────────────────────
    adam_step: i32,
    /// Monotonic step counter for Philox PRNG seeding (τ sampling).
    rng_step: u32,
    total_params: usize,
}

impl GpuIqnHead {
    /// Create a new GPU IQN dual-head trainer.
    ///
    /// Compiles all 7 CUDA kernels, initializes online + target weights with
    /// Xavier/Glorot, and pre-allocates all GPU buffers.
    pub fn new(
        stream: Arc<CudaStream>,
        config: GpuIqnConfig,
    ) -> Result<Self, MLError> {
        let total_params = config.total_params();
        let b = config.batch_size;
        let n = config.num_quantiles;
        let h = config.hidden_dim;
        let tba = config.total_branch_actions();

        // Compile all 9 kernels
        let (trunk_fwd, fwd_loss, bwd, gnorm, adam, fwd_only, ema, decode_act, sample_taus_k) =
            compile_iqn_kernels(&stream, &config)?;

        // Allocate and initialize weights (Xavier/Glorot)
        let online_params = init_iqn_xavier_weights(&stream, &config)?;
        // Target params start as a copy of online params
        let target_params = clone_cuda_slice(&stream, &online_params, total_params)?;

        // Adam state (zeroed)
        let m_buf = alloc_f32(&stream, total_params, "iqn_m")?;
        let v_buf = alloc_f32(&stream, total_params, "iqn_v")?;
        let grad_buf = alloc_f32(&stream, total_params, "iqn_grad")?;
        let grad_norm_buf = alloc_f32(&stream, 1, "iqn_grad_norm")?;
        let d_h_s2_buf = alloc_f32(&stream, b * h, "iqn_d_h_s2")?;

        // Per-step buffers — fixed τ midpoints (QR-DQN style).
        // τ_i = (2i - 1) / (2N) for i = 1..N. Deterministic → CUDA Graph compatible.
        // No per-step random sampling. IQN cosine embedding still operates on these.
        let mut online_taus = alloc_f32(&stream, b * n, "iqn_online_taus")?;
        let mut target_taus = alloc_f32(&stream, b * n, "iqn_target_taus")?;
        {
            let fixed: Vec<f32> = (0..n)
                .map(|i| (2.0 * (i as f32) + 1.0) / (2.0 * n as f32))
                .collect();
            let mut tiled = Vec::with_capacity(b * n);
            for _ in 0..b {
                tiled.extend_from_slice(&fixed);
            }
            stream.memcpy_htod(&tiled, &mut online_taus)
                .map_err(|e| MLError::ModelError(format!("IQN fixed tau upload: {e}")))?;
            stream.memcpy_htod(&tiled, &mut target_taus)
                .map_err(|e| MLError::ModelError(format!("IQN fixed tau upload: {e}")))?;
        }
        let branch_actions = stream.alloc_zeros::<i32>(b * 3).map_err(|e| {
            MLError::ModelError(format!("IQN alloc branch_actions: {e}"))
        })?;
        let target_h_s2 = alloc_f32(&stream, b * h, "iqn_target_h_s2")?;

        // Activation saves
        let save_embed = alloc_f32(&stream, b * n * h, "iqn_save_embed")?;
        let save_combined = alloc_f32(&stream, b * n * h, "iqn_save_combined")?;
        let save_q_online = alloc_f32(&stream, b * n * tba, "iqn_save_q_online")?;
        let save_q_target = alloc_f32(&stream, b * n * tba, "iqn_save_q_target")?;

        // Loss buffers
        let per_sample_loss = alloc_f32(&stream, b, "iqn_per_sample_loss")?;
        let total_loss = alloc_f32(&stream, 1, "iqn_total_loss")?;

        // Persistent rewards/dones buffers (avoids per-step temp allocs)
        let rewards_buf = alloc_f32(&stream, b, "iqn_rewards")?;
        let dones_buf = alloc_f32(&stream, b, "iqn_dones")?;

        let vram_bytes = (total_params * 6 + b * n * 2 + b * 3
            + b * h + b * n * h * 2 + b * n * tba * 2
            + b * 3 + 2) * 4;

        info!(
            hidden_dim = h,
            embed_dim = config.embed_dim,
            num_quantiles = n,
            kappa = config.kappa,
            total_params,
            total_branch_actions = tba,
            vram_kb = vram_bytes / 1024,
            "GpuIqnHead initialized: 9 kernels compiled, Xavier weights"
        );

        Ok(Self {
            config,
            stream,
            trunk_forward_kernel: trunk_fwd,
            forward_loss_kernel: fwd_loss,
            backward_kernel: bwd,
            grad_norm_kernel: gnorm,
            adam_kernel: adam,
            forward_kernel: fwd_only,
            ema_kernel: ema,
            decode_actions_kernel: decode_act,
            sample_taus_kernel: sample_taus_k,
            online_params,
            target_params,
            m_buf,
            v_buf,
            grad_buf,
            grad_norm_buf,
            d_h_s2_buf,
            online_taus,
            target_taus,
            branch_actions,
            target_h_s2,
            rewards_buf,
            dones_buf,
            save_embed,
            save_combined,
            save_q_online,
            save_q_target,
            per_sample_loss,
            total_loss,
            adam_step: 0,
            rng_step: 0,
            total_params,
        })
    }

    /// Run one GPU-native IQN training step.
    ///
    /// All data sourced from DQN trainer's GPU buffers via DtoD copies.
    /// No CPU arrays, no HtoD uploads for batch data.
    ///
    /// Full pipeline:
    /// 1. Sample τ values and upload (small, ~32KB)
    /// 2. Decode flat actions → branch actions on GPU via decode kernel
    /// 3. DtoD copy rewards/dones from DQN trainer buffers
    /// 4. Compute target h_s2 via trunk forward kernel
    /// 5. Forward + loss
    /// 6. Backward + grad norm + Adam
    ///
    /// Returns the batch-mean quantile Huber loss.
    pub fn train_iqn_step_gpu(
        &mut self,
        online_h_s2: &CudaSlice<f32>,
        next_states_buf: &CudaSlice<f32>,
        target_dueling: &DuelingWeightSet,
        dqn_actions_buf: &CudaSlice<i32>,
        dqn_rewards_buf: &CudaSlice<f32>,
        dqn_dones_buf: &CudaSlice<f32>,
    ) -> Result<f32, MLError> {
        let b = self.config.batch_size;

        // Layer 3a: τ values are fixed midpoints, pre-computed in constructor.
        // No per-step sampling → deterministic → CUDA Graph compatible.

        // 2. GPU decode: flat actions [B] → branch actions [B, 3]
        let decode_blocks = (b + 255) / 256;
        let decode_config = LaunchConfig {
            grid_dim: (decode_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        let batch_size_i32 = b as i32;

        // Safety: dqn_actions_buf[B] i32, branch_actions[B*3] i32, same context.
        unsafe {
            self.stream
                .launch_builder(&self.decode_actions_kernel)
                .arg(dqn_actions_buf)
                .arg(&mut self.branch_actions)
                .arg(&batch_size_i32)
                .launch(decode_config)
                .map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?;
        }

        // 3. DtoD copy rewards/dones from DQN trainer buffers
        use super::gpu_dqn_trainer::{raw_device_ptr, dtod_copy};
        let f32_bytes = std::mem::size_of::<f32>();
        let r_src = raw_device_ptr(dqn_rewards_buf, &self.stream);
        let r_dst = raw_device_ptr(&self.rewards_buf, &self.stream);
        dtod_copy(r_dst, r_src, b * f32_bytes, &self.stream, 0, "iqn_rewards_dtod")?;

        let d_src = raw_device_ptr(dqn_dones_buf, &self.stream);
        let d_dst = raw_device_ptr(&self.dones_buf, &self.stream);
        dtod_copy(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?;

        // 4-9. Execute core training pipeline (shared logic)
        self.execute_training_pipeline(online_h_s2, next_states_buf, target_dueling)
    }

    /// Core IQN training pipeline: trunk forward → forward+loss → backward → Adam.
    ///
    /// Assumes `branch_actions`, `rewards_buf`, `dones_buf` are already populated
    /// on GPU (by either `train_iqn_step_gpu` or the legacy HtoD path).
    fn execute_training_pipeline(
        &mut self,
        online_h_s2: &CudaSlice<f32>,
        next_states_buf: &CudaSlice<f32>,
        target_dueling: &DuelingWeightSet,
    ) -> Result<f32, MLError> {
        let b = self.config.batch_size;

        // 4. Compute target h_s2 via trunk forward kernel
        let shmem_bytes = self.config.shared_h1 * 4;
        let trunk_config = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: shmem_bytes as u32,
        };

        // Safety: all pointers are valid GPU allocations on the same context.
        // next_states_buf points to the DQN trainer's pre-allocated buffer.
        // target_dueling weights are extracted from VarMap CudaSlice buffers.
        unsafe {
            self.stream
                .launch_builder(&self.trunk_forward_kernel)
                .arg(next_states_buf)
                .arg(&target_dueling.w_s1)
                .arg(&target_dueling.b_s1)
                .arg(&target_dueling.w_s2)
                .arg(&target_dueling.b_s2)
                .arg(&mut self.target_h_s2)
                .arg(&(b as i32))
                .launch(trunk_config)
                .map_err(|e| MLError::ModelError(format!("IQN trunk forward kernel: {e}")))?;
        }

        // 5. Zero total_loss, grad_norm, and gradient buffer on GPU (memset, no CPU alloc)
        self.stream.memset_zeros(&mut self.total_loss)
            .map_err(|e| MLError::ModelError(format!("IQN zero total_loss: {e}")))?;
        self.stream.memset_zeros(&mut self.grad_norm_buf)
            .map_err(|e| MLError::ModelError(format!("IQN zero grad_norm: {e}")))?;
        self.stream.memset_zeros(&mut self.grad_buf)
            .map_err(|e| MLError::ModelError(format!("IQN zero grads: {e}")))?;
        self.stream.memset_zeros(&mut self.d_h_s2_buf)
            .map_err(|e| MLError::ModelError(format!("IQN zero d_h_s2: {e}")))?;

        let batch_size_i32 = b as i32;
        let gamma = self.config.gamma;

        // 6. Forward + loss kernel: one warp per sample
        let fwd_config = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: 0,
        };

        // Safety: all buffer pointers are valid GPU allocations on the same context
        // and stream. online_h_s2 points to the DQN trainer's save_h_s2 buffer.
        unsafe {
            self.stream
                .launch_builder(&self.forward_loss_kernel)
                .arg(online_h_s2)
                .arg(&self.target_h_s2)
                .arg(&self.online_taus)
                .arg(&self.target_taus)
                .arg(&self.branch_actions)
                .arg(&self.rewards_buf)
                .arg(&self.dones_buf)
                .arg(&gamma)
                .arg(&self.online_params)
                .arg(&self.target_params)
                .arg(&mut self.per_sample_loss)
                .arg(&mut self.total_loss)
                .arg(&mut self.save_embed)
                .arg(&mut self.save_combined)
                .arg(&mut self.save_q_online)
                .arg(&mut self.save_q_target)
                .arg(&batch_size_i32)
                .launch(fwd_config)
                .map_err(|e| MLError::ModelError(format!("IQN forward+loss kernel: {e}")))?;
        }

        // 7. Backward kernel: one warp per sample
        let bwd_config = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: 0,
        };

        // Safety: all pointers valid, save buffers written by forward kernel on same stream.
        unsafe {
            self.stream
                .launch_builder(&self.backward_kernel)
                .arg(online_h_s2)
                .arg(&self.save_embed)
                .arg(&self.save_combined)
                .arg(&self.save_q_online)
                .arg(&self.save_q_target)
                .arg(&self.online_taus)
                .arg(&self.branch_actions)
                .arg(&self.online_params)
                .arg(&mut self.grad_buf)
                .arg(&mut self.d_h_s2_buf) // IQN trunk gradient output
                .arg(&batch_size_i32)
                .launch(bwd_config)
                .map_err(|e| MLError::ModelError(format!("IQN backward kernel: {e}")))?;
        }

        // 8. Grad norm kernel
        let total_params_i32 = self.total_params as i32;
        let norm_blocks = (self.total_params + 255) / 256;
        let norm_config = LaunchConfig {
            grid_dim: (norm_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        // Safety: grad_buf has total_params elements, grad_norm_buf has 1.
        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_kernel)
                .arg(&self.grad_buf)
                .arg(&mut self.grad_norm_buf)
                .arg(&total_params_i32)
                .launch(norm_config)
                .map_err(|e| MLError::ModelError(format!("IQN grad_norm kernel: {e}")))?;
        }

        // 9. Adam update kernel
        self.adam_step += 1;
        let adam_blocks = (self.total_params + 255) / 256;
        let adam_config = LaunchConfig {
            grid_dim: (adam_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let lr = self.config.lr;
        let beta1 = self.config.beta1;
        let beta2 = self.config.beta2;
        let eps = self.config.epsilon;
        let wd = self.config.weight_decay;
        let mgn = self.config.max_grad_norm;
        let t = self.adam_step;

        // Safety: all buffers have total_params elements on the same context.
        unsafe {
            self.stream
                .launch_builder(&self.adam_kernel)
                .arg(&mut self.online_params)
                .arg(&mut self.grad_buf)
                .arg(&mut self.m_buf)
                .arg(&mut self.v_buf)
                .arg(&self.grad_norm_buf)
                .arg(&lr)
                .arg(&beta1)
                .arg(&beta2)
                .arg(&eps)
                .arg(&wd)
                .arg(&mgn)
                .arg(&t)
                .arg(&total_params_i32)
                .launch(adam_config)
                .map_err(|e| MLError::ModelError(format!("IQN adam kernel: {e}")))?;
        }

        // Loss stays on GPU — no per-step readback. Read only when needed via iqn_loss().
        // The old per-step DtoH sync was a hidden serialization point.
        Ok(0.0) // placeholder — actual loss available via read_loss()
    }

    /// Update target IQN weights via Polyak EMA.
    ///
    /// `target[i] = (1 - tau) * target[i] + tau * online[i]`
    pub fn target_ema_update(&mut self, tau: f32) -> Result<(), MLError> {
        let n = self.total_params;
        let blocks = (n + 255) / 256;
        let config = LaunchConfig {
            grid_dim: (blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let n_i32 = n as i32;

        // Safety: target_params and online_params have total_params elements each.
        unsafe {
            self.stream
                .launch_builder(&self.ema_kernel)
                .arg(&mut self.target_params)
                .arg(&self.online_params)
                .arg(&tau)
                .arg(&n_i32)
                .launch(config)
                .map_err(|e| MLError::ModelError(format!("IQN EMA kernel: {e}")))?;
        }

        Ok(())
    }

    /// Current Adam step count.
    pub fn adam_step(&self) -> i32 {
        self.adam_step
    }

    /// Per-sample IQN loss buffer reference (for PER weighting).
    pub fn per_sample_loss(&self) -> &CudaSlice<f32> {
        &self.per_sample_loss
    }

    /// IQN gradient w.r.t. shared trunk h_s2 [B, hidden_dim].
    /// After backward, this contains dL_iqn/d(h_s2) — add to C51's trunk gradient
    /// so the trunk receives bounded IQN Huber gradients alongside C51.
    pub fn d_h_s2(&self) -> &CudaSlice<f32> {
        &self.d_h_s2_buf
    }

    /// Read IQN loss from GPU (call sparingly — synchronizes the stream).
    /// Use for epoch-end logging, NOT per-step monitoring.
    pub fn read_loss(&self) -> Result<f32, MLError> {
        let mut host = [0.0_f32];
        self.stream.memcpy_dtoh(&self.total_loss, &mut host)
            .map_err(|e| MLError::ModelError(format!("IQN loss readback: {e}")))?;
        Ok(host[0])
    }

    /// Compute CVaR-based position scaling from IQN quantiles.
    ///
    /// Runs the IQN forward-only kernel on `h_s2` (trunk activation),
    /// then computes CVaR at `alpha` (e.g. 0.05 = 5th percentile) for
    /// each sample's selected exposure action.
    ///
    /// Returns a GPU buffer `[batch_size]` with scaling factors in [0.25, 1.0]:
    /// - CVaR ≥ 0 (expected profit even in worst case) → scale = 1.0
    /// - CVaR < 0 (expected loss in worst case) → scale = max(0.25, 1.0 + CVaR * 5.0)
    ///
    /// The scale is applied to target_position in the env_step kernel.
    pub fn compute_cvar_scales(
        &mut self,
        h_s2: &CudaSlice<f32>,
        actions: &CudaSlice<i32>,
        batch_size: usize,
        alpha: f32,
    ) -> Result<CudaSlice<f32>, MLError> {
        let b = batch_size;
        let n = self.config.num_quantiles;
        let tba = self.config.total_branch_actions();

        // 1. Sample τ values for inference
        let total_taus = (b * n) as i32;
        let tau_blocks = (b * n + 255) / 256;
        let tau_config = LaunchConfig {
            grid_dim: (tau_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        let rng_step = self.rng_step as u32;
        self.rng_step += 1;
        unsafe {
            self.stream
                .launch_builder(&self.sample_taus_kernel)
                .arg(&mut self.online_taus)
                .arg(&total_taus)
                .arg(&rng_step)
                .launch(tau_config)
                .map_err(|e| MLError::ModelError(format!("IQN CVaR sample_taus: {e}")))?;
        }

        // 2. Run IQN forward-only kernel
        let fwd_config = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: 0,
        };
        let batch_i32 = b as i32;
        unsafe {
            self.stream
                .launch_builder(&self.forward_kernel)
                .arg(h_s2)
                .arg(&self.online_taus)
                .arg(&self.online_params)
                .arg(&mut self.save_q_online)
                .arg(&batch_i32)
                .launch(fwd_config)
                .map_err(|e| MLError::ModelError(format!("IQN CVaR forward: {e}")))?;
        }

        // 3. Compute CVaR on GPU — ZERO CPU readback.
        // Inline CUDA kernel: one thread per sample, insertion-sort 32 quantiles
        // in registers, compute CVaR as mean of lowest alpha_count values.
        let mut scales_buf = self.stream.alloc_zeros::<f32>(b)
            .map_err(|e| MLError::ModelError(format!("IQN CVaR alloc scales: {e}")))?;
        let alpha_f = alpha;
        let n_i32 = n as i32;
        let tba_i32 = tba as i32;
        let b0_i32 = self.config.branch_0_size as i32;
        let b1b2 = (self.config.branch_1_size * self.config.branch_2_size) as i32;
        let cvar_blocks = (b + 255) / 256;
        let cvar_config = LaunchConfig {
            grid_dim: (cvar_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        // Compile the CVaR kernel inline (cached after first call)
        static CVAR_KERNEL_SRC: &str = r#"
extern "C" __global__ void iqn_cvar_kernel(
    const float* __restrict__ q_values,  // [B, N_TAU, TBA]
    const int*   __restrict__ actions,   // [B] factored action indices
    float*       scales_out,             // [B] output CVaR scales
    int B, int N_TAU, int TBA, int B0_SIZE, int B1B2, float ALPHA
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= B) return;

    // Decode exposure index from factored action
    int exposure_idx = actions[i] / B1B2;
    if (exposure_idx >= B0_SIZE) exposure_idx = B0_SIZE - 1;

    // Extract quantile values for this sample's exposure action
    // q_values layout: [B, N_TAU, TBA]
    int alpha_count = (int)(ALPHA * (float)N_TAU);
    if (alpha_count < 1) alpha_count = 1;

    // Find the alpha_count smallest quantile values using partial sort.
    // For N_TAU=32 and alpha_count=1-2, just find the minimum(s).
    float cvar_sum = 0.0f;
    if (alpha_count == 1) {
        // Fast path: just find the minimum
        float min_val = 1e30f;
        for (int t = 0; t < N_TAU; t++) {
            float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx];
            if (v < min_val) min_val = v;
        }
        cvar_sum = min_val;
    } else {
        // General path: find alpha_count smallest values
        // Use insertion into a small sorted array (alpha_count <= ~5 for typical alpha=0.05)
        float sorted[8]; // max alpha_count for alpha=0.25, N_TAU=32
        int sorted_len = 0;
        for (int t = 0; t < N_TAU; t++) {
            float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx];
            if (sorted_len < alpha_count) {
                // Insert into sorted array
                int pos = sorted_len;
                while (pos > 0 && sorted[pos-1] > v) {
                    sorted[pos] = sorted[pos-1];
                    pos--;
                }
                sorted[pos] = v;
                sorted_len++;
            } else if (v < sorted[sorted_len-1]) {
                // Replace the largest in sorted
                int pos = sorted_len - 1;
                while (pos > 0 && sorted[pos-1] > v) {
                    sorted[pos] = sorted[pos-1];
                    pos--;
                }
                sorted[pos] = v;
            }
        }
        for (int j = 0; j < alpha_count; j++) cvar_sum += sorted[j];
    }

    float cvar = cvar_sum / (float)alpha_count;

    // Scale: [0.25, 1.0]. CVaR >= 0 -> full size. CVaR < 0 -> reduce.
    float scale = (cvar >= 0.0f) ? 1.0f : fmaxf(0.25f, 1.0f + cvar * 5.0f);
    scales_out[i] = scale;
}
"#;
        // Compile once, cache in OnceLock
        use std::sync::OnceLock;
        static CVAR_PTX: OnceLock<Result<cudarc::nvrtc::Ptx, String>> = OnceLock::new();
        let context = self.stream.context();
        let ptx = CVAR_PTX.get_or_init(|| {
            crate::cuda_pipeline::compile_ptx_for_device(CVAR_KERNEL_SRC, &context)
        });
        let ptx = ptx.as_ref().map_err(|e| MLError::ModelError(format!("CVaR kernel PTX: {e}")))?;
        let module = context.load_module(ptx.clone())
            .map_err(|e| MLError::ModelError(format!("CVaR kernel module: {e}")))?;
        let cvar_kernel = module.load_function("iqn_cvar_kernel")
            .map_err(|e| MLError::ModelError(format!("CVaR kernel load: {e}")))?;

        unsafe {
            self.stream
                .launch_builder(&cvar_kernel)
                .arg(&self.save_q_online)
                .arg(actions)
                .arg(&mut scales_buf)
                .arg(&batch_i32)
                .arg(&n_i32)
                .arg(&tba_i32)
                .arg(&b0_i32)
                .arg(&b1b2)
                .arg(&alpha_f)
                .launch(cvar_config)
                .map_err(|e| MLError::ModelError(format!("CVaR kernel launch: {e}")))?;
        }

        Ok(scales_buf)
    }
}

// ---------------------------------------------------------------------------
// Kernel compilation
// ---------------------------------------------------------------------------

/// Compile all 9 IQN CUDA kernels with dimension defines.
fn compile_iqn_kernels(
    stream: &Arc<CudaStream>,
    config: &GpuIqnConfig,
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
    let dim_overrides = format!(
        "#define IQN_HIDDEN {hidden}\n\
         #define IQN_EMBED_DIM {embed}\n\
         #define IQN_NUM_QUANTILES {nq}\n\
         #define IQN_KAPPA {kappa:.6}f\n\
         #define BRANCH_0_SIZE {b0}\n\
         #define BRANCH_1_SIZE {b1}\n\
         #define BRANCH_2_SIZE {b2}\n\
         #define IQN_STATE_DIM {sd}\n\
         #define IQN_SHARED_H1 {h1}\n",
        hidden = config.hidden_dim,
        embed = config.embed_dim,
        nq = config.num_quantiles,
        kappa = config.kappa,
        b0 = config.branch_0_size,
        b1 = config.branch_1_size,
        b2 = config.branch_2_size,
        sd = config.state_dim,
        h1 = config.shared_h1,
    );

    let kernel_src = include_str!("iqn_dual_head_kernel.cu");
    let full_source = format!("{dim_overrides}\n{kernel_src}");

    info!(
        hidden = config.hidden_dim,
        embed_dim = config.embed_dim,
        num_quantiles = config.num_quantiles,
        kappa = config.kappa,
        total_params = config.total_params(),
        "GpuIqnHead: compiling 9 IQN kernels"
    );

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

    let load = |name: &str| -> Result<CudaFunction, MLError> {
        module.load_function(name).map_err(|e| {
            MLError::ModelError(format!("{name} load: {e}"))
        })
    };

    let trunk_fwd = load("iqn_trunk_forward_kernel")?;
    let fwd_loss = load("iqn_forward_loss_kernel")?;
    let backward = load("iqn_backward_kernel")?;
    let grad_norm = load("iqn_grad_norm_kernel")?;
    let adam = load("iqn_adam_kernel")?;
    let fwd_only = load("iqn_forward_kernel")?;
    let ema = load("iqn_ema_kernel")?;
    let decode_act = load("iqn_decode_actions_kernel")?;
    let sample_taus = load("iqn_sample_taus_kernel")?;

    info!("GpuIqnHead: 9 kernels compiled and loaded");
    Ok((trunk_fwd, fwd_loss, backward, grad_norm, adam, fwd_only, ema, decode_act, sample_taus))
}

// ---------------------------------------------------------------------------
// Weight initialization
// ---------------------------------------------------------------------------

/// Initialize IQN weights with Xavier/Glorot uniform distribution.
fn init_iqn_xavier_weights(
    stream: &Arc<CudaStream>,
    config: &GpuIqnConfig,
) -> Result<CudaSlice<f32>, MLError> {
    use rand::Rng;

    let total = config.total_params();
    let h = config.hidden_dim;
    let d = config.embed_dim;
    let b0 = config.branch_0_size;
    let b1 = config.branch_1_size;
    let b2 = config.branch_2_size;
    let mut weights = vec![0.0_f32; total];
    let mut rng = rand::thread_rng();
    let mut offset = 0;

    // W_embed [H, D]
    let limit = (6.0_f64 / (h + d) as f64).sqrt() as f32;
    for w in &mut weights[offset..offset + h * d] {
        *w = rng.gen_range(-limit..limit);
    }
    offset += h * d;
    // b_embed [H] — zero
    offset += h;

    // Branch output layers
    for branch_size in [b0, b1, b2] {
        let limit = (6.0_f64 / (branch_size + h) as f64).sqrt() as f32;
        // W_bd [branch_size, H]
        for w in &mut weights[offset..offset + branch_size * h] {
            *w = rng.gen_range(-limit..limit);
        }
        offset += branch_size * h;
        // b_bd [branch_size] — zero
        offset += branch_size;
    }

    debug_assert_eq!(offset, total, "IQN weight layout mismatch");

    // Upload to GPU
    let mut params_buf = alloc_f32(stream, total, "iqn_online_params")?;
    stream.memcpy_htod(&weights, &mut params_buf)
        .map_err(|e| MLError::ModelError(format!("IQN weight upload: {e}")))?;

    Ok(params_buf)
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Allocate a zeroed f32 buffer on GPU.
fn alloc_f32(
    stream: &Arc<CudaStream>,
    n: usize,
    name: &str,
) -> Result<CudaSlice<f32>, MLError> {
    stream.alloc_zeros::<f32>(n).map_err(|e| {
        MLError::ModelError(format!("GpuIqn alloc {name} ({n} f32): {e}"))
    })
}

/// Clone a CudaSlice by copying device → host → device.
fn clone_cuda_slice(
    stream: &Arc<CudaStream>,
    src: &CudaSlice<f32>,
    n: usize,
) -> Result<CudaSlice<f32>, MLError> {
    let mut host = vec![0.0_f32; n];
    stream.memcpy_dtoh(src, &mut host) // gpu-exit: weight clone for target network
        .map_err(|e| MLError::ModelError(format!("IQN clone D→H: {e}")))?;
    let mut dst = alloc_f32(stream, n, "iqn_target_params")?;
    stream.memcpy_htod(&host, &mut dst)
        .map_err(|e| MLError::ModelError(format!("IQN clone H→D: {e}")))?;
    Ok(dst)
}

