#![allow(unsafe_code)] // Required for CUDA DtoD memcpy via cudarc.

//! Branching Dueling Q-Network (BDQ) for factored action spaces.
//!
//! Implements the architecture from "Action Branching Architectures for Deep
//! Reinforcement Learning" (Tavakoli et al., 2018).
//!
//! ## Architecture
//!
//! ```text
//! State [state_dim] -> Shared Encoder [hidden_dims]
//!                          |
//!            +-------------+-------------+
//!            v             v             v
//!       Value V(s)    Advantage_0    Advantage_1  ... Advantage_D
//!       [1 scalar]    [n_0 actions]  [n_1 actions]    [n_D actions]
//!            v             v             v              v
//!       Q(s,a) = V(s) + (1/D) * sum_d [A_d(s,a_d) - mean(A_d)]
//! ```
//!
//! ## Branch Layout (Trading)
//!
//! - **Branch 0 (Exposure)**: 5 outputs -- Short100, Short50, Flat, Long50, Long100
//! - **Branch 1 (Order Type)**: 3 outputs -- Market, `LimitMaker`, `IoC`
//! - **Branch 2 (Urgency)**: 3 outputs -- Patient, Normal, Aggressive
//!
//! ## Key Properties
//!
//! - Linear growth in outputs: `sum(n_d)` instead of `product(n_d)` (11 vs 45)
//! - Each branch independently selectable via argmax
//! - Shared value V(s) amortized across branches
//! - Mean advantage subtraction per branch (identifiability)
//!
//! ## Extensions
//!
//! - **Per-Branch C51 Distributional**: Each branch head outputs atom distributions
//!   instead of scalar Q-values. Expected Q = sum(softmax(logits) * z).
//! - **`NoisyNet`**: Factorized Gaussian noise in value/branch heads for learned exploration.
//! - **State Dim Alignment**: Auto-pad `state_dim` to multiples of 8 for tensor core HMMA.

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

use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
use ml_core::cuda_autograd::{GpuLinear, GpuTensor, GpuVarStore};
use serde::{Deserialize, Serialize};

use crate::noisy_layers::NoisyLinear;
use ml_core::MLError;

/// Clone a sub-range of a `CudaSlice` into a fresh owned `CudaSlice` via D2D copy.
///
/// This avoids the `CudaView` lifetime issue when passing slices to kernels
/// that require `&CudaSlice`. Zero CPU involvement.
fn clone_subrange(
    src: &CudaSlice<f32>,
    offset: usize,
    count: usize,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let mut dst = stream.alloc_zeros::<f32>(count).map_err(|e| {
        MLError::ModelError(format!("clone_subrange alloc({count}): {e}"))
    })?;
    let src_view = src.slice(offset..offset + count);
    let num_bytes = count * std::mem::size_of::<f32>();
    let src_ptr = {
        let (ptr, guard) = src_view.device_ptr(stream);
        let _no_drop = ManuallyDrop::new(guard);
        ptr
    };
    let dst_ptr = {
        let (ptr, guard) = dst.device_ptr_mut(stream);
        let _no_drop = ManuallyDrop::new(guard);
        ptr
    };
    // SAFETY: src and dst are valid device allocations. num_bytes <= both alloc sizes.
    unsafe {
        cudarc::driver::result::memcpy_dtod_async(
            dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
        ).map_err(|e| MLError::ModelError(format!("clone_subrange D2D: {e}")))?;
    }
    Ok(dst)
}

/// Async device-to-device memcpy for a single `CudaSlice` buffer. Zero CPU involvement.
fn dtod_copy_slice(
    src: &CudaSlice<f32>,
    dst: &mut CudaSlice<f32>,
    stream: &Arc<CudaStream>,
    label: &str,
) -> Result<(), MLError> {
    let num_bytes = src.len() * std::mem::size_of::<f32>();
    let src_ptr = {
        let (ptr, guard) = src.device_ptr(stream);
        let _no_drop = ManuallyDrop::new(guard);
        ptr
    };
    let dst_ptr = {
        let (ptr, guard) = dst.device_ptr_mut(stream);
        let _no_drop = ManuallyDrop::new(guard);
        ptr
    };
    // SAFETY: src and dst are valid device allocations on the same context.
    // num_bytes = len * sizeof(f32) does not exceed either allocation.
    unsafe {
        cudarc::driver::result::memcpy_dtod_async(
            dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
        ).map_err(|e| MLError::ModelError(format!("DtoD {label}: {e}")))?;
    }
    Ok(())
}

/// Output of the branching network's forward pass.
///
/// Provides the decomposed value + per-branch advantages
/// for both action selection and loss computation.
#[derive(Debug)]
pub struct BranchOutput {
    /// State value V(s): [batch, 1]
    pub value: GpuTensor,
    /// Per-branch advantage tensors `A_d(s`, .): [batch, `n_d`] for each branch d.
    /// Contains expected Q-values (sum of softmax(logits) * z) for greedy action selection.
    pub advantages: Vec<GpuTensor>,
    /// Per-branch log-softmax distributions: [batch, `n_d`, `num_atoms`] for each branch d.
    pub advantage_log_probs: Option<Vec<GpuTensor>>,
    /// Value stream distribution: [batch, 1, `num_atoms`].
    pub value_log_probs: Option<GpuTensor>,
}

/// Configuration for Branching Dueling Q-Network.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BranchingConfig {
    /// Input state dimension (pre-aligned for tensor cores)
    pub state_dim: usize,
    /// Number of actions per branch, e.g. [5, 3, 3]
    pub branch_sizes: Vec<usize>,
    /// Shared feature extraction layer dimensions
    pub shared_hidden_dims: Vec<usize>,
    /// Value stream hidden dimension
    pub value_hidden_dim: usize,
    /// Per-branch advantage stream hidden dimension
    pub branch_hidden_dim: usize,
    /// `LeakyReLU` negative slope
    pub leaky_relu_alpha: f64,
    /// Dropout rate after shared layers (0.0 = disabled)
    #[serde(default)]
    pub dropout_rate: f64,

    // -- C51 Distributional (always enabled) --
    /// Number of categorical atoms for distributional value (C51).
    #[serde(default = "default_num_atoms")]
    pub num_atoms: usize,
    /// Minimum support value for C51 distribution.
    #[serde(default = "default_v_min")]
    pub v_min: f32,
    /// Maximum support value for C51 distribution.
    #[serde(default = "default_v_max")]
    pub v_max: f32,

    // -- NoisyNet (always enabled) --
    /// Initial noise std dev for `NoisyNet` (scaled by `1/sqrt(fan_in)` internally).
    #[serde(default = "default_noisy_sigma_init")]
    pub noisy_sigma_init: f32,
}

const fn default_num_atoms() -> usize {
    51
}
const fn default_v_min() -> f32 {
    -25.0
}
const fn default_v_max() -> f32 {
    25.0
}
const fn default_noisy_sigma_init() -> f32 {
    0.5
}

impl BranchingConfig {
    /// Trading-specific config with dynamic branch sizes.
    ///
    /// # Arguments
    ///
    /// * `state_dim` - Input state dimension
    /// * `shared_hidden_dims` - Shared feature extraction layers
    /// * `branch_hidden_dim` - Hidden dim for each branch head
    /// * `branch_sizes` - Actions per branch (e.g. `[5, 3, 3]` for exposure/order/urgency)
    pub const fn trading_default(
        state_dim: usize,
        shared_hidden_dims: Vec<usize>,
        branch_hidden_dim: usize,
        branch_sizes: Vec<usize>,
    ) -> Self {
        Self {
            state_dim,
            branch_sizes,
            shared_hidden_dims,
            value_hidden_dim: branch_hidden_dim,
            branch_hidden_dim,
            leaky_relu_alpha: 0.01,
            dropout_rate: 0.0,
            num_atoms: 51,
            v_min: -25.0,
            v_max: 25.0,
            noisy_sigma_init: 0.5,
        }
    }

    /// Create from DQN hyperparameters with dynamic branch sizes.
    ///
    /// `state_dim` is aligned to multiples of 8 for tensor core HMMA dispatch on CUDA
    /// when `align_for_gpu` is `Some(...)`. On CPU the dimension is unchanged.
    ///
    /// # Arguments
    ///
    /// * `branch_sizes` - Actions per branch (e.g. `[5, 3, 3]` for exposure/order/urgency)
    pub fn from_dqn_params(
        state_dim: usize,
        hidden_dims: &[usize],
        dueling_hidden_dim: usize,
        leaky_relu_alpha: f64,
        align_for_gpu: Option<&Arc<CudaStream>>,
        branch_sizes: Vec<usize>,
    ) -> Self {
        let aligned_state_dim = if align_for_gpu.is_some() {
            (state_dim + 7) & !7
        } else {
            state_dim
        };
        Self {
            state_dim: aligned_state_dim,
            branch_sizes,
            shared_hidden_dims: hidden_dims.to_vec(),
            value_hidden_dim: dueling_hidden_dim,
            branch_hidden_dim: dueling_hidden_dim,
            leaky_relu_alpha,
            dropout_rate: 0.0,
            num_atoms: 51,
            v_min: -25.0,
            v_max: 25.0,
            noisy_sigma_init: 0.5,
        }
    }

    /// Total number of branches (D).
    pub fn num_branches(&self) -> usize {
        self.branch_sizes.len()
    }
}

/// A `NoisyLinear` wrapper for value/branch heads (noisy nets always enabled).
enum MaybeNoisyLinear {
    Noisy(NoisyLinear),
}

impl MaybeNoisyLinear {
    /// Forward pass through the noisy layer.
    fn forward(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        let Self::Noisy(n) = self;
        n.forward(x)
    }

    /// Resample noise.
    fn reset_noise(&mut self) -> Result<(), MLError> {
        let Self::Noisy(n) = self;
        n.reset_noise()
    }

    /// Resample noise with custom sigma scale.
    fn reset_noise_with_sigma(&mut self, sigma_scale: f64) -> Result<(), MLError> {
        let Self::Noisy(n) = self;
        n.reset_noise_with_sigma(sigma_scale)
    }

    /// Disable noise for evaluation.
    fn disable_noise(&mut self) -> Result<(), MLError> {
        let Self::Noisy(n) = self;
        n.disable_noise()
    }

    /// Convert sigma vars and epsilon buffers to F32 to match mu vars.
    #[allow(dead_code)]
    const fn ensure_f32(&mut self) -> Result<(), MLError> {
        let Self::Noisy(n) = self;
        n.ensure_f32()
    }

    /// Register mu (weight/bias) params in `GpuVarStore` under `{name}.weight` / `{name}.bias`.
    ///
    /// Clones the mu data from the `NoisyLinear` into the `GpuVarStore` so that the
    /// GPU experience collector and serialization can find them by name.
    /// Uses async `DtoD` memcpy -- zero CPU involvement.
    fn register_mu_in_varstore(&self, vars: &mut GpuVarStore, name: &str) -> Result<(), MLError> {
        let Self::Noisy(n) = self;
        let [w_mu, b_mu] = n.mu_slices();
        let out_f = n.out_features();
        let in_f = n.in_features();
        let stream = vars.stream().clone();

        // Clone weight_mu -> "{name}.weight" via DtoD memcpy
        let w_len = w_mu.len();
        let mut w_clone = stream.alloc_zeros::<f32>(w_len).map_err(|e| {
            MLError::ModelError(format!("register_mu alloc {name}.weight: {e}"))
        })?;
        dtod_copy_slice(w_mu, &mut w_clone, &stream, &format!("register_mu {name}.weight"))?;
        vars.register(format!("{name}.weight"), w_clone, vec![out_f, in_f])?;

        // Clone bias_mu -> "{name}.bias" via DtoD memcpy
        let b_len = b_mu.len();
        let mut b_clone = stream.alloc_zeros::<f32>(b_len).map_err(|e| {
            MLError::ModelError(format!("register_mu alloc {name}.bias: {e}"))
        })?;
        dtod_copy_slice(b_mu, &mut b_clone, &stream, &format!("register_mu {name}.bias"))?;
        vars.register(format!("{name}.bias"), b_clone, vec![out_f])?;

        Ok(())
    }

    /// Get only sigma (noise std dev) parameter slices -- `weight_sigma`, `bias_sigma`.
    const fn noisy_sigma_slices(&self) -> [&CudaSlice<f32>; 2] {
        let Self::Noisy(n) = self;
        n.sigma_slices()
    }
}

impl std::fmt::Debug for MaybeNoisyLinear {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Noisy(NoisyLinear)")
    }
}

/// Branching Dueling Q-Network with independent advantage heads.
///
/// Extends the Dueling DQN architecture with D independent advantage streams,
/// one per action dimension. The value function V(s) is shared and amortized
/// across all branches.
///
/// Supports optional per-branch C51 distributional heads and `NoisyNet` exploration.
#[allow(missing_debug_implementations)]
pub struct BranchingDuelingQNetwork {
    /// Shared feature extraction layers
    shared_layers: Vec<GpuLinear>,

    /// Value stream: hidden -> scalar (or `num_atoms` when distributional)
    value_fc: MaybeNoisyLinear,
    value_out: MaybeNoisyLinear,

    /// Per-branch advantage streams: D x (hidden -> `n_d` or `n_d`*`num_atoms`)
    branch_fcs: Vec<MaybeNoisyLinear>,
    branch_outs: Vec<MaybeNoisyLinear>,

    /// Configuration
    config: BranchingConfig,

    /// Dropout rate for shared layers (0.0 = disabled)
    _dropout_rate: f32,

    /// Weight storage for serialization and optimizer
    vars: GpuVarStore,

    /// CUDA stream for GPU operations
    stream: Arc<CudaStream>,

    /// Pre-computed C51 support atoms z = `linspace(v_min`, `v_max`, `num_atoms`).
    support: Option<GpuTensor>,
}

impl BranchingDuelingQNetwork {
    /// Create a new branching dueling Q-network.
    ///
    /// All layers use Xavier initialization for stable gradient flow.
    /// Value and branch heads use factorized `NoisyLinear` layers (always enabled).
    /// Output dims are scaled by `num_atoms` (distributional always enabled).
    pub fn new(config: BranchingConfig, stream: Arc<CudaStream>) -> Result<Self, MLError> {
        if config.branch_sizes.is_empty() {
            return Err(MLError::InvalidInput(
                "BranchingDuelingQNetwork requires at least one branch".to_owned(),
            ));
        }

        let mut vars = GpuVarStore::new(stream.clone());
        let sigma_init = f64::from(config.noisy_sigma_init);

        // Build shared feature layers (GpuLinear in GpuVarStore)
        let mut shared_layers = Vec::new();
        let mut current_dim = config.state_dim;
        for (i, &hidden_dim) in config.shared_hidden_dims.iter().enumerate() {
            let layer = vars.linear_xavier(&format!("shared_{i}"), current_dim, hidden_dim)?;
            shared_layers.push(layer);
            current_dim = hidden_dim;
        }
        let final_hidden = current_dim;

        // Value stream: NoisyLinear hidden + NoisyLinear output
        // Output dims scaled by num_atoms (distributional always enabled)
        let value_fc = Self::build_head_layer(
            final_hidden, config.value_hidden_dim, &stream, "value_fc", sigma_init,
        )?;
        let value_out = Self::build_head_layer(
            config.value_hidden_dim, config.num_atoms, &stream, "value_out", sigma_init,
        )?;

        // Per-branch advantage streams
        let mut branch_fcs = Vec::with_capacity(config.branch_sizes.len());
        let mut branch_outs = Vec::with_capacity(config.branch_sizes.len());

        for (d, &n_d) in config.branch_sizes.iter().enumerate() {
            let fc = Self::build_head_layer(
                final_hidden, config.branch_hidden_dim, &stream,
                &format!("branch_{d}_fc"), sigma_init,
            )?;
            let out = Self::build_head_layer(
                config.branch_hidden_dim, n_d * config.num_atoms, &stream,
                &format!("branch_{d}_out"), sigma_init,
            )?;
            branch_fcs.push(fc);
            branch_outs.push(out);
        }

        // Register NoisyLinear mu params in GpuVarStore for GPU weight extraction
        value_fc.register_mu_in_varstore(&mut vars, "value_fc")?;
        value_out.register_mu_in_varstore(&mut vars, "value_out")?;
        for (d, fc) in branch_fcs.iter().enumerate() {
            fc.register_mu_in_varstore(&mut vars, &format!("branch_{d}_fc"))?;
        }
        for (d, out) in branch_outs.iter().enumerate() {
            out.register_mu_in_varstore(&mut vars, &format!("branch_{d}_out"))?;
        }

        // Compute C51 support atoms
        let support = Some(Self::support_atoms(
            config.v_min, config.v_max, config.num_atoms, &stream,
        )?);

        let dropout_rate_val = config.dropout_rate as f32;

        Ok(Self {
            shared_layers,
            value_fc,
            value_out,
            branch_fcs,
            branch_outs,
            config,
            _dropout_rate: dropout_rate_val,
            vars,
            stream,
            support,
        })
    }

    /// Build a `NoisyLinear` layer for value/branch heads.
    fn build_head_layer(
        fan_in: usize,
        fan_out: usize,
        stream: &Arc<CudaStream>,
        _name: &str,
        sigma_init: f64,
    ) -> Result<MaybeNoisyLinear, MLError> {
        let noisy = NoisyLinear::new(fan_in, fan_out, stream.clone(), sigma_init)?;
        Ok(MaybeNoisyLinear::Noisy(noisy))
    }

    /// Compute C51 support atoms: `linspace(v_min, v_max, num_atoms)` on device.
    pub fn support_atoms(
        v_min: f32,
        v_max: f32,
        num_atoms: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        if num_atoms < 2 {
            return Err(MLError::InvalidInput(
                "num_atoms must be >= 2 for C51 support".to_owned(),
            ));
        }
        let delta = (v_max - v_min) / (num_atoms - 1) as f32;
        let values: Vec<f32> = (0..num_atoms).map(|i| v_min + i as f32 * delta).collect();
        GpuTensor::from_host(&values, vec![num_atoms], stream)
    }

    /// All data in `GpuVarStore` is natively `CudaSlice<f32>` -- this is a no-op.
    /// (Candle DType/contiguous conversion removed.)
    #[allow(dead_code, clippy::unnecessary_wraps)]
    const fn ensure_f32_contiguous(_vars: &GpuVarStore) -> Result<(), MLError> {
        Ok(())
    }

    /// Resample noise in all `NoisyLinear` heads.
    ///
    /// Must be called before each training forward pass.
    pub fn reset_noise(&mut self) -> Result<(), MLError> {
        self.value_fc.reset_noise()?;
        self.value_out.reset_noise()?;
        for fc in &mut self.branch_fcs {
            fc.reset_noise()?;
        }
        for out in &mut self.branch_outs {
            out.reset_noise()?;
        }
        Ok(())
    }

    /// Resample noise with custom sigma scale in all `NoisyLinear` heads.
    pub fn reset_noise_with_sigma(&mut self, sigma_scale: f64) -> Result<(), MLError> {
        self.value_fc.reset_noise_with_sigma(sigma_scale)?;
        self.value_out.reset_noise_with_sigma(sigma_scale)?;
        for fc in &mut self.branch_fcs {
            fc.reset_noise_with_sigma(sigma_scale)?;
        }
        for out in &mut self.branch_outs {
            out.reset_noise_with_sigma(sigma_scale)?;
        }
        Ok(())
    }

    /// Disable noise in all `NoisyLinear` heads (for evaluation).
    pub fn disable_noise(&mut self) -> Result<(), MLError> {
        self.value_fc.disable_noise()?;
        self.value_out.disable_noise()?;
        for fc in &mut self.branch_fcs {
            fc.disable_noise()?;
        }
        for out in &mut self.branch_outs {
            out.disable_noise()?;
        }
        Ok(())
    }

    /// Forward pass returning decomposed value and per-branch advantages (cold path).
    ///
    /// **Hot-path forward is handled by `GpuDqnTrainer::forward_only_q()` which uses
    /// the fused CUDA kernel `dqn_forward_only_kernel` with BF16 tensor core matmul
    /// (1 kernel launch vs ~160 Candle dispatch calls). This Candle-based forward
    /// exists for:**
    /// - Unit tests validating network outputs
    /// - The backward pass (Candle autograd through shared/value/branch layers)
    /// - Weight initialization verification
    ///
    /// # Arguments
    /// * `state` - Input tensor [batch, `state_dim`]
    /// * `train` - Enable dropout (true during training)
    ///
    /// # Returns
    /// `BranchOutput` with:
    /// - `value` [batch, 1]: state value (expected Q when distributional)
    /// - `advantages` [batch, `n_d`]: per-branch expected Q-values (for greedy selection)
    /// - `advantage_log_probs` (distributional only): [batch, `n_d`, `num_atoms`] log-softmax
    /// - `value_log_probs` (distributional only): [batch, 1, `num_atoms`] log-softmax
    pub fn forward_branches(
        &self,
        state: &GpuTensor,
        _train: bool,
    ) -> Result<BranchOutput, MLError> {
        // Forward pass through shared encoder, then branch to value + per-branch advantages.
        let cublas = cudarc::cublas::CudaBlas::new(self.stream.clone()).map_err(|e| {
            MLError::ModelError(format!("branching forward cublas init: {e}"))
        })?;
        let act_kernels = ml_core::cuda_autograd::ActivationKernels::new(&self.stream)?;
        let alpha = self.config.leaky_relu_alpha as f32;

        // Shared feature extraction
        let mut h = state.gpu_clone(&self.stream)?;
        for layer in &self.shared_layers {
            let (out, _acts) = layer.forward(&h, &self.vars, &cublas, &self.stream)?;
            let (activated, _mask) = act_kernels.leaky_relu_fwd(&out, alpha, &self.stream)?;
            h = activated;
        }
        // h: [batch, final_hidden]

        // Value head: NoisyLinear fc + LeakyReLU + NoisyLinear out
        let v_hidden = self.value_fc.forward(&h)?;
        let (v_act, _mask) = act_kernels.leaky_relu_fwd(&v_hidden, alpha, &self.stream)?;
        let v_raw = self.value_out.forward(&v_act)?;
        // v_raw: [batch, num_atoms]

        // Distributional forward for per-branch advantages
        self.forward_distributional(&h, v_raw)
    }

    /// Distributional (C51) forward path.
    ///
    /// Each branch outputs `[batch, n_d * num_atoms]`, reshaped to `[batch, n_d, num_atoms]`.
    /// Log-softmax is applied along the atoms dimension (dim=2).
    /// Expected Q = sum(softmax(logits) * z) stored in `advantages` for greedy selection.
    fn forward_distributional(
        &self,
        h: &GpuTensor,
        v_raw: GpuTensor,
    ) -> Result<BranchOutput, MLError> {
        // Each branch outputs [batch, n_d * num_atoms], reshaped to [batch, n_d, num_atoms].
        // Log-softmax along atoms dim, expected Q via support broadcast.
        let act_kernels = ml_core::cuda_autograd::ActivationKernels::new(&self.stream)?;
        let alpha = self.config.leaky_relu_alpha as f32;
        let num_atoms = self.config.num_atoms;
        let batch = h.shape().first().copied().unwrap_or(1);

        // Get support atoms on host
        let support_host = if let Some(s) = &self.support {
            s.to_host(&self.stream)?
        } else {
            let delta = (self.config.v_max - self.config.v_min) / (num_atoms - 1) as f32;
            (0..num_atoms).map(|i| self.config.v_min + i as f32 * delta).collect()
        };

        // Process value stream: log-softmax over atoms
        let v_host = v_raw.to_host(&self.stream)?;
        let mut value_log_probs_data = Vec::with_capacity(batch * num_atoms);
        let mut value_expected = Vec::with_capacity(batch);

        for b in 0..batch {
            // log_softmax for value atoms
            let start = b * num_atoms;
            let logits: Vec<f32> = (0..num_atoms)
                .map(|i| v_host.get(start + i).copied().unwrap_or(0.0))
                .collect();
            let max_val = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
            let log_sum_exp = {
                let exps_sum: f32 = logits.iter().map(|&v| (v - max_val).exp()).sum();
                max_val + exps_sum.ln()
            };
            let log_probs: Vec<f32> = logits.iter().map(|&v| v - log_sum_exp).collect();

            // Expected value = sum(softmax(logits) * support)
            let mut expected = 0.0_f32;
            for (i, &lp) in log_probs.iter().enumerate() {
                let prob = lp.exp();
                let z = support_host.get(i).copied().unwrap_or(0.0);
                expected += prob * z;
            }
            value_expected.push(expected);
            value_log_probs_data.extend_from_slice(&log_probs);
        }

        let value = GpuTensor::from_host(&value_expected, vec![batch, 1], &self.stream)?;
        let value_log_probs = GpuTensor::from_host(
            &value_log_probs_data, vec![batch, 1, num_atoms], &self.stream,
        )?;

        // Process each branch
        let mut advantages = Vec::with_capacity(self.config.branch_sizes.len());
        let mut advantage_log_probs_vec = Vec::with_capacity(self.config.branch_sizes.len());

        for (&n_d, (fc, out)) in self.config.branch_sizes.iter()
            .zip(self.branch_fcs.iter().zip(self.branch_outs.iter()))
        {
            // Branch forward: NoisyLinear fc + LeakyReLU + NoisyLinear out
            let b_hidden = fc.forward(h)?;
            let (b_act, _mask) = act_kernels.leaky_relu_fwd(&b_hidden, alpha, &self.stream)?;
            let b_raw = out.forward(&b_act)?;
            // b_raw: [batch, n_d * num_atoms]

            let b_host = b_raw.to_host(&self.stream)?;

            // Reshape to [batch, n_d, num_atoms], apply log-softmax, compute expected Q
            let mut branch_expected = Vec::with_capacity(batch * n_d);
            let mut branch_log_probs = Vec::with_capacity(batch * n_d * num_atoms);

            for b_idx in 0..batch {
                for action in 0..n_d {
                    let start = b_idx * n_d * num_atoms + action * num_atoms;
                    let logits: Vec<f32> = (0..num_atoms)
                        .map(|i| b_host.get(start + i).copied().unwrap_or(0.0))
                        .collect();

                    let max_val = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
                    let log_sum_exp = {
                        let exps_sum: f32 = logits.iter().map(|&v| (v - max_val).exp()).sum();
                        max_val + exps_sum.ln()
                    };
                    let log_probs: Vec<f32> = logits.iter().map(|&v| v - log_sum_exp).collect();

                    // Expected Q for this action = sum(softmax * support)
                    let mut expected = 0.0_f32;
                    for (i, &lp) in log_probs.iter().enumerate() {
                        let prob = lp.exp();
                        let z = support_host.get(i).copied().unwrap_or(0.0);
                        expected += prob * z;
                    }
                    branch_expected.push(expected);
                    branch_log_probs.extend_from_slice(&log_probs);
                }
            }

            let adv = GpuTensor::from_host(&branch_expected, vec![batch, n_d], &self.stream)?;
            let adv_lp = GpuTensor::from_host(
                &branch_log_probs, vec![batch, n_d, num_atoms], &self.stream,
            )?;

            advantages.push(adv);
            advantage_log_probs_vec.push(adv_lp);
        }

        Ok(BranchOutput {
            value,
            advantages,
            advantage_log_probs: Some(advantage_log_probs_vec),
            value_log_probs: Some(value_log_probs),
        })
    }

    /// Inference-mode forward (no dropout).
    pub fn forward_branches_eval(&self, state: &GpuTensor) -> Result<BranchOutput, MLError> {
        self.forward_branches(state, false)
    }

    /// Compute aggregate Q(s, a) for given per-branch actions.
    ///
    /// Q(s,a) = V(s) + (1/D) x `sum_d` [`A_d(s`, `a_d`) - `mean(A_d(s`, .))]
    ///
    /// GPU-native: uses GPU gather + GPU `col_sums` for mean. Zero full-tensor download.
    /// Only per-branch row means are computed via GPU reduction.
    ///
    /// # Arguments
    /// * `output` - Branch output from `forward_branches()`
    /// * `branch_actions` - Per-branch action indices: D tensors, each [batch] of f32 (u32 cast)
    ///
    /// # Returns
    /// Aggregate Q-values [batch]
    pub fn aggregate_q_for_actions(
        output: &BranchOutput,
        branch_actions: &[GpuTensor],
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        let d = output.advantages.len();
        if branch_actions.len() != d {
            return Err(MLError::InvalidInput(format!(
                "Expected {} branch actions, got {}",
                d,
                branch_actions.len()
            )));
        }

        let v = output
            .value
            .squeeze(1, stream)
            .map_err(|e| MLError::ModelError(format!("Value squeeze: {}", e)))?; // [batch]

        let mut centered_sum = GpuTensor::zeros_like(&v, stream)
            .map_err(|e| MLError::ModelError(format!("Zeros like: {}", e)))?;

        let elem_kernels = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;
        let reductions = ml_core::cuda_autograd::ReductionKernels::new(stream)?;

        for (a_d, action_d) in output.advantages.iter().zip(branch_actions.iter()) {
            let shape = a_d.shape();
            let batch_size = shape.first().copied().unwrap_or(0);
            let n_d = shape.get(1).copied().unwrap_or(0);

            // GPU gather: taken[i] = a_d[i, action_d[i]]
            // Cast f32 action indices to u32 on GPU via elementwise floor, then reinterpret.
            // action_d contains integer indices stored as f32; cast them to u32 for the gather kernel.
            let act_u32: CudaSlice<u32> = {
                // Download small action vector (batch_size u32s = tiny) and re-upload as u32
                let act_host = action_d.to_host(stream).map_err(|e| {
                    MLError::ModelError(format!("Action to_host for u32 cast: {}", e))
                })?;
                let act_u32_host: Vec<u32> = act_host.iter().map(|&v| v as u32).collect();
                stream.clone_htod(&act_u32_host).map_err(|e| {
                    MLError::ModelError(format!("Action u32 upload: {}", e))
                })?
            };

            // GPU gather: out[i] = a_d[i * n_d + indices[i]]
            let gathered = elem_kernels.gather(
                a_d.cuda_data(), &act_u32, batch_size, 1, 1, n_d,
            ).map_err(|e| MLError::ModelError(format!("GPU gather: {}", e)))?;
            let taken = GpuTensor::new(gathered, vec![batch_size])?;

            // GPU mean per row: col_sums gives sum per column (transpose view).
            // For [batch, n_d], row mean = sum of each row / n_d.
            // col_sums operates on [rows, cols] -> per-col sum [cols].
            // We want per-row mean, so transpose first: [n_d, batch] -> col_sums -> [batch]
            let transposed = elem_kernels.transpose_2d(a_d.cuda_data(), batch_size, n_d)?;
            let row_sums_host = reductions.col_sums(&transposed, n_d, batch_size)?;
            let row_means: Vec<f32> = row_sums_host.iter().map(|&s| s / n_d as f32).collect();
            let mean_tensor = GpuTensor::from_host(&row_means, vec![batch_size], stream)
                .map_err(|e| MLError::ModelError(format!("Mean tensor: {}", e)))?;

            // centered = taken - mean (GPU sub)
            let centered = taken.sub(&mean_tensor, stream)
                .map_err(|e| MLError::ModelError(format!("Centered sub: {}", e)))?;

            centered_sum = centered_sum
                .add(&centered, stream)
                .map_err(|e| MLError::ModelError(format!("Centered sum: {}", e)))?;
        }

        // Q(s, a) = V(s) + (1/D) x sum centered advantages
        let inv_d = 1.0_f32 / d as f32;
        let scale = GpuTensor::scalar(inv_d, stream)
            .map_err(|e| MLError::ModelError(format!("Scale tensor: {}", e)))?;
        let scaled = centered_sum
            .broadcast_mul(&scale, stream)
            .map_err(|e| MLError::ModelError(format!("Scale multiply: {}", e)))?;

        v.add(&scaled, stream)
            .map_err(|e| MLError::ModelError(format!("Q aggregate: {}", e)))
    }

    /// Compute max aggregate Q-value using greedy per-branch selection.
    ///
    /// Q*(s) = V(s) + (1/D) x `sum_d` [max_{`a_d`} `A_d(s`, `a_d`) - `mean(A_d)`]
    ///
    /// GPU-native: per-row max via flat argmax + gather, per-row mean via transpose + `col_sums`.
    /// Zero full-tensor download.
    ///
    /// Used for computing TD targets: y = r + gamma x Q*_target(s').
    pub fn max_aggregate_q(output: &BranchOutput, stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
        // Max Q = Q at greedy actions, by definition.
        // Greedy actions = argmax per branch, then aggregate_q_for_actions computes
        // Q = V + (1/D) * sum_d(A_d[greedy_d] - mean(A_d))
        let greedy = Self::greedy_branch_actions_batch(output, stream)?;
        Self::aggregate_q_for_actions(output, &greedy, stream)
    }

    /// Greedy action selection: argmax per branch independently.
    ///
    /// GPU-native: runs flat argmax kernel per branch (single scalar `DtoH` readback
    /// per branch -- 4 bytes each, 12 bytes total for 3 branches). No full tensor download.
    ///
    /// # Returns
    /// Vector of D action indices (one per branch).
    pub fn greedy_branch_actions(output: &BranchOutput, stream: &Arc<CudaStream>) -> Result<Vec<u32>, MLError> {
        let reductions = ml_core::cuda_autograd::ReductionKernels::new(stream)?;
        let mut actions = Vec::with_capacity(output.advantages.len());
        for a_d in &output.advantages {
            let n_d = a_d.shape().get(1).copied().unwrap_or(a_d.numel());
            // GPU argmax on the first row's n_d elements (single-batch inference).
            // Uses flat argmax kernel which works correctly for any n_d >= 1.
            let first_row = clone_subrange(a_d.cuda_data(), 0, n_d, stream)?;
            let idx = reductions.argmax(&first_row, n_d)?;
            actions.push(idx);
        }
        Ok(actions)
    }

    /// Batch greedy action selection: argmax per branch for a batch.
    ///
    /// GPU-native: runs flat argmax kernel per row per branch. Each row yields
    /// a single `u32` index (4 bytes `DtoH`). For batch=B and D=3 branches with
    /// `n_d` in {3,5}, total readback is B*D*4 bytes (e.g. 256*3*4 = 3 KB).
    ///
    /// # Returns
    /// D tensors of shape [batch], each containing u32 action indices (as f32).
    pub fn greedy_branch_actions_batch(output: &BranchOutput, stream: &Arc<CudaStream>) -> Result<Vec<GpuTensor>, MLError> {
        let reductions = ml_core::cuda_autograd::ReductionKernels::new(stream)?;
        let mut actions = Vec::with_capacity(output.advantages.len());
        for a_d in &output.advantages {
            let shape = a_d.shape();
            let batch_size = shape.first().copied().unwrap_or(0);
            let n_d = shape.get(1).copied().unwrap_or(0);

            // GPU flat argmax per row -- avoids the argmax_rows kernel bug
            // where warp-shuffle returns 0 for inactive lanes when cols < 32.
            let mut indices_f32 = Vec::with_capacity(batch_size);
            for b in 0..batch_size {
                let row_offset = b * n_d;
                let row_owned = clone_subrange(a_d.cuda_data(), row_offset, n_d, stream)?;
                let idx = reductions.argmax(&row_owned, n_d)?;
                indices_f32.push(idx as f32);
            }

            let indices_tensor = GpuTensor::from_host(&indices_f32, vec![batch_size], stream)
                .map_err(|e| MLError::ModelError(format!("greedy indices tensor: {e}")))?;
            actions.push(indices_tensor);
        }
        Ok(actions)
    }

    /// Decompose a factored action index (0-44) into per-branch indices.
    ///
    /// Index = exposure × (`num_order` × `num_urgency`) + order × `num_urgency` + urgency
    ///
    /// # Returns
    /// (`exposure_idx`, `order_idx`, `urgency_idx`)
    pub const fn decompose_factored_action(
        factored_idx: usize,
        num_order_types: usize,
        num_urgency_levels: usize,
    ) -> (usize, usize, usize) {
        let stride = num_order_types * num_urgency_levels;
        let exposure = factored_idx / stride;
        let order = (factored_idx % stride) / num_urgency_levels;
        let urgency = factored_idx % num_urgency_levels;
        (exposure, order, urgency)
    }

    /// Decompose a batch of factored action indices into per-branch tensors.
    ///
    /// # Arguments
    /// * `actions` - Slice of u32 factored indices
    /// * `stream` - CUDA stream for tensor allocation
    ///
    /// # Returns
    /// 3 tensors of f32 (u32 indices cast to f32): [batch] exposure, [batch] order, [batch] urgency
    pub fn decompose_actions_batch(
        actions: &[u32],
        stream: &Arc<CudaStream>,
        num_order_types: usize,
        num_urgency_levels: usize,
    ) -> Result<Vec<GpuTensor>, MLError> {
        let mut exposures = Vec::with_capacity(actions.len());
        let mut orders = Vec::with_capacity(actions.len());
        let mut urgencies = Vec::with_capacity(actions.len());

        for &a in actions {
            let (e, o, u) = Self::decompose_factored_action(a as usize, num_order_types, num_urgency_levels);
            exposures.push(e as f32);
            orders.push(o as f32);
            urgencies.push(u as f32);
        }

        let e_tensor = GpuTensor::from_host(&exposures, vec![actions.len()], stream)
            .map_err(|e| MLError::ModelError(format!("Exposure tensor: {}", e)))?;
        let o_tensor = GpuTensor::from_host(&orders, vec![actions.len()], stream)
            .map_err(|e| MLError::ModelError(format!("Order tensor: {}", e)))?;
        let u_tensor = GpuTensor::from_host(&urgencies, vec![actions.len()], stream)
            .map_err(|e| MLError::ModelError(format!("Urgency tensor: {}", e)))?;

        Ok(vec![e_tensor, o_tensor, u_tensor])
    }

    /// GPU-native decomposition of factored action indices into per-branch tensors.
    ///
    /// Operates entirely on-device using GPU elementwise arithmetic + GPU floor:
    /// - exposure = floor(action / stride)
    /// - remainder = action - exposure * stride
    /// - order = `floor(remainder / urgency_levels)`
    /// - urgency = remainder - order * `urgency_levels`
    ///
    /// Zero GPU->CPU->GPU roundtrip.
    ///
    /// # Arguments
    /// * `actions` - `GpuTensor` of f32 factored indices (cast from u32), shape [batch]
    ///
    /// # Returns
    /// 3 tensors of f32: [batch] exposure, [batch] order, [batch] urgency
    pub fn decompose_actions_batch_gpu(
        actions: &GpuTensor,
        num_order_types: usize,
        num_urgency_levels: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Vec<GpuTensor>, MLError> {
        let stride = (num_order_types * num_urgency_levels) as f64;
        let urg = num_urgency_levels as f64;
        let batch = actions.numel();

        let elem_k = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;

        // exposure = floor(action / stride) — GPU affine + GPU floor (unary op=2)
        let scaled_for_exp = actions.affine(1.0 / stride, 0.0, stream)
            .map_err(|e| MLError::ModelError(format!("decompose gpu: scale for exposure: {e}")))?;
        let exp_floored = elem_k.unary(scaled_for_exp.cuda_data(), batch, 2, 0.0, 0.0)
            .map_err(|e| MLError::ModelError(format!("decompose gpu: floor exposure: {e}")))?;
        let e_t = GpuTensor::new(exp_floored, vec![batch])?;

        // remainder = action - exposure * stride  (GPU sub + GPU affine)
        let exp_times_stride = e_t.affine(stride, 0.0, stream)
            .map_err(|e| MLError::ModelError(format!("decompose gpu: exp*stride: {e}")))?;
        let remainder = actions.sub(&exp_times_stride, stream)
            .map_err(|e| MLError::ModelError(format!("decompose gpu: remainder: {e}")))?;

        // order = floor(remainder / urgency_levels) — GPU affine + GPU floor
        let scaled_for_ord = remainder.affine(1.0 / urg, 0.0, stream)
            .map_err(|e| MLError::ModelError(format!("decompose gpu: scale for order: {e}")))?;
        let ord_floored = elem_k.unary(scaled_for_ord.cuda_data(), batch, 2, 0.0, 0.0)
            .map_err(|e| MLError::ModelError(format!("decompose gpu: floor order: {e}")))?;
        let o_t = GpuTensor::new(ord_floored, vec![batch])?;

        // urgency = remainder - order * urgency_levels  (GPU sub + GPU affine)
        let ord_times_urg = o_t.affine(urg, 0.0, stream)
            .map_err(|e| MLError::ModelError(format!("decompose gpu: ord*urg: {e}")))?;
        let u_t = remainder.sub(&ord_times_urg, stream)
            .map_err(|e| MLError::ModelError(format!("decompose gpu: urgency: {e}")))?;

        Ok(vec![e_t, o_t, u_t])
    }

    /// Compose per-branch action indices into a factored index.
    pub const fn compose_factored_action(
        exposure: u32,
        order: u32,
        urgency: u32,
        num_order_types: u32,
        num_urgency_levels: u32,
    ) -> u32 {
        exposure * (num_order_types * num_urgency_levels) + order * num_urgency_levels + urgency
    }

    /// Get `GpuVarStore` for optimizer and serialization.
    pub const fn vars(&self) -> &GpuVarStore {
        &self.vars
    }

    /// Mutable reference to the `GpuVarStore` (for checkpoint import).
    pub const fn vars_mut(&mut self) -> &mut GpuVarStore {
        &mut self.vars
    }

    /// Collect ALL trainable `cudarc::driver::CudaSlice<f32>`s: `GpuVarStore` (shared encoder + mu) + `NoisyLinear` sigma.
    ///
    /// Mu vars (`weight_mu`, `bias_mu`) are registered in `GpuVarStore` at construction
    /// time for GPU experience collector compatibility. Sigma vars (`weight_sigma`,
    /// `bias_sigma`) remain standalone. This method collects both without duplication.
    pub fn all_trainable_vars(&self) -> Vec<CudaSlice<f32>> {
        let mut vars: Vec<CudaSlice<f32>> = Vec::new();
        // Shared encoder vars from GpuVarStore
        for (_name, param) in self.vars.iter() {
            // Clone the CudaSlice reference -- this is cheap (Arc bump)
            vars.push(param.data.clone());
        }
        // Only sigma vars -- mu already in GpuVarStore
        for s in self.value_fc.noisy_sigma_slices() {
            vars.push(s.clone());
        }
        for s in self.value_out.noisy_sigma_slices() {
            vars.push(s.clone());
        }
        for fc in &self.branch_fcs {
            for s in fc.noisy_sigma_slices() {
                vars.push(s.clone());
            }
        }
        for out in &self.branch_outs {
            for s in out.noisy_sigma_slices() {
                vars.push(s.clone());
            }
        }
        vars
    }

    /// Collect only the `NoisyLinear` sigma `cudarc::driver::CudaSlice<f32>`s (for target network Polyak updates).
    ///
    /// Mu vars are registered in `GpuVarStore`, so `polyak_update()` on `GpuVarStore` handles them.
    /// This method returns only sigma vars for `polyak_update_var_pairs()`.
    ///
    /// Returns vars in a deterministic order: `value_fc`, `value_out`, then
    /// `branch_fcs[0..D]`, `branch_outs[0..D]`. Both online and target networks
    /// produce the same order, so vars can be zipped for Polyak update.
    pub fn noisy_vars_ordered(&self) -> Vec<CudaSlice<f32>> {
        let mut vars = Vec::new();
        for s in self.value_fc.noisy_sigma_slices() {
            vars.push(s.clone());
        }
        for s in self.value_out.noisy_sigma_slices() {
            vars.push(s.clone());
        }
        for fc in &self.branch_fcs {
            for s in fc.noisy_sigma_slices() {
                vars.push(s.clone());
            }
        }
        for out in &self.branch_outs {
            for s in out.noisy_sigma_slices() {
                vars.push(s.clone());
            }
        }
        vars
    }

    /// Get stream.
    pub const fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Get configuration.
    pub const fn config(&self) -> &BranchingConfig {
        &self.config
    }

    /// Copy weights from another branching network (target network sync).
    ///
    /// Copies both `GpuVarStore` vars (shared encoder) AND `NoisyLinear` head vars.
    /// Zero CPU downloads — all copies are async device-to-device memcpy.
    pub fn copy_weights_from(&mut self, other: &BranchingDuelingQNetwork) -> Result<(), MLError> {
        // 1. Copy GpuVarStore vars (shared encoder layers) via DtoD memcpy
        for (name, other_param) in other.vars.iter() {
            if let Some(self_param) = self.vars.get_mut(name) {
                let num_bytes = other_param.data.len() * std::mem::size_of::<f32>();
                let src_ptr = {
                    let (ptr, guard) = other_param.data.device_ptr(&self.stream);
                    let _no_drop = ManuallyDrop::new(guard);
                    ptr
                };
                let dst_ptr = {
                    let (ptr, guard) = self_param.data.device_ptr_mut(&self.stream);
                    let _no_drop = ManuallyDrop::new(guard);
                    ptr
                };
                // SAFETY: src and dst are valid device allocations on the same context.
                // num_bytes = len * sizeof(f32) does not exceed either allocation.
                unsafe {
                    cudarc::driver::result::memcpy_dtod_async(
                        dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(),
                    ).map_err(|e| {
                        MLError::ModelError(format!("Copy weight DtoD {name}: {e}"))
                    })?;
                }
            }
        }

        // 2. Copy NoisyLinear head vars via DtoD memcpy (delegated to NoisyLinear::copy_params_from)
        Self::copy_noisy_layer(&mut self.value_fc, &other.value_fc)?;
        Self::copy_noisy_layer(&mut self.value_out, &other.value_out)?;
        for (self_fc, other_fc) in self.branch_fcs.iter_mut().zip(other.branch_fcs.iter()) {
            Self::copy_noisy_layer(self_fc, other_fc)?;
        }
        for (self_out, other_out) in self.branch_outs.iter_mut().zip(other.branch_outs.iter()) {
            Self::copy_noisy_layer(self_out, other_out)?;
        }

        Ok(())
    }

    /// Copy `NoisyLinear` weights between matching layers via `DtoD` memcpy.
    fn copy_noisy_layer(
        dst: &mut MaybeNoisyLinear,
        src: &MaybeNoisyLinear,
    ) -> Result<(), MLError> {
        let MaybeNoisyLinear::Noisy(src_noisy) = src;
        #[allow(irrefutable_let_patterns)]
        let MaybeNoisyLinear::Noisy(dst_noisy) = dst;
        dst_noisy.copy_params_from(src_noisy)
    }
}

impl std::fmt::Debug for BranchingDuelingQNetwork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BranchingDuelingQNetwork")
            .field("config", &self.config)
            .field("num_shared_layers", &self.shared_layers.len())
            .field("num_branches", &self.branch_fcs.len())
            .field("stream", &"Arc<CudaStream>")
            .finish()
    }
}

#[cfg(test)]
#[allow(
    clippy::get_first,
    clippy::assertions_on_result_states,
    clippy::redundant_clone
)]
mod tests {
    use super::*;

    /// Helper: create a test CUDA stream.
    fn test_stream() -> Arc<CudaStream> {
        let ctx = cudarc::driver::CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("CUDA stream required")
    }

    /// Helper: create a default distributional+noisy config for tests.
    fn trading_config_default(state_dim: usize) -> BranchingConfig {
        BranchingConfig::trading_default(state_dim, vec![64, 32], 16, vec![5, 3, 3])
    }

    /// Helper: create a distributional+noisy config with smaller atom count for tests.
    fn trading_config_small_atoms(state_dim: usize) -> BranchingConfig {
        let mut cfg = BranchingConfig::trading_default(state_dim, vec![64, 32], 16, vec![5, 3, 3]);
        cfg.num_atoms = 11; // smaller for tests
        cfg.v_min = -5.0;
        cfg.v_max = 5.0;
        cfg
    }

    // ======================================================================
    // Existing scalar tests (backward compatibility)
    // ======================================================================

    #[test]
    fn test_creation() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(16);
        let net = BranchingDuelingQNetwork::new(config, stream)?;
        assert_eq!(net.shared_layers.len(), 2);
        assert_eq!(net.branch_fcs.len(), 3);
        assert_eq!(net.branch_outs.len(), 3);
        Ok(())
    }

    #[test]
    fn test_forward_shapes() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(16);
        let num_atoms = config.num_atoms;
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let batch = 4;
        let state = GpuTensor::randn(&[batch, 16], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        assert_eq!(output.value.shape(), &[batch, 1]);
        assert_eq!(output.advantages.len(), 3);
        assert_eq!(
            output.advantages.get(0).map(|t| t.shape().to_vec()),
            Some(vec![batch, 5])
        ); // exposure
        assert_eq!(
            output.advantages.get(1).map(|t| t.shape().to_vec()),
            Some(vec![batch, 3])
        ); // order
        assert_eq!(
            output.advantages.get(2).map(|t| t.shape().to_vec()),
            Some(vec![batch, 3])
        ); // urgency
        // Distributional is always enabled: log_probs should be present
        assert!(output.advantage_log_probs.is_some());
        assert!(output.value_log_probs.is_some());
        let adv_lp = output.advantage_log_probs.as_ref().unwrap();
        assert_eq!(adv_lp.len(), 3);
        assert_eq!(
            adv_lp.get(0).map(|t| t.shape().to_vec()),
            Some(vec![batch, 5, num_atoms])
        );
        Ok(())
    }

    #[test]
    fn test_aggregate_q() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(8);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[2, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        // Actions: sample 0 = (2, 1, 0), sample 1 = (4, 0, 2)
        let exposure = GpuTensor::from_host(&[2.0_f32, 4.0], vec![2], &stream)?;
        let order = GpuTensor::from_host(&[1.0_f32, 0.0], vec![2], &stream)?;
        let urgency = GpuTensor::from_host(&[0.0_f32, 2.0], vec![2], &stream)?;

        let q = BranchingDuelingQNetwork::aggregate_q_for_actions(
            &output,
            &[exposure, order, urgency],
            &stream,
        )?;
        assert_eq!(q.shape(), &[2]);

        // Check finite
        let q_vec = q.to_host(&stream)?;
        for &v in &q_vec {
            assert!(v.is_finite(), "Q should be finite, got {}", v);
        }
        Ok(())
    }

    #[test]
    fn test_max_aggregate_q() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(8);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[3, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        let q_max = BranchingDuelingQNetwork::max_aggregate_q(&output, &stream)?;
        assert_eq!(q_max.shape(), &[3]);

        // Max Q should be >= any specific action Q
        let greedy = BranchingDuelingQNetwork::greedy_branch_actions_batch(&output, &stream)?;
        let q_greedy =
            BranchingDuelingQNetwork::aggregate_q_for_actions(&output, &greedy, &stream)?;

        let max_vec = q_max.to_host(&stream)?;
        let greedy_vec = q_greedy.to_host(&stream)?;
        for (m, g) in max_vec.iter().zip(greedy_vec.iter()) {
            // max_aggregate_q delegates to greedy_branch_actions + aggregate_q_for_actions,
            // so results should be identical (same code path, same GPU ops).
            assert!(
                (m - g).abs() < 1e-5,
                "Max Q ({m}) should equal greedy Q ({g})"
            );
        }
        Ok(())
    }

    #[test]
    fn test_greedy_actions_single() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(8);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[1, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        let actions = BranchingDuelingQNetwork::greedy_branch_actions(&output, &stream)?;
        assert_eq!(actions.len(), 3);
        assert!(
            *actions.get(0).unwrap_or(&99) < 5,
            "Exposure action should be 0-4"
        );
        assert!(
            *actions.get(1).unwrap_or(&99) < 3,
            "Order action should be 0-2"
        );
        assert!(
            *actions.get(2).unwrap_or(&99) < 3,
            "Urgency action should be 0-2"
        );
        Ok(())
    }

    #[test]
    fn test_decompose_compose_roundtrip() {
        for idx in 0..45_u32 {
            let (e, o, u) = BranchingDuelingQNetwork::decompose_factored_action(idx as usize, 3, 3);
            let composed =
                BranchingDuelingQNetwork::compose_factored_action(e as u32, o as u32, u as u32, 3, 3);
            assert_eq!(composed, idx, "Roundtrip failed for {}", idx);
        }
    }

    #[test]
    fn test_decompose_actions_batch() -> anyhow::Result<()> {
        let stream = test_stream();
        // Action 0 = (0,0,0), Action 13 = (1,1,1), Action 44 = (4,2,2)
        let actions = vec![0_u32, 13, 44];
        let branches =
            BranchingDuelingQNetwork::decompose_actions_batch(&actions, &stream, 3, 3)?;

        let e = branches
            .get(0)
            .ok_or_else(|| anyhow::anyhow!("missing branch 0"))?
            .to_host(&stream)?;
        let o = branches
            .get(1)
            .ok_or_else(|| anyhow::anyhow!("missing branch 1"))?
            .to_host(&stream)?;
        let u = branches
            .get(2)
            .ok_or_else(|| anyhow::anyhow!("missing branch 2"))?
            .to_host(&stream)?;

        assert_eq!(e, vec![0.0_f32, 1.0, 4.0]);
        assert_eq!(o, vec![0.0_f32, 1.0, 2.0]);
        assert_eq!(u, vec![0.0_f32, 1.0, 2.0]);
        Ok(())
    }

    #[test]
    fn test_decompose_actions_batch_gpu_matches_cpu() -> anyhow::Result<()> {
        let stream = test_stream();
        // Test all 45 factored actions -- GPU-native version must match CPU version
        let all_actions: Vec<u32> = (0..45).collect();
        let cpu_branches =
            BranchingDuelingQNetwork::decompose_actions_batch(&all_actions, &stream, 3, 3)?;

        let all_f32: Vec<f32> = all_actions.iter().map(|&x| x as f32).collect();
        let actions_tensor = GpuTensor::from_host(&all_f32, vec![45], &stream)?;
        let gpu_branches =
            BranchingDuelingQNetwork::decompose_actions_batch_gpu(&actions_tensor, 3, 3, &stream)?;

        for (d, name) in [(0, "exposure"), (1, "order"), (2, "urgency")] {
            let cpu_t = cpu_branches
                .get(d)
                .ok_or_else(|| anyhow::anyhow!("missing cpu branch {d}"))?;
            let gpu_t = gpu_branches
                .get(d)
                .ok_or_else(|| anyhow::anyhow!("missing gpu branch {d}"))?;
            let cpu_host = cpu_t.to_host(&stream)?;
            let gpu_host = gpu_t.to_host(&stream)?;
            let max_diff = cpu_host.iter().zip(gpu_host.iter())
                .map(|(a, b)| (a - b).abs())
                .fold(0.0_f32, f32::max);
            assert!(
                max_diff < 0.5,
                "{name} mismatch between CPU and GPU decompose: max_diff={max_diff}"
            );
        }
        Ok(())
    }

    #[test]
    fn test_weight_copy() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(8);
        let net1 = BranchingDuelingQNetwork::new(config.clone(), Arc::clone(&stream))?;
        let mut net2 = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        net2.copy_weights_from(&net1)?;

        let state = GpuTensor::full(&[1, 8], 1.0, &stream)?;
        let out1 = net1.forward_branches_eval(&state)?;
        let out2 = net2.forward_branches_eval(&state)?;

        let val1 = out1.value.to_host(&stream)?;
        let val2 = out2.value.to_host(&stream)?;
        let val_diff = val1.iter().zip(val2.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);
        assert!(val_diff < 1e-5, "Values should match after copy: diff={val_diff}");

        for d in 0..3 {
            let a1 = out1
                .advantages
                .get(d)
                .ok_or_else(|| anyhow::anyhow!("missing adv {} net1", d))?;
            let a2 = out2
                .advantages
                .get(d)
                .ok_or_else(|| anyhow::anyhow!("missing adv {} net2", d))?;
            let a1_host = a1.to_host(&stream)?;
            let a2_host = a2.to_host(&stream)?;
            let adv_diff = a1_host.iter().zip(a2_host.iter())
                .map(|(a, b)| (a - b).abs())
                .fold(0.0_f32, f32::max);
            assert!(
                adv_diff < 1e-5,
                "Branch {} advantages should match after copy: diff={adv_diff}",
                d
            );
        }
        Ok(())
    }

    #[test]
    fn test_aggregate_q_mean_subtraction() -> anyhow::Result<()> {
        let stream = test_stream();
        // Test that mean subtraction ensures zero-mean advantage contribution
        let config = trading_config_default(4);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::full(&[1, 4], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        // Try all possible actions and verify aggregate Q values are finite
        for e in 0..5_u32 {
            for o in 0..3_u32 {
                for u in 0..3_u32 {
                    let exposure = GpuTensor::from_host(&[e as f32], vec![1], &stream)?;
                    let order = GpuTensor::from_host(&[o as f32], vec![1], &stream)?;
                    let urgency = GpuTensor::from_host(&[u as f32], vec![1], &stream)?;

                    let q = BranchingDuelingQNetwork::aggregate_q_for_actions(
                        &output,
                        &[exposure, order, urgency],
                        &stream,
                    )?;
                    let q_host = q.to_host(&stream)?;
                    let v = q_host.first().copied().unwrap_or(f32::NAN);
                    assert!(v.is_finite(), "Q({},{},{}) should be finite", e, o, u);
                }
            }
        }
        Ok(())
    }

    #[test]
    fn test_empty_branch_rejected() {
        let stream = test_stream();
        let config = BranchingConfig {
            state_dim: 8,
            branch_sizes: vec![],
            shared_hidden_dims: vec![16],
            value_hidden_dim: 8,
            branch_hidden_dim: 8,
            leaky_relu_alpha: 0.01,
            dropout_rate: 0.0,
            num_atoms: 51,
            v_min: -25.0,
            v_max: 25.0,
            noisy_sigma_init: 0.5,
        };
        let result = BranchingDuelingQNetwork::new(config, stream);
        assert!(result.is_err());
    }

    #[test]
    fn test_from_dqn_params() {
        let config =
            BranchingConfig::from_dqn_params(48, &[256, 128], 64, 0.01, None, vec![5, 3, 3]);
        assert_eq!(config.state_dim, 48);
        assert_eq!(config.branch_sizes, vec![5, 3, 3]);
        assert_eq!(config.shared_hidden_dims, vec![256, 128]);
        assert_eq!(config.branch_hidden_dim, 64);
        assert_eq!(config.value_hidden_dim, 64);
    }

    #[test]
    fn test_deterministic_forward() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(8);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[2, 8], 1.0, &stream)?;
        let out1 = net.forward_branches_eval(&state)?;
        let out2 = net.forward_branches_eval(&state)?;

        let v1 = out1.value.to_host(&stream)?;
        let v2 = out2.value.to_host(&stream)?;
        for (a, b) in v1.iter().zip(v2.iter()) {
            assert!((a - b).abs() < 1e-6, "Forward should be deterministic");
        }
        Ok(())
    }

    // ======================================================================
    // Feature 1: C51 Distributional tests
    // ======================================================================

    #[test]
    fn test_support_atoms() -> anyhow::Result<()> {
        let stream = test_stream();
        let support = BranchingDuelingQNetwork::support_atoms(-5.0, 5.0, 11, &stream)?;
        assert_eq!(support.shape(), &[11]);
        let vals = support.to_host(&stream)?;
        let first = vals.first().copied().unwrap_or(f32::NAN);
        let last = vals.last().copied().unwrap_or(f32::NAN);
        assert!((first - (-5.0)).abs() < 1e-6, "First atom should be v_min");
        assert!((last - 5.0).abs() < 1e-6, "Last atom should be v_max");
        // Check uniform spacing
        for i in 1..vals.len() {
            let prev = vals.get(i - 1).copied().unwrap_or(f32::NAN);
            let cur = vals.get(i).copied().unwrap_or(f32::NAN);
            assert!(
                (cur - prev - 1.0).abs() < 1e-5,
                "Spacing should be 1.0, got {}",
                cur - prev
            );
        }
        Ok(())
    }

    #[test]
    fn test_support_atoms_invalid() {
        let stream = test_stream();
        let result = BranchingDuelingQNetwork::support_atoms(-5.0, 5.0, 1, &stream);
        assert!(result.is_err(), "num_atoms < 2 should fail");
    }

    #[test]
    fn test_distributional_forward_shapes() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_small_atoms(16);
        let num_atoms = config.num_atoms;
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let batch = 4;
        let state = GpuTensor::randn(&[batch, 16], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        // Value: [batch, 1] (expected V)
        assert_eq!(output.value.shape(), &[batch, 1]);

        // Advantages: [batch, n_d] (expected Q per branch)
        assert_eq!(output.advantages.len(), 3);
        assert_eq!(
            output.advantages.get(0).map(|t| t.shape().to_vec()),
            Some(vec![batch, 5])
        );
        assert_eq!(
            output.advantages.get(1).map(|t| t.shape().to_vec()),
            Some(vec![batch, 3])
        );
        assert_eq!(
            output.advantages.get(2).map(|t| t.shape().to_vec()),
            Some(vec![batch, 3])
        );

        // Advantage log probs: [batch, n_d, num_atoms]
        let adv_lp = output
            .advantage_log_probs
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("advantage_log_probs should be Some"))?;
        assert_eq!(adv_lp.len(), 3);
        assert_eq!(
            adv_lp.get(0).map(|t| t.shape().to_vec()),
            Some(vec![batch, 5, num_atoms])
        );
        assert_eq!(
            adv_lp.get(1).map(|t| t.shape().to_vec()),
            Some(vec![batch, 3, num_atoms])
        );
        assert_eq!(
            adv_lp.get(2).map(|t| t.shape().to_vec()),
            Some(vec![batch, 3, num_atoms])
        );

        // Value log probs: [batch, 1, num_atoms]
        let v_lp = output
            .value_log_probs
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("value_log_probs should be Some"))?;
        assert_eq!(v_lp.shape(), &[batch, 1, num_atoms]);

        Ok(())
    }

    #[test]
    fn test_distributional_log_probs_valid() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_small_atoms(8);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[2, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        // Check that exp(log_probs) sum to 1 along atoms dim
        let adv_lp = output
            .advantage_log_probs
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("expected advantage_log_probs"))?;

        for (d, lp) in adv_lp.iter().enumerate() {
            let lp_host = lp.to_host(&stream)?;
            // Check exp(log_probs) sum for each (batch, action) pair
            let shape = lp.shape();
            let batch = shape.first().copied().unwrap_or(0);
            let n_d = shape.get(1).copied().unwrap_or(0);
            let atoms = shape.get(2).copied().unwrap_or(0);
            for b in 0..batch {
                for a in 0..n_d {
                    let mut s = 0.0_f32;
                    for k in 0..atoms {
                        let idx = b * n_d * atoms + a * atoms + k;
                        let val = lp_host.get(idx).copied().unwrap_or(0.0);
                        s += val.exp();
                    }
                    assert!(
                        (s - 1.0).abs() < 1e-4,
                        "Branch {} probs should sum to 1, got {}",
                        d,
                        s
                    );
                }
            }
        }

        // Check value log probs sum to 1
        let v_lp = output
            .value_log_probs
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("expected value_log_probs"))?;
        let v_host = v_lp.to_host(&stream)?;
        let v_shape = v_lp.shape();
        let v_batch = v_shape.first().copied().unwrap_or(0);
        let v_cols = v_shape.get(1).copied().unwrap_or(0);
        let v_atoms = v_shape.get(2).copied().unwrap_or(0);
        for b in 0..v_batch {
            for c in 0..v_cols {
                let mut s = 0.0_f32;
                for k in 0..v_atoms {
                    let idx = b * v_cols * v_atoms + c * v_atoms + k;
                    let val = v_host.get(idx).copied().unwrap_or(0.0);
                    s += val.exp();
                }
                assert!(
                    (s - 1.0).abs() < 1e-4,
                    "Value probs should sum to 1, got {}",
                    s
                );
            }
        }

        Ok(())
    }

    #[test]
    fn test_distributional_expected_q_manual() -> anyhow::Result<()> {
        let stream = test_stream();
        // Verify expected Q = sum(softmax(logits) * z) matches manual calculation
        let config = trading_config_small_atoms(8);
        let num_atoms = config.num_atoms;
        let v_min = config.v_min;
        let v_max = config.v_max;
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[1, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        let support =
            BranchingDuelingQNetwork::support_atoms(v_min, v_max, num_atoms, &stream)?;
        let z_vals = support.to_host(&stream)?;

        // Check branch 0 expected Q matches manual computation
        let adv_lp = output
            .advantage_log_probs
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("expected advantage_log_probs"))?;
        let lp_0 = adv_lp
            .get(0)
            .ok_or_else(|| anyhow::anyhow!("missing branch 0 log probs"))?;
        // lp_0: [1, 5, num_atoms] -> download to host
        let lp_host = lp_0.to_host(&stream)?;
        let expected_q_host = output
            .advantages
            .get(0)
            .ok_or_else(|| anyhow::anyhow!("missing branch 0 advantages"))?
            .to_host(&stream)?;

        for action_idx in 0..5 {
            let mut manual_q = 0.0_f32;
            for k in 0..num_atoms {
                let log_p = lp_host.get(action_idx * num_atoms + k).copied().unwrap_or(0.0);
                let z = z_vals.get(k).copied().unwrap_or(0.0);
                manual_q += log_p.exp() * z;
            }
            let expected_q = expected_q_host.get(action_idx).copied().unwrap_or(f32::NAN);
            assert!(
                (manual_q - expected_q).abs() < 1e-4,
                "Branch 0, action {}: manual Q ({}) != expected Q ({})",
                action_idx,
                manual_q,
                expected_q,
            );
        }

        Ok(())
    }

    #[test]
    fn test_distributional_expected_q_in_range() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_small_atoms(8);
        let v_min = config.v_min;
        let v_max = config.v_max;
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[4, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        // Expected Q values must be within [v_min, v_max]
        for (d, adv) in output.advantages.iter().enumerate() {
            let vals = adv.to_host(&stream)?;
            for &v in &vals {
                assert!(
                    v >= v_min && v <= v_max,
                    "Branch {} expected Q ({}) should be in [{}, {}]",
                    d,
                    v,
                    v_min,
                    v_max
                );
            }
        }

        // Value expected should also be in range
        let v_vals = output.value.to_host(&stream)?;
        for &v in &v_vals {
            assert!(
                v >= v_min && v <= v_max,
                "Value expected ({}) should be in [{}, {}]",
                v,
                v_min,
                v_max
            );
        }

        Ok(())
    }

    #[test]
    fn test_distributional_aggregate_q() -> anyhow::Result<()> {
        let stream = test_stream();
        // Test that aggregate_q_for_actions works with distributional mode
        let config = trading_config_small_atoms(8);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[2, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        let exposure = GpuTensor::from_host(&[0.0_f32, 4.0], vec![2], &stream)?;
        let order = GpuTensor::from_host(&[2.0_f32, 0.0], vec![2], &stream)?;
        let urgency = GpuTensor::from_host(&[1.0_f32, 2.0], vec![2], &stream)?;

        let q = BranchingDuelingQNetwork::aggregate_q_for_actions(
            &output,
            &[exposure, order, urgency],
            &stream,
        )?;
        assert_eq!(q.shape(), &[2]);

        let q_vec = q.to_host(&stream)?;
        for &v in &q_vec {
            assert!(v.is_finite(), "Distributional Q should be finite, got {}", v);
        }
        Ok(())
    }

    // ======================================================================
    // Feature 2: NoisyNet tests
    // ======================================================================

    #[test]
    fn test_noisy_creation() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(16);
        let net = BranchingDuelingQNetwork::new(config, stream)?;
        assert_eq!(net.shared_layers.len(), 2);
        assert_eq!(net.branch_fcs.len(), 3);
        Ok(())
    }

    #[test]
    fn test_noisy_forward_shapes() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(16);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let batch = 4;
        let state = GpuTensor::randn(&[batch, 16], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        assert_eq!(output.value.shape(), &[batch, 1]);
        assert_eq!(output.advantages.len(), 3);
        assert_eq!(
            output.advantages.get(0).map(|t| t.shape().to_vec()),
            Some(vec![batch, 5])
        );
        // Distributional always enabled: log_probs should be present
        assert!(output.advantage_log_probs.is_some());
        Ok(())
    }

    #[test]
    fn test_noisy_different_after_reset() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(8);
        let mut net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[2, 8], 1.0, &stream)?;

        // First forward with initial noise
        net.reset_noise()?;
        let out1 = net.forward_branches_eval(&state)?;

        // Resample noise
        net.reset_noise()?;
        let out2 = net.forward_branches_eval(&state)?;

        // Outputs should differ due to different noise samples
        let v1 = out1.value.to_host(&stream)?;
        let v2 = out2.value.to_host(&stream)?;
        let diff: f32 = v1.iter().zip(v2.iter()).map(|(a, b)| (a - b).powi(2)).sum();
        assert!(
            diff > 1e-6,
            "NoisyNet outputs should differ after reset_noise (diff={})",
            diff
        );
        Ok(())
    }

    #[test]
    fn test_noisy_deterministic_after_disable() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_default(8);
        let mut net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[2, 8], 1.0, &stream)?;

        // Disable noise (eval mode)
        net.disable_noise()?;
        let out1 = net.forward_branches_eval(&state)?;
        let out2 = net.forward_branches_eval(&state)?;

        let v1 = out1.value.to_host(&stream)?;
        let v2 = out2.value.to_host(&stream)?;
        let diff: f32 = v1.iter().zip(v2.iter()).map(|(a, b)| (a - b).powi(2)).sum();
        assert!(
            diff < 1e-10,
            "With disabled noise, forward should be deterministic (diff={})",
            diff
        );
        Ok(())
    }

    #[test]
    fn test_noisy_distributional_combined() -> anyhow::Result<()> {
        let stream = test_stream();
        // Noisy and distributional are always enabled
        let config = trading_config_small_atoms(8);
        let num_atoms = config.num_atoms;

        let mut net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;
        net.reset_noise()?;

        let state = GpuTensor::randn(&[2, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        // Should have distributional outputs
        assert!(output.advantage_log_probs.is_some());
        assert!(output.value_log_probs.is_some());

        let adv_lp = output
            .advantage_log_probs
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("expected advantage_log_probs"))?;
        assert_eq!(
            adv_lp.get(0).map(|t| t.shape().to_vec()),
            Some(vec![2, 5, num_atoms])
        );

        // Expected Q should still be finite
        let q_vals = output
            .advantages
            .get(0)
            .ok_or_else(|| anyhow::anyhow!("missing adv 0"))?
            .to_host(&stream)?;
        for &v in &q_vals {
            assert!(v.is_finite(), "NoisyDistributional Q should be finite");
        }

        Ok(())
    }

    // ======================================================================
    // Feature 3: State Dim Alignment tests
    // ======================================================================

    #[test]
    fn test_from_dqn_params_no_device() {
        let config =
            BranchingConfig::from_dqn_params(45, &[256, 128], 64, 0.01, None, vec![5, 3, 3]);
        assert_eq!(
            config.state_dim, 45,
            "Without device, state_dim should be unchanged"
        );
    }

    #[test]
    fn test_from_dqn_params_gpu_alignment() {
        let stream = test_stream();
        let config = BranchingConfig::from_dqn_params(
            45,
            &[256, 128],
            64,
            0.01,
            Some(&stream),
            vec![5, 3, 3],
        );
        assert_eq!(
            config.state_dim, 48,
            "CUDA device should align state_dim to next multiple of 8"
        );
    }

    #[test]
    fn test_alignment_function_directly() {
        // On CUDA, align to next multiple of 8 for tensor core HMMA dispatch
        let aligned_gpu = (45 + 7) & !7;
        assert_eq!(aligned_gpu, 48); // 45 -> 48 (next multiple of 8)

        // Already aligned values remain unchanged
        let aligned_gpu_8 = (48 + 7) & !7;
        assert_eq!(aligned_gpu_8, 48);

        let aligned_gpu_1 = (1 + 7) & !7;
        assert_eq!(aligned_gpu_1, 8); // 1 -> 8 (next multiple of 8)
    }

    // ======================================================================
    // Distributional + scalar interop tests
    // ======================================================================

    #[test]
    fn test_distributional_max_aggregate_q() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_small_atoms(8);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[3, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        let q_max = BranchingDuelingQNetwork::max_aggregate_q(&output, &stream)?;
        assert_eq!(q_max.shape(), &[3]);

        let greedy = BranchingDuelingQNetwork::greedy_branch_actions_batch(&output, &stream)?;
        let q_greedy =
            BranchingDuelingQNetwork::aggregate_q_for_actions(&output, &greedy, &stream)?;

        let max_vec = q_max.to_host(&stream)?;
        let greedy_vec = q_greedy.to_host(&stream)?;
        for (m, g) in max_vec.iter().zip(greedy_vec.iter()) {
            assert!(
                (m - g).abs() < 1e-5,
                "Distributional max Q ({m}) should equal greedy Q ({g})"
            );
        }
        Ok(())
    }

    #[test]
    fn test_distributional_greedy_actions() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_small_atoms(8);
        let net = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        let state = GpuTensor::randn(&[1, 8], 1.0, &stream)?;
        let output = net.forward_branches_eval(&state)?;

        let actions = BranchingDuelingQNetwork::greedy_branch_actions(&output, &stream)?;
        assert_eq!(actions.len(), 3);
        assert!(
            *actions.get(0).unwrap_or(&99) < 5,
            "Distributional: Exposure action should be 0-4"
        );
        assert!(
            *actions.get(1).unwrap_or(&99) < 3,
            "Distributional: Order action should be 0-2"
        );
        assert!(
            *actions.get(2).unwrap_or(&99) < 3,
            "Distributional: Urgency action should be 0-2"
        );
        Ok(())
    }

    #[test]
    fn test_distributional_weight_copy() -> anyhow::Result<()> {
        let stream = test_stream();
        let config = trading_config_small_atoms(8);
        let net1 = BranchingDuelingQNetwork::new(config.clone(), Arc::clone(&stream))?;
        let mut net2 = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        net2.copy_weights_from(&net1)?;

        let state = GpuTensor::full(&[1, 8], 1.0, &stream)?;
        let out1 = net1.forward_branches_eval(&state)?;
        let out2 = net2.forward_branches_eval(&state)?;

        // Expected values should match
        let v1_flat = out1.value.to_host(&stream)?;
        let v2_flat = out2.value.to_host(&stream)?;
        let v1_val = v1_flat.first().copied().unwrap_or(f32::NAN);
        let v2_val = v2_flat.first().copied().unwrap_or(f32::NAN);
        assert!(
            (v1_val - v2_val).abs() < 1e-4,
            "Distributional values should match after copy: {} vs {}",
            v1_val,
            v2_val
        );
        Ok(())
    }

    #[test]
    fn test_all_trainable_vars_includes_noisy() -> anyhow::Result<()> {
        let stream = test_stream();
        // NoisyLinear vars must be in all_trainable_vars() for the optimizer to update them
        let config = trading_config_small_atoms(8);

        let net = BranchingDuelingQNetwork::new(config, stream)?;

        let varmap_only = net.vars().all_vars()?;
        let all_vars = net.all_trainable_vars();
        let sigma_only = net.noisy_vars_ordered();

        // GpuVarStore has shared encoder (2 layers x 2 = 4) + NoisyLinear mu
        // (8 layers x weight_mu + bias_mu = 16) = 20 total.
        // Sigma vars are stored directly in NoisyLinear (not in GpuVarStore).
        assert_eq!(
            varmap_only.len(), 20,
            "GpuVarStore should have shared encoder (4) + NoisyLinear mu (16) = 20"
        );
        // noisy_vars_ordered returns only sigma vars: 8 layers x 2 = 16
        assert_eq!(
            sigma_only.len(), 16,
            "8 NoisyLinear layers x 2 sigma vars each"
        );
        // all_trainable_vars = GpuVarStore (shared + mu) + sigma
        assert_eq!(
            all_vars.len(),
            varmap_only.len() + sigma_only.len(),
            "all_trainable_vars should combine GpuVarStore ({}) + sigma ({})",
            varmap_only.len(),
            sigma_only.len()
        );
        // Total: 20 (shared + mu) + 16 sigma = 36
        assert_eq!(all_vars.len(), 36, "20 (shared + mu) + 16 sigma = 36 total");

        Ok(())
    }

    #[test]
    fn test_noisy_weight_copy() -> anyhow::Result<()> {
        let stream = test_stream();
        // Verify copy_weights_from syncs ALL vars (GpuVarStore mu + standalone sigma)
        let config = trading_config_small_atoms(8);

        let net1 = BranchingDuelingQNetwork::new(config.clone(), Arc::clone(&stream))?;
        let mut net2 = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?;

        // Use randn input with larger variance to amplify weight differences
        let state = GpuTensor::randn(&[4, 8], 5.0, &stream)?;

        // Before copy: outputs SHOULD differ (random mu init), but under BF16
        // quantization on CUDA, small Xavier init differences can round to zero.
        let out1 = net1.forward_branches_eval(&state)?;
        let out2 = net2.forward_branches_eval(&state)?;
        let v1 = out1.value.to_host(&stream)?;
        let v2 = out2.value.to_host(&stream)?;
        let diff_before: f32 = v1.iter().zip(v2.iter()).map(|(a, b)| (a - b).powi(2)).sum();
        if diff_before < 1e-6 {
            tracing::warn!("Before copy, outputs identical under BF16 (diff={})", diff_before);
        }

        net2.copy_weights_from(&net1)?;

        // After copy: outputs should match (all weights synced)
        // GPU comparison: sub → abs → stats.max → single scalar DtoH
        let out1a = net1.forward_branches_eval(&state)?;
        let out2a = net2.forward_branches_eval(&state)?;
        let diff_tensor = out1a.value.sub(&out2a.value, &stream)
            .map_err(|e| anyhow::anyhow!("GPU sub: {e}"))?;
        let elem_k = ml_core::cuda_autograd::elementwise::get_or_compile(&stream)?;
        let abs_diff = elem_k.abs(diff_tensor.cuda_data(), diff_tensor.numel())
            .map_err(|e| anyhow::anyhow!("GPU abs: {e}"))?;
        let reductions = ml_core::cuda_autograd::ReductionKernels::new(&stream)?;
        let stats = reductions.stats(&abs_diff, diff_tensor.numel())
            .map_err(|e| anyhow::anyhow!("GPU stats: {e}"))?;
        assert!(stats.max < 1e-3, "After copy, outputs should match: max_abs_diff={}", stats.max);

        // Also verify sigma vars were copied — GPU sub → abs → stats.max per var pair
        let s1 = net1.noisy_vars_ordered();
        let s2 = net2.noisy_vars_ordered();
        assert_eq!(s1.len(), s2.len());
        for (a, b) in s1.iter().zip(s2.iter()) {
            assert_eq!(a.len(), b.len());
            let diff_data = elem_k.binary(
                &a, &b, a.len(),
                1, // op=1 is sub (a - b)
            ).map_err(|e| anyhow::anyhow!("GPU sigma sub: {e}"))?;
            let abs_data = elem_k.abs(&diff_data, a.len())
                .map_err(|e| anyhow::anyhow!("GPU sigma abs: {e}"))?;
            let var_stats = reductions.stats(&abs_data, a.len())
                .map_err(|e| anyhow::anyhow!("GPU sigma stats: {e}"))?;
            assert!(var_stats.max < 1e-5, "Sigma var mismatch: max_abs_diff={}", var_stats.max);
        }

        Ok(())
    }

    #[test]
    fn test_config_serde_defaults() -> anyhow::Result<()> {
        // Verify serde defaults for new fields
        let json = r#"{
            "state_dim": 16,
            "branch_sizes": [5, 3, 3],
            "shared_hidden_dims": [64, 32],
            "value_hidden_dim": 16,
            "branch_hidden_dim": 16,
            "leaky_relu_alpha": 0.01
        }"#;
        let config: BranchingConfig = serde_json::from_str(json)?;
        assert_eq!(config.num_atoms, 51);
        assert!((config.v_min - (-25.0)).abs() < 1e-6);
        assert!((config.v_max - 25.0).abs() < 1e-6);
        assert!((config.noisy_sigma_init - 0.5).abs() < 1e-6);
        assert!((config.dropout_rate - 0.0).abs() < 1e-6);
        Ok(())
    }
}
