//! ACTUAL Working Deep Q-Network Implementation
//!
//! This module provides a complete, working DQN implementation with:
//! - Real mathematical operations using candle-core v0.9.1
//! - Experience replay buffer with proper memory management
//! - Epsilon-greedy exploration with decay
//! - Target network updates with soft/hard copying
//! - Proper Q-learning update with Bellman equation
//! - NO productions, todo!(), or unimplemented!() macros

use std::collections::VecDeque;
use std::sync::{Arc, Mutex};

// Xavier init used by branching network; noisy layers use their own initialization.
// (Arc imported above with Mutex)
use cudarc::driver::CudaStream;
use ml_core::cuda_autograd::{GpuTensor, GpuAdamW};
use ml_core::native_types::NativeDType;
use rand::{thread_rng, Rng};
use serde::{Deserialize, Serialize};
use common::metrics::training_metrics;
use tracing::debug;

use super::{Experience, ExposureLevel, FactoredAction, OrderType, Urgency};
use ml_core::MLError;

/// Configuration for the `DQN`
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DQNConfig {
    /// Number of actions
    pub num_actions: usize,
    /// Hidden layer dimensions
    pub hidden_dims: Vec<usize>,
    /// Learning rate
    pub learning_rate: f64,
    /// Discount factor (gamma)
    pub gamma: f32,
    /// Exploration parameters
    pub epsilon_start: f32,
    pub epsilon_end: f32,
    pub epsilon_decay: f32,
    /// Experience replay parameters
    pub replay_buffer_capacity: usize,
    /// Original buffer capacity before `AutoReplaySizer` (for gradient collapse warmup guard)
    pub collapse_warmup_capacity: usize,
    pub batch_size: usize,
    pub min_replay_size: usize,
    /// Target network update frequency
    pub target_update_freq: usize,
    /// Delta parameter for Huber loss
    pub huber_delta: f32,
    /// `LeakyReLU` negative slope (alpha) to prevent dead neurons
    pub leaky_relu_alpha: f64,
    /// Gradient clipping max norm (Wave 11 Bug #1 fix)
    pub gradient_clip_norm: f64,

    // Target update configuration: Cosine-annealed EMA (BYOL/MoCo v3)
    // τ(t) = τ_final - (τ_final - τ_base) * (cos(π·t/T) + 1) / 2
    // Early training: τ ≈ τ_base (fast adaptation) → Late training: τ ≈ τ_final (stability)
    /// Base EMA coefficient -- used early in training for fast target tracking
    pub tau: f64,
    /// Final EMA coefficient -- approached late in training for stability.
    /// Set equal to `tau` to disable annealing (reverts to fixed Polyak).
    pub tau_final: f64,
    /// Total training steps for cosine schedule (T). 0 = use fixed tau (no annealing).
    pub tau_anneal_steps: u64,

    // Rainbow DQN warmup period
    /// Number of steps to collect experiences with random exploration before training begins
    /// Rainbow DQN standard: 80,000 steps (prevents early overfitting to sparse data)
    pub warmup_steps: usize,

    // N-Step Returns configuration
    /// Number of steps for multi-step TD learning (n-step returns)
    /// n=1: standard single-step TD (default DQN)
    /// n>1: multi-step returns (better credit assignment, used in Rainbow DQN)
    /// Rainbow DQN standard: 3 steps
    pub n_steps: usize,

    // Prioritized Experience Replay (PER) configuration
    /// Initial trading capital (for portfolio tracking)
    pub initial_capital: f64,
    /// PER alpha parameter (prioritization exponent)
    pub per_alpha: f64,
    /// PER beta start value (importance sampling weight)
    pub per_beta_start: f64,
    /// PER beta maximum value
    pub per_beta_max: f64,
    /// Number of steps to anneal beta from start to max
    pub per_beta_annealing_steps: usize,
    /// Maximum memory allocation for GPU PER replay buffer (bytes).
    /// Computed at runtime from detected GPU VRAM via
    /// `GpuHardwareInfo::per_max_buffer_bytes()`.
    /// Overridden by `DqnTrainer` constructor on CUDA.
    pub per_max_memory_bytes: usize,

    // Dueling Networks configuration
    /// Hidden dimension for dueling advantage stream
    pub dueling_hidden_dim: usize,

    // Distributional RL (C51) configuration (always enabled)
    /// Number of atoms for value distribution
    pub num_atoms: usize,
    /// Minimum value for distribution support
    pub v_min: f32,
    /// Maximum value for distribution support
    pub v_max: f32,

    // Noisy Networks configuration (always enabled)
    /// Initial standard deviation for noisy layers
    pub noisy_sigma_init: f64,

    // Q-value clipping is unconditional in forward pass; gradients still flow through
    // because clamp is applied after forward (inference only).
    /// Minimum Q-value bound (clipped during forward pass)
    /// Default: -500.0 (based on portfolio normalization baseline ±375 × 1.33 safety margin)
    pub q_value_clip_min: f32,
    /// Maximum Q-value bound (clipped during forward pass)
    /// Default: 500.0 (based on portfolio normalization baseline ±375 × 1.33 safety margin)
    pub q_value_clip_max: f32,

    // WAVE 23: Early Stopping P0 Fix #1 - Adaptive gradient collapse threshold
    /// Gradient collapse threshold multiplier (learning-rate aware)
    /// Threshold = `learning_rate` × `gradient_collapse_multiplier`
    /// Default: 100.0 (was 0.1, causing 1000+ false alarms per trial)
    ///
    /// Rationale:
    /// - Typical LR: 1e-5 to 1e-4
    /// - Old threshold (0.1): 1e-6 to 1e-5 (TOO STRICT, false alarms)
    /// - New threshold (100.0): 1e-3 to 1e-2 (appropriate for healthy gradients)
    ///
    /// Expected Impact: Eliminate 1000+ false gradient collapse warnings per trial
    pub gradient_collapse_multiplier: f64,

    /// Number of consecutive epochs with gradient collapse before early stopping
    /// Default: 5 (prevents stopping on transient spikes)
    pub gradient_collapse_patience: usize,

    /// SAC-style entropy coefficient for Q-value softmax regularization.
    /// Adds -coeff * H(softmax(Q)) to the loss, preventing action collapse.
    /// Default: 0.001. Range for hyperopt: [0.001, 0.1].
    pub entropy_coefficient: f64,

    /// Minimum epsilon for exploration when noisy nets are enabled.
    /// Standard noisy nets set epsilon=0.0, but 45-action spaces need a floor.
    /// Default: 0.05. Range for hyperopt: [0.02, 0.10].
    pub noisy_epsilon_floor: f32,

    /// Coefficient β for count-based exploration bonus.
    /// Adds bonus = β × sqrt(ln(N) / (1 + `n_a`)) to Q-values before argmax.
    /// Set to 0.0 to disable UCB exploration entirely (production inference path).
    /// Default: 0.1. Range for hyperopt: [0.01, 0.5].
    pub count_bonus_coefficient: f64,

    // Conservative Q-Learning (CQL) for offline RL (Kumar et al. 2020)
    pub cql_alpha: f64,

    // IQN (Implicit Quantile Networks) for distributional RL (Dabney et al. 2018b).
    // Production training routes IQN unconditionally through `cuda_pipeline::gpu_iqn_head`
    // + `iqn_dual_head_kernel`; this `iqn_lambda` is the weight on the IQN aux head in
    // the dual C51+IQN loss (consumed by `cuda_pipeline/gpu_dqn_trainer.rs`).
    /// Lambda weight for IQN loss in dual-head ensemble.
    pub iqn_lambda: f32,

    // Branching DQN (Phase C+): independent Q-heads for factored action space
    /// Hidden dimension for branching advantage heads.
    /// Independent from `dueling_hidden_dim` so hyperopt can tune them separately.
    /// Default: 128.
    pub branch_hidden_dim: usize,

    /// Number of order type actions for branching DQN (e.g. `Market`, `LimitMaker`, `IoC`).
    /// Default: 3. Flows to `CountBonus` and `BranchingConfig`.
    pub num_order_types: usize,

    /// Number of urgency level actions for branching DQN (e.g. `Patient`, `Normal`, `Aggressive`).
    /// Default: 3. Flows to `CountBonus` and `BranchingConfig`.
    pub num_urgency_levels: usize,

    /// Hidden layer width for curiosity forward dynamics model.
    /// Default: 128. Must match `CUR_HIDDEN` in CUDA kernel.
    pub curiosity_hidden_dim: usize,

    /// Number of market features for curiosity model input/output.
    /// Default: 42 (`MARKET_DIM`). Must match CUDA `common_device_functions.cuh`.
    pub curiosity_market_dim: usize,

    // CVaR (Conditional Value at Risk): risk-aware quantile (consumed by cuda_pipeline `cvar_alpha_buf`).
    pub cvar_alpha: f32,

    // Agent-level trading parameters (from legacy agent::DQNConfig)
    /// BUG #7 FIX: Minimum profit threshold (profit must exceed cost * factor, default 1.5 = 50% margin)
    pub minimum_profit_factor: f32,
    /// WAVE 30: L2 weight decay for `AdamW` optimizer (default: 1e-4, range: [1e-5, 1e-3])
    /// Prevents overfitting by penalizing large weights via L2 regularization.
    pub weight_decay: f64,

    /// Adam optimizer epsilon.
    pub adam_epsilon: f64,

    /// Dropout rate for shared layers in dueling network (0.0 = disabled)
    /// Applied after each shared layer activation during training forward pass only.
    pub dropout_rate: f64,

}

impl Default for DQNConfig {
    fn default() -> Self {
        Self {
            num_actions: 4,  // 4 direction levels (Short/Hold/Long/Flat) — 4-branch refactor
            hidden_dims: vec![256, 256],
            learning_rate: 3e-5,   // Conservative default for stable training (matches conservative())
            gamma: 0.95,
            epsilon_start: 1.0,
            epsilon_end: 0.01,
            epsilon_decay: 0.995,
            replay_buffer_capacity: 100_000,
            collapse_warmup_capacity: 100_000,
            batch_size: 1024,  // H100 target (AutoBatchSizer caps for smaller GPUs)
            min_replay_size: 1000,
            target_update_freq: 1000,
            huber_delta: 1.0,
            leaky_relu_alpha: 0.01,
            gradient_clip_norm: 10.0,  // 2025 best practice: balanced clipping (not too aggressive)
            tau: 0.005,
            tau_final: 0.0005,  // 10x slower at convergence for offline RL stability
            tau_anneal_steps: 100_000,  // cosine anneal over 100K gradient steps
            warmup_steps: 1000,
            n_steps: 1,
            initial_capital: 100_000.0,

            per_alpha: 0.6,
            per_beta_start: 0.4,
            per_beta_max: 1.0,
            per_beta_annealing_steps: 100_000,
            per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
            dueling_hidden_dim: 128,
            num_atoms: 51,
            // v_min/v_max computed from reward_scale=10, gamma=0.95:
            // Q* ~ 10 / (1-0.95) * 1.2 = 240, clamped to [20, 300] → ±240
            v_min: -(10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            v_max: (10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            noisy_sigma_init: 0.5,
            q_value_clip_min: -200.0,  // Reward v6: tighter than old -500 but covers v_range + safety margin
            q_value_clip_max: 200.0,
            gradient_collapse_multiplier: 2.0,
            gradient_collapse_patience: 100,

            // Entropy regularization: prevent 5-action exposure collapse
            entropy_coefficient: 0.001, // Must be ≤ reward magnitude (~0.001). Old 0.1 was 100x reward → entropy dominated learning.
            noisy_epsilon_floor: 0.05,
            count_bonus_coefficient: 0.1,

            // CQL: Enabled by default for offline training
            cql_alpha: 0.1,  // Reduced from 1.0: full strength added ~ln(5)=1.6 loss penalty, can crush Q-value differentiation

            iqn_lambda: 0.25,

            // Branching: Enabled by default (45 factored actions via 3 independent heads)
            branch_hidden_dim: 128,
            num_order_types: 3,
            num_urgency_levels: 3,

            // Curiosity model dimensions
            curiosity_hidden_dim: 128,
            curiosity_market_dim: 42,

            cvar_alpha: 0.05,

            // Agent-level trading parameters
            minimum_profit_factor: 1.5, // BUG #7 FIX: 50% margin above breakeven
            weight_decay: 1e-4, // WAVE 30: Standard L2 regularization strength
            adam_epsilon: 1e-8, // f32 Adam (master weights are f32)
            dropout_rate: 0.0, // Disabled by default; hyperopt wires from dropout_initial

            // Mixed precision: disabled by default (auto-detected at runtime)

        }
    }
}

impl DQNConfig {
    /// Compute a SHA-256 hash of the architecture-critical config fields.
    ///
    /// These fields determine tensor shapes in the network. Loading a checkpoint
    /// saved with different architecture parameters will silently produce garbage
    /// or panic on shape mismatch. This hash is embedded in safetensors metadata
    /// on save and validated on load.
    ///
    /// Hashed fields: `STATE_DIM` (constant), `num_actions`, `hidden_dims` (length + values),
    /// `dueling_hidden_dim`, `num_atoms`, `iqn_lambda`.
    pub fn architecture_hash(&self) -> String {
        use sha2::{Digest, Sha256};
        let mut hasher = Sha256::new();
        hasher.update(ml_core::state_layout::STATE_DIM.to_le_bytes());
        hasher.update(self.num_actions.to_le_bytes());
        hasher.update(self.hidden_dims.len().to_le_bytes());
        for &dim in &self.hidden_dims {
            hasher.update(dim.to_le_bytes());
        }
        hasher.update(self.dueling_hidden_dim.to_le_bytes());
        hasher.update(self.num_atoms.to_le_bytes());
        hasher.update(self.num_order_types.to_le_bytes());
        hasher.update(self.num_urgency_levels.to_le_bytes());
        hasher.update(self.iqn_lambda.to_le_bytes());
        hex::encode(hasher.finalize())
    }

    /// Build safetensors metadata for checkpoint files.
    ///
    /// Returns a `HashMap<String, String>` containing the architecture hash
    /// and human-readable config summary for debugging shape mismatches.
    pub fn checkpoint_metadata(&self) -> std::collections::HashMap<String, String> {
        let mut meta = std::collections::HashMap::new();
        meta.insert("dqn.arch_hash".to_owned(), self.architecture_hash());
        meta.insert("dqn.state_dim".to_owned(), ml_core::state_layout::STATE_DIM.to_string());
        meta.insert("dqn.num_actions".to_owned(), self.num_actions.to_string());
        meta.insert("dqn.hidden_dims".to_owned(), format!("{:?}", self.hidden_dims));
        meta.insert("dqn.num_order_types".to_owned(), self.num_order_types.to_string());
        meta.insert("dqn.num_urgency_levels".to_owned(), self.num_urgency_levels.to_string());
        meta.insert("dqn.dueling_hidden_dim".to_owned(), self.dueling_hidden_dim.to_string());
        meta.insert("dqn.num_atoms".to_owned(), self.num_atoms.to_string());
        meta.insert("dqn.iqn_lambda".to_owned(), self.iqn_lambda.to_string());
        meta.insert("dqn.gamma".to_owned(), self.gamma.to_string());
        meta
    }

    /// Load `DQNConfig` from a safetensors checkpoint file.
    ///
    /// Reads architecture-critical fields (`state_dim`, `num_actions`, `hidden_dims`,
    /// Rainbow flags) from the file's metadata. Non-architecture fields (LR,
    /// epsilon, buffer capacity) use defaults since they don't affect inference.
    pub fn from_safetensors_file(path: &std::path::Path) -> Result<Self, ml_core::MLError> {
        let raw_bytes = std::fs::read(path).map_err(|e| {
            ml_core::MLError::CheckpointError(format!("Failed to read safetensors: {}", e))
        })?;
        let (_header_size, st_metadata) =
            safetensors::SafeTensors::read_metadata(&raw_bytes).map_err(|e| {
                ml_core::MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e))
            })?;

        let Some(meta) = st_metadata.metadata() else {
            return Err(ml_core::MLError::CheckpointError(
                "Safetensors checkpoint has no metadata — re-train to embed config".to_owned(),
            ));
        };

        Self::from_checkpoint_metadata(meta)
    }

    /// Reconstruct a `DQNConfig` from safetensors checkpoint metadata.
    ///
    /// Reads architecture-critical fields (`state_dim`, `num_actions`, `hidden_dims`,
    /// Rainbow flags) from the metadata map. Non-architecture fields (LR, epsilon,
    /// buffer capacity) use defaults since they don't affect inference.
    ///
    /// Returns `Err` if required fields are missing or malformed.
    pub fn from_checkpoint_metadata(
        metadata: &std::collections::HashMap<String, String>,
    ) -> Result<Self, ml_core::MLError> {
        let get = |key: &str| -> Result<&str, ml_core::MLError> {
            metadata.get(key).map(|s| s.as_str()).ok_or_else(|| {
                ml_core::MLError::CheckpointError(format!(
                    "Checkpoint metadata missing '{}' — re-train to embed config", key
                ))
            })
        };

        let saved_state_dim: usize = get("dqn.state_dim")?
            .parse()
            .map_err(|e| ml_core::MLError::CheckpointError(format!("Bad dqn.state_dim: {}", e)))?;
        if saved_state_dim != ml_core::state_layout::STATE_DIM {
            return Err(ml_core::MLError::CheckpointError(format!(
                "Checkpoint state_dim={} but STATE_DIM={}. Old checkpoints are invalid.",
                saved_state_dim, ml_core::state_layout::STATE_DIM,
            )));
        }
        let num_actions: usize = get("dqn.num_actions")?
            .parse()
            .map_err(|e| ml_core::MLError::CheckpointError(format!("Bad dqn.num_actions: {}", e)))?;

        // Parse hidden_dims from debug format "[256, 256]"
        let hidden_dims_str = get("dqn.hidden_dims")?;
        let hidden_dims: Vec<usize> = hidden_dims_str
            .trim_matches(|c| c == '[' || c == ']')
            .split(',')
            .filter(|s| !s.trim().is_empty())
            .map(|s| s.trim().parse::<usize>())
            .collect::<Result<Vec<_>, _>>()
            .map_err(|e| ml_core::MLError::CheckpointError(format!("Bad dqn.hidden_dims: {}", e)))?;

        let parse_usize = |key: &str| -> Result<usize, ml_core::MLError> {
            get(key)?.parse().map_err(|e| ml_core::MLError::CheckpointError(format!("Bad {}: {}", key, e)))
        };

        // Architecture flags (use_dueling / use_distributional / use_noisy_nets / use_branching
        // / use_iqn / iqn_embedding_dim / iqn_num_quantiles / iqn_kappa) from older checkpoints
        // are silently dropped — every one of these is unconditionally enabled in the live
        // architecture. The cuda_pipeline dual head owns IQN; branching + dueling + distributional
        // + NoisyNets are the only architecture, no alternative code path remains.
        let num_order_types = metadata.get("dqn.num_order_types")
            .and_then(|s| s.parse().ok())
            .unwrap_or(3); // backward compat
        let num_urgency_levels = metadata.get("dqn.num_urgency_levels")
            .and_then(|s| s.parse().ok())
            .unwrap_or(3); // backward compat
        let dueling_hidden_dim = parse_usize("dqn.dueling_hidden_dim")?;
        let num_atoms = parse_usize("dqn.num_atoms")?;
        let iqn_lambda: f32 = metadata.get("dqn.iqn_lambda")
            .and_then(|s| s.parse().ok())
            .unwrap_or(0.25); // backward compat: default to 0.25

        let gamma: f32 = metadata.get("dqn.gamma")
            .and_then(|s| s.parse().ok())
            .unwrap_or(0.95);

        Ok(Self {
            num_actions,
            hidden_dims,
            dueling_hidden_dim,
            num_atoms,
            num_order_types,
            num_urgency_levels,
            iqn_lambda,
            gamma,
            ..Self::default()
        })
    }

    /// Validate that a checkpoint's metadata matches this config's architecture.
    ///
    /// Returns `Ok(())` if the hash matches. Returns an error if metadata is
    /// missing or the architecture hash differs.
    pub fn validate_checkpoint_metadata(
        &self,
        metadata: &Option<std::collections::HashMap<String, String>>,
    ) -> Result<(), ml_core::MLError> {
        let Some(meta) = metadata else {
            return Err(ml_core::MLError::CheckpointError(
                "Checkpoint has no architecture metadata - cannot validate. \
                 Re-train to generate a checkpoint with embedded config hash."
                    .to_owned(),
            ));
        };

        let Some(saved_hash) = meta.get("dqn.arch_hash") else {
            return Err(ml_core::MLError::CheckpointError(
                "Checkpoint metadata missing dqn.arch_hash field - cannot validate. \
                 Re-train to generate a checkpoint with embedded config hash."
                    .to_owned(),
            ));
        };

        let current_hash = self.architecture_hash();
        if *saved_hash != current_hash {
            let saved_actions = meta.get("dqn.num_actions").map(|s| s.as_str()).unwrap_or("?");
            let saved_hidden = meta.get("dqn.hidden_dims").map(|s| s.as_str()).unwrap_or("?");
            return Err(ml_core::MLError::CheckpointError(format!(
                "Architecture mismatch: checkpoint was saved with num_actions={}, \
                 hidden_dims={} (hash={}) but current config has num_actions={}, \
                 hidden_dims={:?} (hash={}) — STATE_DIM={}",
                saved_actions, saved_hidden, saved_hash,
                self.num_actions, self.hidden_dims, current_hash,
                ml_core::state_layout::STATE_DIM,
            )));
        }

        Ok(())
    }

    /// Create `DQN` config from central configuration system
    ///
    /// CRITICAL: Eliminates dangerous hardcoded defaults
    pub fn from_config_manager(
        _config_manager: &config::ConfigManager,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        // Use emergency defaults since specific DQN configs may not be available
        tracing::warn!(
            "Using emergency DQN config defaults - DQN configs not available in ServiceConfig"
        );
        Ok(Self::emergency_safe_defaults())
    }

    /// Aggressive configuration for maximum exploration and learning
    ///
    /// Used for testing noisy networks and rapid experimentation
    pub fn aggressive() -> Self {
        Self {
            num_actions: 3,
            hidden_dims: vec![256, 256, 128],  // P1 FIX: Increased capacity for 45-action space (2x output params/action)
            learning_rate: 1e-3,
            gamma: 0.99,
            epsilon_start: 1.0,
            epsilon_end: 0.1,
            epsilon_decay: 0.99,
            replay_buffer_capacity: 50000,
            batch_size: 64,
            min_replay_size: 500,
            target_update_freq: 500,
            huber_delta: 10.0,  // Conservative default (hyperopt can scale to 15-40)
            leaky_relu_alpha: 0.01,
            gradient_clip_norm: 10.0,  // 2025 best practice: balanced clipping

            tau: 0.005,
            tau_final: 0.0005,
            tau_anneal_steps: 100_000,
            warmup_steps: 1000,
            n_steps: 3,
            initial_capital: 100000.0,

            per_alpha: 0.6,
            per_beta_start: 0.4,
            per_beta_max: 1.0,
            per_beta_annealing_steps: 100000,
            per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
            dueling_hidden_dim: 512,  // Wave 11.4: Large hidden dim for aggressive config
            num_atoms: 51,
            v_min: -25.0,  // DSR Q-values: rewards ±2 with gamma=0.92 → Q ≈ ±25
            v_max: 25.0,   // DSR Q-values: rewards ±2 with gamma=0.92 → Q ≈ ±25
            noisy_sigma_init: 0.5,

            // BUG #37 FIX: Q-value clipping (prevents step-level explosions)
            q_value_clip_min: -1000.0,  // Wider bounds for aggressive exploration
            q_value_clip_max: 1000.0,

            // WAVE 23 P0 Fix #1: Adaptive gradient collapse threshold
            gradient_collapse_multiplier: 100.0,
            gradient_collapse_patience: 5,

            cql_alpha: 0.5,
            iqn_lambda: 0.25,
            cvar_alpha: 0.05,

            minimum_profit_factor: 1.5,
            weight_decay: 1e-4,
            adam_epsilon: 1e-8, // f32 Adam (master weights are f32)

            ..Default::default()
        }
    }

    /// Conservative production defaults (balanced safety + performance)
    ///
    /// Recommended for production use with proven DQN baseline
    pub fn conservative() -> Self {
        Self {
            num_actions: 3,
            hidden_dims: vec![256, 256, 128],  // P1 FIX: Increased capacity for 45-action space (2x output params/action)
            learning_rate: 3e-5,   // Conservative default for stable training
            gamma: 0.95,
            epsilon_start: 0.3,
            epsilon_end: 0.05,
            epsilon_decay: 0.995,
            replay_buffer_capacity: 10000,
            batch_size: 1024,  // H100 target (AutoBatchSizer caps for smaller GPUs)
            min_replay_size: 1000,
            target_update_freq: 1000,
            huber_delta: 10.0,  // Conservative default (hyperopt can scale to 15-40)
            leaky_relu_alpha: 0.01,
            gradient_clip_norm: 10.0,  // 2025 best practice: balanced clipping
            tau: 0.005,
            tau_final: 0.0005,
            tau_anneal_steps: 100_000,
            warmup_steps: 0,
            n_steps: 1, // Default to single-step TD (most stable)
            initial_capital: 100000.0,

            per_alpha: 0.6,
            per_beta_start: 0.4,
            per_beta_max: 1.0,
            per_beta_annealing_steps: 100000,
            per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
            dueling_hidden_dim: 64,
            num_atoms: 51,
            // v_min/v_max computed from reward_scale=10, gamma=0.95:
            // Q* ~ 10 / (1-0.95) * 1.2 = 240, clamped to [20, 300] → ±240
            v_min: -(10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            v_max: (10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            noisy_sigma_init: 0.5,

            // BUG #37 FIX: Q-value clipping (prevents step-level explosions)
            q_value_clip_min: -200.0,  // Reward v6: tighter than old -500 but covers v_range + safety margin
            q_value_clip_max: 200.0,

            // WAVE 23 P0 Fix #1: Adaptive gradient collapse threshold
            gradient_collapse_multiplier: 100.0,
            gradient_collapse_patience: 5,

            cql_alpha: 2.0,
            iqn_lambda: 0.25,
            // Conservative variant: tight tail (α=0.05) for risk-averse defaults;
            // CVaR action selection wired via numeric `cvar_alpha` (commit B)
            cvar_alpha: 0.05,

            minimum_profit_factor: 2.0, // Higher safety margin for conservative config
            weight_decay: 1e-4,
            adam_epsilon: 1e-8, // f32 Adam (master weights are f32)

            ..Default::default()
        }
    }

    /// EMERGENCY FALLBACK: Ultra-conservative `DQN` defaults
    ///
    /// WARNING: These defaults prioritize safety over performance
    pub fn emergency_safe_defaults() -> Self {
        tracing::warn!("Using emergency DQN defaults - check configuration system immediately!");
        Self {
            num_actions: 3,                  // Conservative action space
            hidden_dims: vec![256, 256, 128], // P1 FIX: Increased capacity for 45-action space (2x output params/action)
            learning_rate: 1e-5,             // Very conservative learning rate
            gamma: 0.9,                      // Conservative discount factor
            epsilon_start: 0.1,              // Low exploration to prevent erratic behavior
            epsilon_end: 0.01,               // Minimal exploration
            epsilon_decay: 0.99,             // Fast decay to reach stable exploitation
            replay_buffer_capacity: 2048,    // Min 1024 for GPU PER floor
            collapse_warmup_capacity: 2048,
            batch_size: 4,                   // Very small batch size
            min_replay_size: 100,            // Minimal replay requirement
            target_update_freq: 100,         // Frequent updates for stability
            huber_delta: 10.0,               // Conservative default (hyperopt can scale to 15-40)
            leaky_relu_alpha: 0.01,          // Standard LeakyReLU alpha
            gradient_clip_norm: 10.0,  // 2025 best practice: balanced clipping

            // WAVE 16 (Agent 36): Target update defaults (SOFT UPDATES for gradient stability)
            tau: 0.005,
            tau_final: 0.001,  // Very conservative annealing for emergency mode
            tau_anneal_steps: 50_000,

            // Rainbow DQN warmup period
            warmup_steps: 0, // No warmup for emergency mode (safety first)

            // N-Step Returns (Wave 11.5)
            n_steps: 1, // Single-step TD for emergency defaults (most stable)

            // Rainbow DQN defaults (emergency mode still needs GPU PER on CUDA)
            initial_capital: 100000.0,

            per_alpha: 0.6,
            per_beta_start: 0.4,
            per_beta_max: 1.0,
            per_beta_annealing_steps: 100000,
            per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
            dueling_hidden_dim: 64,
            num_atoms: 51,
            // v_min/v_max computed from reward_scale=10, gamma=0.95:
            // Q* ~ 10 / (1-0.95) * 1.2 = 240, clamped to [20, 300] → ±240
            v_min: -(10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            v_max: (10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            noisy_sigma_init: 0.5,

            // BUG #37 FIX: Q-value clipping (prevents step-level explosions)
            q_value_clip_min: -200.0,  // Reward v6: tighter than old -500 but covers v_range + safety margin
            q_value_clip_max: 200.0,

            // WAVE 23 P0 Fix #1: Adaptive gradient collapse threshold
            gradient_collapse_multiplier: 100.0,
            gradient_collapse_patience: 5,

            entropy_coefficient: 0.001, // Must be ≤ reward magnitude (~0.001). Old 0.01 was 10x reward → entropy dominated learning.
            noisy_epsilon_floor: 0.05,
            count_bonus_coefficient: 0.1,

            cql_alpha: 5.0,
            iqn_lambda: 0.25,
            branch_hidden_dim: 128,
            num_order_types: 3,
            num_urgency_levels: 3,
            curiosity_hidden_dim: 128,
            curiosity_market_dim: 42,
            cvar_alpha: 0.05,

            minimum_profit_factor: 1.5,
            weight_decay: 1e-4,
            adam_epsilon: 1e-8, // f32 Adam (master weights are f32)
            dropout_rate: 0.0,

        }
    }
}

/// Result of computing gradients for a mini-batch without stepping the optimizer.
#[allow(missing_debug_implementations)]
pub struct GradientResult {
    /// Scalar loss value for this mini-batch.
    pub loss: f32,
    /// Gradient norm after clipping.
    pub grad_norm: f32,
    /// Gradient tensors keyed by parameter name.
    pub grads: std::collections::BTreeMap<String, GpuTensor>,
    /// TD errors for PER priority updates.
    pub td_errors: Vec<f32>,
    /// Replay buffer indices for PER priority updates.
    pub indices: Vec<usize>,
    /// GPU-resident TD errors (`GpuPrioritized` path).
    pub td_errors_gpu: Option<GpuTensor>,
    /// GPU-resident buffer indices (`GpuPrioritized` path).
    pub indices_gpu: Option<cudarc::driver::CudaSlice<u32>>,
}

/// Experience replay buffer for `DQN`
#[derive(Debug)]
pub struct ExperienceReplayBuffer {
    buffer: VecDeque<Experience>,
    capacity: usize,
}

impl ExperienceReplayBuffer {
    /// Create new replay buffer
    pub fn new(capacity: usize) -> Self {
        Self {
            buffer: VecDeque::with_capacity(capacity),
            capacity,
        }
    }

    /// Add experience to buffer
    pub fn push(&mut self, experience: Experience) {
        if self.buffer.len() >= self.capacity {
            self.buffer.pop_front();
        }
        self.buffer.push_back(experience);
    }

    /// Sample random batch of experiences
    pub fn sample(&self, batch_size: usize) -> Result<Vec<Experience>, MLError> {
        if self.buffer.len() < batch_size {
            return Err(MLError::TrainingError(format!(
                "Not enough experiences in buffer: {} < {}",
                self.buffer.len(),
                batch_size
            )));
        }

        let mut rng = thread_rng();
        let mut batch = Vec::with_capacity(batch_size);

        for _ in 0..batch_size {
            let idx = rng.gen_range(0..self.buffer.len());
            batch.push(self.buffer[idx].clone());
        }

        Ok(batch)
    }

    /// Get current buffer size
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    /// Check if buffer is empty
    pub fn is_empty(&self) -> bool {
        self.buffer.is_empty()
    }

    /// Check if buffer can sample
    pub fn can_sample(&self, min_size: usize) -> bool {
        self.buffer.len() >= min_size
    }

    /// Resize buffer capacity (adaptive buffer sizing)
    ///
    /// Keeps most recent experiences, drops oldest if shrinking.
    ///
    /// # Arguments
    ///
    /// * `new_capacity` - New buffer capacity
    ///
    /// # Returns
    ///
    /// * `Ok(())` - Resize successful
    /// * `Err` - Resize failed
    pub fn resize(&mut self, new_capacity: usize) -> Result<(), MLError> {
        if new_capacity == self.capacity {
            return Ok(()); // No change needed
        }
        
        // Keep most recent experiences if shrinking
        while self.buffer.len() > new_capacity {
            self.buffer.pop_front();
        }
        
        self.capacity = new_capacity;
        Ok(())
    }

    /// BUG #38 FIX: Clear all experiences from the buffer
    pub fn clear(&mut self) {
        self.buffer.clear();
    }
    
    /// Get current buffer capacity
    pub const fn capacity(&self) -> usize {
        self.capacity
    }
}

/// Deep Q-Network implementation
#[allow(missing_debug_implementations)]
pub struct DQN {
    /// `DQN` configuration
    pub config: DQNConfig,
    /// Dueling Q-network (legacy, unused when distributional+dueling hybrid is active)
    pub dueling_q_network: Option<super::dueling::DuelingQNetwork>,
    /// Dueling target network (legacy, unused when distributional+dueling hybrid is active)
    pub dueling_target_network: Option<super::dueling::DuelingQNetwork>,
    /// Distributional dueling Q-network (Wave 11.6)
    pub dist_dueling_q_network: Option<super::distributional_dueling::DistributionalDuelingQNetwork>,
    /// Distributional dueling target network (Wave 11.6: for stable training)
    pub dist_dueling_target_network: Option<super::distributional_dueling::DistributionalDuelingQNetwork>,
    /// Categorical distribution for distributional RL (always initialized)
    pub categorical_dist: Option<super::distributional::CategoricalDistribution>,
    /// GPU-resident PER replay buffer (public for trainer access)
    pub memory: super::replay_buffer_type::StagedGpuBuffer,
    /// Current exploration rate
    epsilon: f32,
    /// Training step counter (gradient updates)
    training_steps: u64,
    /// Total environment steps counter (includes warmup period)
    total_steps: u64,
    /// Optimizer for main network
    optimizer: Option<GpuAdamW>,
    /// CUDA stream for all GPU operations
    stream: Arc<CudaStream>,
    /// Gradient clipping max norm (Wave 11 Bug #1 fix)
    _gradient_clip_norm: f64,
    /// Recent actions for entropy penalty calculation (sliding window)
    recent_actions: VecDeque<FactoredAction>,
    /// N-step buffer for multi-step TD learning (Wave 11.5)
    nstep_buffer: Option<Arc<Mutex<super::NStepBuffer>>>,

    /// Branching DQN (Phase C+): 3 independent Q-heads for factored actions
    pub branching_q_network: Option<super::branching::BranchingDuelingQNetwork>,
    /// Branching target network for stable TD targets
    pub branching_target_network: Option<super::branching::BranchingDuelingQNetwork>,

    // WAVE 23 P0 Fix #2: Early stopping on gradient collapse
    /// Consecutive epochs with gradient collapse (resets when gradients recover)
    gradient_collapse_counter: usize,

    // WAVE 23 P0 Fix #4: Early stopping on Q-value divergence
    /// Consecutive epochs with Q-value divergence (resets when Q-values normalize)
    q_value_divergence_counter: usize,

    /// UCB count-based exploration bonus for action diversity
    count_bonus: super::count_bonus::CountBonus,

    /// When true, skip monitoring in `forward()` to avoid GPU->CPU sync in training hot path
    training_forward_active: bool,

    /// Noise sigma scale for annealed exploration (1.0 = default, decreases over training)
    noise_sigma_scale: f64,

    /// Mixture-of-Experts gate subnetwork: state[42] → 64 → 8 → softmax.
    /// Initialized to zero so initial g(s) = uniform 1/K for every state.
    /// `None` until Phase 3 wires it up.
    pub moe_gate: Option<super::moe::MoeGate>,

    /// 8 expert MLPs: each is h_s1[256] → 64 → h_s2[256] bottleneck.
    /// Independent Xavier init per expert for initial differentiation.
    /// `None` until Phase 3 wires it up.
    pub moe_experts: Option<Vec<super::moe::MoeExpert>>,
}

impl DQN {
    /// Create new `DQN` with auto-detected CUDA device.
    pub fn new(config: DQNConfig) -> Result<Self, MLError> {
        let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| {
            MLError::DeviceError(format!("Failed to open CUDA device 0: {e}"))
        })?;
        // cuMemAllocAsync disabled at cudarc CudaContext level — see vendor/cudarc comment.
        let stream = ctx.new_stream().map_err(|e| {
            MLError::DeviceError(format!("Failed to create CUDA stream: {e}"))
        })?;
        Self::new_with_stream(config, stream)
    }

    /// Create new `DQN` with a specific device (extracts CUDA stream).
    #[allow(clippy::cognitive_complexity, clippy::too_many_lines)]
    pub fn new_on_device(config: DQNConfig, device: ml_core::device::MlDevice) -> Result<Self, MLError> {
        let stream = Arc::clone(device.cuda_stream()?);
        Self::new_with_stream(config, stream)
    }

    /// Create new `DQN` with a specific CUDA stream.
    #[allow(clippy::cognitive_complexity, clippy::too_many_lines)]
    fn new_with_stream(config: DQNConfig, stream: Arc<CudaStream>) -> Result<Self, MLError> {
        // Ensure sync allocation for cublasLtMatmul compatibility
        // cuMemAllocAsync disabled at cudarc CudaContext level — see vendor/cudarc comment.
        if config.num_actions == 0 {
            return Err(MLError::ConfigError("DQN requires num_actions > 0".to_owned()));
        }

        // Distributional + Dueling hybrid (always enabled — mandatory Rainbow DQN component)
        let (dueling_q_network, dueling_target_network, dist_dueling_q_network, dist_dueling_target_network) = {
            let dist_dueling_config = super::distributional_dueling::DistributionalDuelingConfig::from_dqn_params(
                ml_core::state_layout::STATE_DIM,
                config.num_actions,
                config.num_atoms,
                &config.hidden_dims,
                config.dueling_hidden_dim,
                config.leaky_relu_alpha,
            );

            // Create main hybrid network
            let hybrid_net = super::distributional_dueling::DistributionalDuelingQNetwork::new(
                dist_dueling_config.clone(),
                Arc::clone(&stream),
            )?;

            // Create hybrid target network
            let hybrid_target = super::distributional_dueling::DistributionalDuelingQNetwork::new(
                dist_dueling_config,
                Arc::clone(&stream),
            )?;

            // Target weights don't matter — fused trainer initializes via flat DtoD.

            (None, None, Some(hybrid_net), Some(hybrid_target))
        };

        // Create categorical distribution (distributional always enabled)
        let categorical_dist = {
            let dist_config = super::distributional::DistributionalConfig {
                num_atoms: config.num_atoms,
                v_min: config.v_min as f64,  // CategoricalDistribution uses f64 internally
                v_max: config.v_max as f64,  // but creates F32 tensors for GPU compatibility
            };
            Some(super::distributional::CategoricalDistribution::new(&dist_config, &stream)?)
        };

        // GPU PER: on CUDA, allocate GPU-resident ring buffer (hard error on failure).
        // PER is mandatory — CUDA is required.
        let memory = super::replay_buffer_type::StagedGpuBuffer::try_with_halving(
            config.replay_buffer_capacity,
            config.per_alpha,
            config.per_beta_start,
            config.per_beta_max,
            config.per_beta_annealing_steps,
            config.per_max_memory_bytes,
            config.batch_size,
            &stream,
        )?;

        // Wave 11.5: Create n-step buffer if n_steps > 1
        let nstep_buffer = (config.n_steps > 1).then(|| {
            Arc::new(Mutex::new(super::NStepBuffer::new(
                config.n_steps,
                config.gamma as f64,
            )))
        });

        // Phase C+: Branching DQN (4 independent Q-heads — always enabled)
        let (branching_q_network, branching_target_network) = {
            let branch_sizes = vec![
                config.num_actions,        // direction: 3
                3,                         // magnitude: Small/Medium/Full
                config.num_order_types,    // order: 3
                config.num_urgency_levels, // urgency: 3
            ];
            let mut branching_config = super::branching::BranchingConfig::from_dqn_params(
                ml_core::state_layout::STATE_DIM,
                &config.hidden_dims,
                config.dueling_hidden_dim,
                config.leaky_relu_alpha,
                Some(&stream), // align_for_gpu on CUDA
                branch_sizes,
            );
            // Wire hyperopt-tunable branch_hidden_dim and distributional params from DQN config
            branching_config.branch_hidden_dim = config.branch_hidden_dim;
            branching_config.v_min = config.v_min;
            branching_config.v_max = config.v_max;
            branching_config.num_atoms = config.num_atoms;
            branching_config.noisy_sigma_init = config.noisy_sigma_init as f32;

            let branching_net = super::branching::BranchingDuelingQNetwork::new(
                branching_config.clone(),
                Arc::clone(&stream),
            )?;

            let branching_target = super::branching::BranchingDuelingQNetwork::new(
                branching_config,
                Arc::clone(&stream),
            )?;
            // Target weights don't matter — fused trainer initializes via flat DtoD.

            (Some(branching_net), Some(branching_target))
        };

        Ok(Self {
            epsilon: config.epsilon_start,
            dueling_q_network,
            dueling_target_network,
            dist_dueling_q_network,
            dist_dueling_target_network,
            categorical_dist,
            memory,
            training_steps: 0,
            total_steps: 0,
            optimizer: None,
            _gradient_clip_norm: config.gradient_clip_norm,
            count_bonus: super::count_bonus::CountBonus::new(
                config.count_bonus_coefficient,
                config.num_actions,
                config.num_order_types,
                config.num_urgency_levels,
            ),
            config,
            stream,
            recent_actions: VecDeque::with_capacity(100),
            nstep_buffer,
            branching_q_network,
            branching_target_network,
            gradient_collapse_counter: 0,
            q_value_divergence_counter: 0,
            training_forward_active: false,
            noise_sigma_scale: 1.0,
            moe_gate: None,
            moe_experts: None,
        })
    }

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

    /// Get state dimension (canonical constant)
    pub const fn get_state_dim(&self) -> usize {
        ml_core::state_layout::STATE_DIM
    }

    /// Check if noisy networks are enabled (always true -- noisy nets are unconditional)
    pub const fn is_using_noisy_nets(&self) -> bool {
        true
    }

    /// Check if dueling networks are enabled (Wave 11.4)
    ///
    /// Returns true if either:
    /// - Dueling-only architecture is active (`dueling_q_network.is_some()`)
    /// - Hybrid architecture is active (`dist_dueling_q_network.is_some()`)
    ///
    /// Dueling is always enabled (mandatory Rainbow DQN component).
    pub const fn is_using_dueling(&self) -> bool {
        self.dueling_q_network.is_some() || self.dist_dueling_q_network.is_some()
    }

    /// Forward pass through main network
    ///
    /// Architecture priority (Wave 11.6):
    /// 1. Hybrid (distributional + dueling) - highest priority
    /// 2. Dueling only
    /// 3. Standard Q-network (fallback)
    pub fn forward(&self, state: &GpuTensor) -> Result<GpuTensor, MLError> {
        // Auto-convert input to correct device and dtype
        let state = state.to_dtype(NativeDType::F32, &self.stream)
            .map_err(|e| MLError::ModelError(e.to_string()))?;
        let state = state
            .gpu_clone(&self.stream)
            .map_err(|e| MLError::ModelError(format!("Failed to move tensor to device: {}", e)))?;

        // Wave 11.6: Priority 1 - Hybrid (distributional + dueling)
        let q_values = if let Some(ref dist_dueling_net) = self.dist_dueling_q_network {
            // Forward returns [batch, num_actions, num_atoms]
            // Need to compute expected Q-values: Q(s,a) = Σ(z_i * p(s,a,z_i))
            let z_probs = dist_dueling_net.forward(&state)?;

            // Get atom values (support of distribution)
            let num_atoms = self.config.num_atoms;
            let v_min = self.config.v_min;
            let v_max = self.config.v_max;
            let delta_z = (v_max - v_min) / (num_atoms as f32 - 1.0);

            let batch_size = state.dim(0)?;
            let num_actions = self.config.num_actions;

            // GPU-native: build atoms tensor [batch, num_actions, num_atoms]
            let atoms_1d = GpuTensor::arange(0.0_f32, num_atoms as f32, 1.0_f32, &self.stream)?
                .affine(delta_z as f64, v_min as f64, &self.stream)?;
            let atoms_unsq = atoms_1d.unsqueeze(0, &self.stream)?.unsqueeze(0, &self.stream)?;
            let atoms_tensor = atoms_unsq.broadcast_as(&[batch_size, num_actions, num_atoms], &self.stream)?;

            // Compute expected Q-values: Q(s,a) = Σ(z_i * p_i)
            let prod = z_probs.broadcast_mul(&atoms_tensor, &self.stream)?;
            // Sum over atoms: [B, A, N] -> [B, A] via GPU matmul with ones vector
            let (b, a, n_at) = prod.dims3()?;
            // Reshape [B, A, N] -> [B*A, N]
            let prod_2d = prod.reshape(vec![b * a, n_at])?;
            // ones [N, 1]
            let ones = GpuTensor::full(&[n_at, 1], 1.0, &self.stream)?;
            // matmul: [B*A, N] x [N, 1] -> [B*A, 1]
            let cublas = cudarc::cublas::CudaBlas::new(self.stream.clone()).map_err(|e| {
                MLError::ModelError(format!("distributional cublas init: {e}"))
            })?;
            let summed_2d = prod_2d.matmul(&ones, &cublas, &self.stream)?;
            // Reshape [B*A, 1] -> [B, A]
            summed_2d.reshape(vec![b, a])?
        } else if let Some(ref dueling_net) = self.dueling_q_network {
            // Wave 11.4: Priority 2 - Dueling only
{
                // DuelingQNetwork takes &[f32], convert from GpuTensor (dead code path — always None)
                let host = state.to_host(&self.stream)?;
                let batch_size = state.dim(0)?;
                let result = dueling_net.forward(&host, batch_size)?;
                GpuTensor::from_host(&result, vec![batch_size, self.config.num_actions], &self.stream)?
            }
        } else {
            return Err(MLError::ModelError(
                "No Q-network available: branching and distributional-dueling networks are both None".to_owned(),
            ));
        };

        // BUG #19 FIX (Wave 16S-V18): Remove clamp - has zero gradient at boundaries
        // Gradient clipping (max_norm=10.0) + Huber loss (delta=10.0) already prevent explosions
        // Clamp causes ∂clamp/∂x = 0 when Q-values hit ±1000 boundaries → instant gradient death
        // Self-regulation through:
        //   1. Gradient clipping (max_norm=10.0) prevents weight explosions
        //   2. Huber loss (delta=10.0) reduces sensitivity to outliers
        //   3. Adam optimizer with momentum provides natural stabilization

        // BUG #37 FIX: Add forward-pass Q-value clipping to prevent step-level explosions
        // Key difference from BUG #19 fix:
        //   - BUG #19: Removed clamp to fix GRADIENT death (backward pass issue)
        //   - BUG #37: Add clamp to fix Q-VALUE explosions (forward pass issue)
        // This is safe because:
        //   1. Clipping happens AFTER forward pass (inference)
        //   2. Gradients flow through during backward (no gradient death)
        //   3. Prevents extreme Q-values from propagating through action selection
        let q_values = {
            let clamped = q_values.clamp(
                self.config.q_value_clip_min,
                self.config.q_value_clip_max,
                &self.stream,
            )?;

            // GPU-only clip monitoring (every 1000 steps, zero full-tensor download)
            if !self.training_forward_active && self.training_steps % 1000 == 0 {
                let n = q_values.numel();
                if n > 0 {
                    if let Ok(ew) = ml_core::cuda_autograd::elementwise::get_or_compile(&self.stream) {
                        if let Ok(reductions) = ml_core::cuda_autograd::ReductionKernels::new(&self.stream) {
                            // GPU: diff = |q - clamped|, threshold at 0.01
                            if let Ok(diff) = ew.binary(q_values.data(), clamped.data(), n, 1) { // sub
                                if let Ok(abs_diff) = ew.abs(&diff, n) {
                                    // Count elements > 0.01: le(0.01) gives 1.0 where <=0.01, so
                                    // NOT-le = n - sum(le) gives count where >0.01
                                    if let Ok(le_mask) = ew.le(&abs_diff, 0.01, n) {
                                        if let Ok(le_count) = reductions.sum(&le_mask, n) {
                                            let clipped_count = n as f32 - le_count;
                                            if clipped_count > 0.5 {
                                                // Only fetch stats when clipping detected
                                                // ALLOWED: stats() reads 5 scalars that exit system
                                                if let Ok(stats) = reductions.stats(q_values.data(), n) {
                                                    tracing::warn!(
                                                        "BUG #37: Clipped {}/{} Q-values at step {} (range before: [{:.2}, {:.2}], bounds: [{:.2}, {:.2}])",
                                                        clipped_count as u32, n, self.training_steps,
                                                        stats.min, stats.max,
                                                        self.config.q_value_clip_min, self.config.q_value_clip_max
                                                    );
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }

            clamped
        };

        // Cast output back to F32 for API compatibility (callers expect f32 tensors)
        let q_values = q_values.to_dtype(NativeDType::F32, &self.stream)?;

        Ok(q_values)
    }

    /// Select action using epsilon-greedy policy
    #[allow(clippy::cognitive_complexity)]
    pub fn select_action(&mut self, state: &[f32]) -> Result<FactoredAction, MLError> {
        // Reset noise for exploration (Noisy Networks / Rainbow DQN, always enabled)
        // Uses sigma schedule when annealing is active (noise_sigma_scale < 1.0)
        if let Some(ref mut branching_net) = self.branching_q_network {
            if (self.noise_sigma_scale - 1.0).abs() < 1e-9 {
                branching_net.reset_noise()?;
            } else {
                branching_net.reset_noise_with_sigma(self.noise_sigma_scale)?;
            }
        }

        // Increment total steps counter (tracks all environment steps including warmup)
        self.total_steps += 1;

        let mut rng = thread_rng();

        // During warmup period: always use random exploration (epsilon=1.0)
        let in_warmup = self.total_steps <= self.config.warmup_steps as u64;

        // Noisy nets use noisy_epsilon_floor as a minimum random exploration rate
        // instead of the full epsilon schedule. This prevents triple-stacking
        // (noisy + epsilon + count) while ensuring some random actions feed
        // diverse experiences into the replay buffer.
        let effective_epsilon = self.get_effective_epsilon();

        // Epsilon-greedy exploration (forced to random during warmup).
        // Per-branch independent epsilon: each branch flips its own coin.
        // P(all greedy) = (1-ε)^3 ≈ 72.9% at ε=0.10
        // P(at least one random) ≈ 27.1% — much better factored coverage.
        let exp_random = in_warmup || rng.gen::<f32>() < effective_epsilon;
        let ord_random = in_warmup || rng.gen::<f32>() < effective_epsilon;
        let urg_random = in_warmup || rng.gen::<f32>() < effective_epsilon;

        let action = if exp_random && ord_random && urg_random {
            // All branches random — skip forward pass.
            // 4-branch: dir(3) × mag(3) → 7-level ExposureLevel, order(3), urgency(3).
            let dir = rng.gen_range(0..3_usize);
            let mag = rng.gen_range(0..3_usize);
            let exposure = ExposureLevel::from_dir_mag(dir, mag)?;
            let order = OrderType::from_index(rng.gen_range(0..3_usize))?;
            let urgency = Urgency::from_index(rng.gen_range(0..3_usize))?;
            FactoredAction { exposure, order, urgency }
        } else {
            // At least one branch is greedy — need forward pass.
            let state_tensor = GpuTensor::from_host(state, vec![state.len()], &self.stream)
                .and_then(|t| t.reshape(vec![1, ml_core::state_layout::STATE_DIM]))
                .map_err(|e| MLError::ModelError(format!("State tensor: {}", e)))?;
            let branching_net = self.branching_q_network.as_ref().ok_or_else(|| {
                MLError::ModelError("Branching enabled but network not initialized".into())
            })?;
            let output = branching_net.forward_branches_eval(&state_tensor)?;
            let greedy = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output, &self.stream)?;

            // 4-branch: greedy[0]=dir(0-2), greedy[1]=mag(0-2) → 7-level ExposureLevel
            let exposure = if exp_random {
                let dir = rng.gen_range(0..3_usize);
                let mag = rng.gen_range(0..3_usize);
                ExposureLevel::from_dir_mag(dir, mag)?
            } else {
                let dir = greedy.first().copied().unwrap_or(1) as usize;
                let mag = greedy.get(1).copied().unwrap_or(1) as usize;
                ExposureLevel::from_dir_mag(dir, mag)?
            };
            // 4-branch greedy: [dir, mag, order, urgency] → order at [2], urgency at [3]
            let order = if ord_random {
                OrderType::from_index(rng.gen_range(0..3_usize))?
            } else {
                OrderType::from_index(greedy.get(2).copied().unwrap_or(0) as usize)?
            };
            let urgency = if urg_random {
                Urgency::from_index(rng.gen_range(0..3_usize))?
            } else {
                Urgency::from_index(greedy.get(3).copied().unwrap_or(1) as usize)?
            };
            FactoredAction { exposure, order, urgency }
        };

        // Record action for UCB exploration tracking. The numeric
        // `count_bonus_coefficient` is the live kill-switch (0.0 disables
        // the bonus contribution at the action selector); recording itself
        // is unconditional and cheap.
        self.count_bonus.record_action(action.exposure as usize);

        // Track action for entropy penalty calculation
        self.track_action(action);

        // Log warmup progress every 10K steps
        if in_warmup && self.total_steps % 10000 == 0 {
            tracing::info!(
                "Warmup: {}/{}K steps ({:.1}% complete)",
                self.total_steps / 1000,
                self.config.warmup_steps / 1000,
                (self.total_steps as f64 / self.config.warmup_steps as f64) * 100.0
            );
        }

        // Log warmup completion
        if self.total_steps == self.config.warmup_steps as u64 {
            tracing::info!(
                "✓ Warmup complete - starting training ({}K steps collected)",
                self.config.warmup_steps / 1000
            );
        }

        Ok(action)
    }

    /// Select action with confidence derived from softmax over Q-values.
    ///
    /// Identical to `select_action` but also returns a confidence score in [0.5, 0.95].
    /// Confidence is the softmax probability of the selected action (temperature=1.0),
    /// clamped to [0.5, 0.95] to avoid overconfidence and underconfidence.
    ///
    /// For production inference (epsilon=0), this gives a meaningful measure of how
    /// much the model "prefers" the selected action over alternatives.
    pub fn select_action_with_confidence(
        &mut self,
        state: &[f32],
    ) -> Result<(FactoredAction, f32), MLError> {
        // Reset noise for exploration (Noisy Networks / Rainbow DQN, always enabled)
        // Uses sigma schedule when annealing is active (same as select_action)
        if let Some(ref mut branching_net) = self.branching_q_network {
            if (self.noise_sigma_scale - 1.0).abs() < 1e-9 {
                branching_net.reset_noise()?;
            } else {
                branching_net.reset_noise_with_sigma(self.noise_sigma_scale)?;
            }
        }

        // Increment total steps counter (tracks all environment steps including warmup)
        self.total_steps += 1;

        let mut rng = thread_rng();

        // During warmup period: always use random exploration (epsilon=1.0)
        let in_warmup = self.total_steps <= self.config.warmup_steps as u64;

        // Noisy nets + epsilon floor (same logic as select_action)
        let effective_epsilon = self.get_effective_epsilon();

        // Epsilon-greedy exploration (forced to random during warmup).
        // Per-branch independent epsilon: each branch flips its own coin.
        let exp_random = in_warmup || rng.gen::<f32>() < effective_epsilon;
        let ord_random = in_warmup || rng.gen::<f32>() < effective_epsilon;
        let urg_random = in_warmup || rng.gen::<f32>() < effective_epsilon;

        let (action, confidence) = if exp_random && ord_random && urg_random {
            // All branches random — skip forward pass.
            // 4-branch: exposure = dir*3 + mag (9 combos), order(3), urgency(3)
            let dir = rng.gen_range(0..3_usize);
            let mag = rng.gen_range(0..3_usize);
            let exposure = ExposureLevel::from_dir_mag(dir, mag)?;
            let order = OrderType::from_index(rng.gen_range(0..3_usize))?;
            let urgency = Urgency::from_index(rng.gen_range(0..3_usize))?;
            let uniform_conf = (1.0_f32 / 81.0).clamp(0.5, 0.95);
            (FactoredAction { exposure, order, urgency }, uniform_conf)
        } else {
            // At least one branch is greedy — need forward pass.
            let state_tensor = GpuTensor::from_host(state, vec![state.len()], &self.stream)
                .and_then(|t| t.reshape(vec![1, ml_core::state_layout::STATE_DIM]))
                .map_err(|e| MLError::ModelError(format!("State tensor: {}", e)))?;
            let branching_net = self.branching_q_network.as_ref().ok_or_else(|| {
                MLError::ModelError("Branching enabled but network not initialized".into())
            })?;
            let output = branching_net.forward_branches_eval(&state_tensor)?;
            let greedy = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output, &self.stream)?;

            // Per-branch confidence: average of per-branch softmax max probs (greedy branches only)
            let mut total_conf = 0.0_f32;
            let mut greedy_count = 0_u32;
            for (i, adv) in output.advantages.iter().enumerate() {
                let is_random = match i { 0 => exp_random, 1 => ord_random, _ => urg_random };
                if !is_random {
                    total_conf += self.softmax_confidence(adv)?;
                    greedy_count += 1;
                }
            }
            let conf = if greedy_count > 0 {
                (total_conf / greedy_count as f32).clamp(0.5, 0.95)
            } else {
                0.5
            };

            // 4-branch: greedy[0]=dir(0-2), greedy[1]=mag(0-2), composite = dir*3+mag
            let exposure = if exp_random {
                let dir = rng.gen_range(0..3_usize);
                let mag = rng.gen_range(0..3_usize);
                ExposureLevel::from_dir_mag(dir, mag)?
            } else {
                let dir = greedy.first().copied().unwrap_or(1) as usize;
                let mag = greedy.get(1).copied().unwrap_or(1) as usize;
                ExposureLevel::from_dir_mag(dir, mag)?
            };
            // 4-branch greedy: [dir, mag, order, urgency] -> order at [2], urgency at [3]
            let order = if ord_random {
                OrderType::from_index(rng.gen_range(0..3_usize))?
            } else {
                OrderType::from_index(greedy.get(2).copied().unwrap_or(0) as usize)?
            };
            let urgency = if urg_random {
                Urgency::from_index(rng.gen_range(0..3_usize))?
            } else {
                Urgency::from_index(greedy.get(3).copied().unwrap_or(1) as usize)?
            };
            (FactoredAction { exposure, order, urgency }, conf)
        };

        // Record action for UCB exploration tracking. The numeric
        // `count_bonus_coefficient` is the live kill-switch.
        self.count_bonus.record_action(action.exposure as usize);

        // Track action for entropy penalty calculation
        self.track_action(action);

        Ok((action, confidence))
    }

    /// Compute softmax confidence for the selected (argmax) action.
    ///
    /// Given Q-values of shape `[1, num_actions]`, computes softmax probabilities
    /// and returns the probability of the highest-scoring action, clamped to [0.5, 0.95].
    fn softmax_confidence(&self, q_values: &GpuTensor) -> Result<f32, MLError> {
        // Inference path: downloads num_actions floats (typically 5 values = 20 bytes)
        // for softmax confidence computation. This exits the system as a single
        // confidence scalar. Acceptable: num_actions is O(1), not O(batch_size).
        // Uses self.stream — no CUDA context/stream creation per call.
        let host = q_values.to_host(&self.stream)?;
        // Numerically stable softmax on host
        let max_val = host.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        let exps: Vec<f32> = host.iter().map(|&v| (v - max_val).exp()).collect();
        let sum_exp: f32 = exps.iter().sum();
        let max_prob = if sum_exp > 0.0 {
            exps.iter().copied().fold(0.0_f32, f32::max) / sum_exp
        } else {
            1.0 / host.len() as f32
        };

        Ok(max_prob.clamp(0.5, 0.95))
    }

    /// Pure greedy action selection with confidence for production inference -- takes `&self`.
    ///
    /// Unlike [`select_action`] / [`select_action_with_confidence`] this method
    /// performs NO mutations:
    /// - No epsilon-greedy exploration (always greedy)
    /// - No noisy-net noise reset
    /// - No step counter increment
    /// - No count-bonus tracking
    /// - No action entropy tracking
    ///
    /// This allows the caller to hold a **read lock** (`RwLock::read()`) instead
    /// of a write lock, enabling concurrent predictions from multiple tokio tasks.
    ///
    /// Returns `(action, confidence)` where confidence is the softmax probability
    /// of the selected action over Q-values, clamped to \[0.5, 0.95\].
    ///
    /// Callers MUST ensure exploration is disabled in the config (epsilon=0,
    /// noisy nets disabled via `disable_noise()`, `count_bonus_coefficient=0.0`, `warmup_steps=0`).
    /// The production `DQNModel` wrapper enforces this at construction time.
    pub fn select_action_inference(&self, state: &[f32]) -> Result<(FactoredAction, f32), MLError> {
        let state_tensor = GpuTensor::from_host(state, vec![state.len()], &self.stream)
            .and_then(|t| t.reshape(vec![1, ml_core::state_layout::STATE_DIM]))
            .map_err(|e| MLError::ModelError(format!("State tensor: {}", e)))?;

        // Branching DQN: greedy argmax per branch with per-branch confidence.
        let branching_net = self.branching_q_network.as_ref().ok_or_else(|| {
            MLError::ModelError("Branching enabled but network not initialized".into())
        })?;
        let output = branching_net.forward_branches_eval(&state_tensor)?;

        let branch_actions = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output, &self.stream)?;
        // 4-branch: [dir, mag, order, urgency] → exposure = dir*3 + mag
        let dir = branch_actions.first().copied().unwrap_or(1) as usize;
        let mag = branch_actions.get(1).copied().unwrap_or(1) as usize;
        let exposure = ExposureLevel::from_dir_mag(dir, mag)?;
        let order = OrderType::from_index(branch_actions.get(2).copied().unwrap_or(0) as usize)?;
        let urgency = Urgency::from_index(branch_actions.get(3).copied().unwrap_or(1) as usize)?;

        // Per-branch confidence: average of per-branch softmax max probs.
        // (Risk-sensitive CVaR confidence on the live C51 distribution lands
        // separately — see `cvar_alpha` / commit B follow-up.)
        let mut total_conf = 0.0_f32;
        for adv in &output.advantages {
            total_conf += self.softmax_confidence(adv)?;
        }
        let conf = (total_conf / output.advantages.len() as f32).clamp(0.5, 0.95);

        Ok((FactoredAction { exposure, order, urgency }, conf))
    }

    /// Compute Q-values for a batch of states via the branching network's exposure
    /// head. Production training trains only the branching head; distributional
    /// IQN runs inside the cuda_pipeline dual head and is not consumed here.
    ///
    /// Returns a `[N, num_actions]` tensor of branching exposure-head Q-values.
    pub fn q_values_for_batch(&self, states: &GpuTensor) -> Result<GpuTensor, MLError> {
        let branching_net = self.branching_q_network.as_ref().ok_or_else(|| {
            MLError::ModelError("Branching enabled but network not initialized".into())
        })?;
        // Ensure input is on the correct device and dtype (matches forward() contract)
        let states = states.to_dtype(NativeDType::F32, &self.stream)
            .map_err(|e| MLError::ModelError(e.to_string()))?;
        let states = states.gpu_clone(&self.stream)
            .map_err(|e| MLError::ModelError(format!("device migration: {e}")))?;
        let output = branching_net.forward_branches_eval(&states)?;
        // advantages[0] = exposure branch Q-values [batch, 5]
        output.advantages.into_iter().next().ok_or_else(|| {
            MLError::ModelError("BranchOutput has no advantage branches".into())
        })
    }

    /// Batch greedy action selection — single GPU forward pass for N states.
    ///
    /// Runs a batched forward pass through the Q-network and returns the greedy
    /// (argmax) action index for each state. No exploration, no side effects.
    ///
    /// # Arguments
    /// * `states` - `GpuTensor` of shape `[N, state_dim]`
    ///
    /// # Returns
    /// `Vec<usize>` of length N containing greedy action indices.
    ///
    /// # Performance
    /// Reduces N individual GPU kernel launches to a single batched forward pass.
    /// With `EVAL_CHUNK_SIZE=1024`, this gives ~1000× fewer kernel launches vs per-bar inference.
    pub fn batch_greedy_actions(&self, states: &GpuTensor) -> Result<GpuTensor, MLError> {
        let q_values = self.q_values_for_batch(states)?;
        let indices = q_values.argmax(1, &self.stream)
            .map_err(|e| MLError::ModelError(format!("Batch argmax failed: {}", e)))?;
        let floats: Vec<f32> = indices.iter().map(|&i| i as f32).collect();
        GpuTensor::from_host(&floats, vec![indices.len()], &self.stream)
    }

    /// Batch softmax (Boltzmann) action selection via Gumbel-max trick — fully GPU-resident.
    ///
    /// Equivalent to sampling from softmax(Q/temperature) per row, but avoids
    /// CPU transfer of Q-values. Uses: argmax(Q/T - log(-log(U))) where U~Uniform(0,1).
    /// At low temperature (~0.1), nearly greedy; provides diversity when Q-values are uniform.
    ///
    /// # Arguments
    /// * `states` - `GpuTensor` of shape `[N, state_dim]`
    /// * `temperature` - Boltzmann temperature (clamped to >= 1e-6)
    pub fn batch_softmax_actions(
        &self,
        states: &GpuTensor,
        temperature: f64,
    ) -> Result<GpuTensor, MLError> {
        let q_values = self.q_values_for_batch(states)?;
        let temp = temperature.max(1e-6) as f32;

        // GPU-native Gumbel-max: Q/T + Gumbel noise on GPU, then GPU argmax per row
        let n = q_values.dim(0)?;
        let num_actions = q_values.dim(1)?;
        let total = n * num_actions;

        // Scale Q-values by 1/temperature on GPU
        let inv_temp = 1.0 / temp;
        let scaled = q_values.affine(inv_temp as f64, 0.0, &self.stream)?;

        // Generate Gumbel noise: -log(-log(U)) where U ~ Uniform(0,1)
        // Use randn as source of randomness, then transform
        let mut rng = thread_rng();
        let mut gumbel_host = Vec::with_capacity(total);
        for _ in 0..total {
            let u: f32 = rng.gen::<f32>().max(1e-10).min(1.0 - 1e-7);
            gumbel_host.push(-((-u.ln()).ln()));
        }
        let gumbel_gpu = GpuTensor::from_host(&gumbel_host, vec![n, num_actions], &self.stream)?;

        // GPU: perturbed = scaled + gumbel
        let perturbed = scaled.add(&gumbel_gpu, &self.stream)?;

        // GPU: argmax per row -> [n] indices
        // ALLOWED: argmax_rows returns Vec<u32> (one u32 per sample that exits system as action)
        let reductions = ml_core::cuda_autograd::ReductionKernels::new(&self.stream)?;
        let indices = reductions.argmax_rows(perturbed.data(), n, num_actions)?;

        // Convert indices to f32 GpuTensor
        let result_f32: Vec<f32> = indices.iter().map(|&i| i as f32).collect();
        GpuTensor::from_host(&result_f32, vec![n], &self.stream)
    }

    /// Gumbel-softmax over 5 exposure-level Q-values, fully GPU-resident.
    ///
    /// DQN outputs [N, 5] Q-values (one per exposure level). Gumbel-max samples
    /// stochastically from the Q-value distribution with temperature control.
    /// Only transfers N u32 indices (exposure levels 0-4) to CPU.
    pub fn batch_hierarchical_softmax_actions(
        &self,
        states: &GpuTensor,
        temperature: f64,
    ) -> Result<GpuTensor, MLError> {
        // Reuse the general Gumbel-max implementation
        self.batch_softmax_actions(states, temperature)
    }

    /// Track action for entropy penalty calculation
    ///
    /// This should be called after action selection (both single and batch modes)
    /// to maintain the sliding window used for diversity penalty.
    pub fn track_action(&mut self, action: FactoredAction) {
        self.recent_actions.push_back(action);
        if self.recent_actions.len() > 100 {
            self.recent_actions.pop_front();
        }
    }

    /// Track multiple actions for entropy penalty calculation (batch version)
    pub fn track_actions_batch(&mut self, actions: &[FactoredAction]) {
        for &action in actions {
            self.track_action(action);
        }
    }

    /// Update the optimizer's learning rate in-place: `new_lr = current_lr * decay_factor`.
    ///
    /// Uses `set_learning_rate()` to preserve Adam momentum/variance state
    /// (first/second moment estimates). Previously this recreated the entire
    /// optimizer, discarding all accumulated m/v buffers — causing the optimizer
    /// to "cold start" after every LR decay step.
    ///
    /// Called by the LR scheduler at epoch boundaries. If the optimizer has not
    /// been initialised yet (no `train_step` calls), this is a no-op.
    pub fn update_learning_rate(&mut self, decay_factor: f64) -> Result<(), MLError> {
        if let Some(ref mut optimizer) = self.optimizer {
            let current_lr = optimizer.learning_rate() as f64;
            let new_lr = current_lr * decay_factor;
            optimizer.set_learning_rate(new_lr as f32);
        }
        Ok(())
    }

    /// Log Q-values for the first state in batch (Wave 10-A4 diagnostic monitoring)
    ///
    /// WAVE 23 P0 Fix #4: Now implements early stopping on Q-value divergence
    #[allow(clippy::cognitive_complexity)]
    pub fn log_q_values(&mut self, states_tensor: &GpuTensor) -> Result<(), MLError> {
        // Get Q-values for first state in batch
        let first_state = states_tensor.narrow(0, 0, 1, &self.stream)?;
        let q_values = self.forward(&first_state)?;
        let n_actions = q_values.dim(1).unwrap_or(0);

        // GPU-only stats computation via fused reduction kernel (zero full-tensor download)
        let n = q_values.numel();
        if n == 0 {
            return self.log_q_values_from_stats(0.0, 0.0, 0.0, 0.0, n_actions);
        }
        // ALLOWED: stats() returns 5 scalars (20 bytes) that exit the system for monitoring
        let reductions = ml_core::cuda_autograd::ReductionKernels::new(&self.stream)?;
        let stats = reductions.stats(q_values.data(), n)?;

        // Delegate to stats-based variant (shared divergence/collapse logic)
        self.log_q_values_from_stats(stats.min, stats.max, stats.mean, stats.variance, n_actions)
    }

    /// Log Q-values from pre-computed GPU statistics (zero tensor readback).
    ///
    /// Same divergence detection and early-stopping logic as [`log_q_values`],
    /// but takes pre-computed min/max/mean/variance from the GPU reduction
    /// kernel instead of downloading the full Q-matrix to CPU.
    #[allow(clippy::cognitive_complexity)]
    pub fn log_q_values_from_stats(
        &mut self,
        q_min: f32,
        q_max: f32,
        q_mean: f32,
        q_variance: f32,
        n_actions: usize,
    ) -> Result<(), MLError> {
        // Push to Prometheus gauges (cheap atomic set, scrapable every 15s)
        training_metrics::set_training_step("dqn", "current", self.training_steps as f64);
        training_metrics::set_q_value_stats("dqn", "current", q_mean as f64, q_max as f64);

        tracing::debug!(
            "Step {} Q-values ({} actions): range=[{:.2}, {:.2}], mean={:.2}",
            self.training_steps, n_actions, q_min, q_max, q_mean
        );

        // Q-value collapse detection
        if q_mean.abs() < 0.0001 && q_variance < 0.0001 {
            tracing::warn!(
                "Q-VALUE COLLAPSE DETECTED at step {}: mean={:.6}, variance={:.6}, range=[{:.6}, {:.6}]",
                self.training_steps, q_mean, q_variance, q_min, q_max
            );
        }

        // Divergence detection (same logic as log_q_values)
        const Q_VALUE_DIVERGENCE_THRESHOLD: f32 = 10000.0;
        const Q_VALUE_DIVERGENCE_PATIENCE: usize = 3;

        if q_max.abs() > Q_VALUE_DIVERGENCE_THRESHOLD
            || q_min.abs() > Q_VALUE_DIVERGENCE_THRESHOLD
        {
            self.q_value_divergence_counter += 1;
            tracing::warn!(
                "Q-VALUE DIVERGENCE: range=[{:.2}, {:.2}] exceeds threshold ±{:.0} at step {} (consecutive: {}/{})",
                q_min, q_max, Q_VALUE_DIVERGENCE_THRESHOLD,
                self.training_steps,
                self.q_value_divergence_counter, Q_VALUE_DIVERGENCE_PATIENCE
            );
            if self.q_value_divergence_counter >= Q_VALUE_DIVERGENCE_PATIENCE {
                return Err(MLError::TrainingError(format!(
                    "Training stopped: Q-value divergence for {} consecutive checks (threshold: ±{:.0}, range: [{:.2}, {:.2}]). \
                     Consider reducing learning rate ({:.2e}), increasing gradient clipping ({:.2}), or reducing Huber delta ({:.2})",
                    self.q_value_divergence_counter,
                    Q_VALUE_DIVERGENCE_THRESHOLD,
                    q_min,
                    q_max,
                    self.config.learning_rate,
                    self.config.gradient_clip_norm,
                    self.config.huber_delta
                )));
            }
        } else if self.q_value_divergence_counter > 0 {
            tracing::info!(
                "Q-values normalized: range=[{:.2}, {:.2}] within threshold ±{:.0}, resetting counter (was: {})",
                q_min, q_max, Q_VALUE_DIVERGENCE_THRESHOLD,
                self.q_value_divergence_counter
            );
            self.q_value_divergence_counter = 0;
        } else {
            // Q-values within normal range, no action needed
        }

        Ok(())
    }

    /// Detect dead neurons and log comprehensive diagnostics (Wave 10-A4)
    ///
    /// WAVE 23 P0 Fix #2: Now implements early stopping on consecutive gradient collapse
    pub fn log_diagnostics(&mut self, grad_norm: f32) -> Result<(), MLError> {
        let dead_pct = self.detect_dead_neurons()?;

        // Push to Prometheus gauges (cheap atomic set, scrapable every 15s)
        training_metrics::set_gradient_norm("dqn", "current", grad_norm as f64);

        tracing::debug!(
            "Step {} Diagnostics: grad_norm={:.2}, dead_neurons={:.2}%",
            self.training_steps,
            grad_norm,
            dead_pct
        );

        // WAVE 24 BUG #43 FIX: Add warmup period for gradient collapse check
        // Problem: Threshold too strict for cold-start training (buffer filling phase)
        // During cold start: replay buffer filling, random Q-network weights, minimal TD errors
        // Expected gradient norms: 0.0001-0.001 (10-100x smaller than normal training)
        // Warmup: Skip check for first 20% of buffer size (prevents false positives)
        let warmup_steps = (self.config.replay_buffer_capacity as f64 * 0.2) as u64;
        let past_warmup = self.training_steps > warmup_steps;

        // WAVE 23 P0 Fix #1: Adaptive gradient collapse threshold (learning-rate aware)
        // Threshold = LR × gradient_collapse_multiplier (default: 100.0)
        // Old threshold (0.1): 1000+ false alarms per trial
        // New threshold (100.0): Scales appropriately with learning rate
        let threshold = self.config.learning_rate as f32 * self.config.gradient_collapse_multiplier as f32;
        if past_warmup && grad_norm < threshold {
            self.gradient_collapse_counter += 1;
            tracing::warn!(
                "⚠️  GRADIENT COLLAPSE: norm={:.6} (threshold: {:.6}) at step {} (consecutive: {}/{})",
                grad_norm,
                threshold,
                self.training_steps,
                self.gradient_collapse_counter,
                self.config.gradient_collapse_patience
            );

            // WAVE 23 P0 Fix #2: Stop training after patience exhausted
            if self.gradient_collapse_counter >= self.config.gradient_collapse_patience {
                return Err(MLError::TrainingError(format!(
                    "Training stopped: Gradient collapse detected for {} consecutive epochs (threshold: {:.6}, actual: {:.6}). \
                     This indicates the model has stopped learning. Consider:\n\
                     1. Increasing learning rate (current: {:.2e})\n\
                     2. Reducing batch size (current: {})\n\
                     3. Checking for data quality issues\n\
                     4. Adjusting network architecture",
                    self.gradient_collapse_counter,
                    threshold,
                    grad_norm,
                    self.config.learning_rate,
                    self.config.batch_size
                )));
            }
        } else {
            // Reset counter when gradients recover
            if self.gradient_collapse_counter > 0 {
                tracing::info!(
                    "✓ Gradients recovered: norm={:.6} (threshold: {:.6}), resetting collapse counter (was: {})",
                    grad_norm,
                    threshold,
                    self.gradient_collapse_counter
                );
                self.gradient_collapse_counter = 0;
            }
        }

        Ok(())
    }

    /// Suppress or enable `forward()` monitoring (GPU→CPU sync avoidance).
    /// Set to `true` before `forward()` calls in the training hot path, `false` after.
    #[allow(clippy::missing_const_for_fn)] // intentionally non-const for future mutation logic
    pub fn set_training_forward_active(&mut self, active: bool) {
        self.training_forward_active = active;
    }

    /// Gradient collapse check WITHOUT dead neuron detection (zero GPU->CPU sync).
    /// Use this in the CUDA guard hot path where `detect_dead_neurons()` would be too expensive.
    /// Dead neuron detection should be called at epoch boundary via `log_diagnostics()`.
    pub fn check_gradient_collapse(&mut self, grad_norm: f32) -> Result<(), MLError> {
        // Push to Prometheus gauges (cheap atomic set)
        training_metrics::set_gradient_norm("dqn", "current", grad_norm as f64);

        tracing::debug!(
            "Step {} guard: grad_norm={:.2}",
            self.training_steps,
            grad_norm,
        );

        // Use collapse_warmup_capacity (original buffer_size before AutoReplaySizer)
        // to prevent unreachable warmup when AutoReplaySizer scales buffer to 10M+
        let warmup_steps = (self.config.collapse_warmup_capacity as f64 * 0.2) as u64;
        let past_warmup = self.training_steps > warmup_steps;

        let threshold =
            self.config.learning_rate as f32 * self.config.gradient_collapse_multiplier as f32;
        if past_warmup && grad_norm < threshold {
            self.gradient_collapse_counter += 1;
            tracing::warn!(
                "GRADIENT COLLAPSE: norm={:.6} (threshold: {:.6}) at step {} (consecutive: {}/{})",
                grad_norm,
                threshold,
                self.training_steps,
                self.gradient_collapse_counter,
                self.config.gradient_collapse_patience
            );

            if self.gradient_collapse_counter >= self.config.gradient_collapse_patience {
                return Err(MLError::TrainingError(format!(
                    "Training stopped: Gradient collapse detected for {} consecutive epochs (threshold: {:.6}, actual: {:.6}). \
                     Consider: 1. Increasing learning rate (current: {:.2e}) 2. Reducing batch size (current: {}) \
                     3. Checking data quality 4. Adjusting network architecture",
                    self.gradient_collapse_counter,
                    threshold,
                    grad_norm,
                    self.config.learning_rate,
                    self.config.batch_size
                )));
            }
        } else if self.gradient_collapse_counter > 0 {
            tracing::info!(
                "Gradients recovered: norm={:.6} (threshold: {:.6}), resetting collapse counter (was: {})",
                grad_norm,
                threshold,
                self.gradient_collapse_counter
            );
            self.gradient_collapse_counter = 0;
        } else {
            // Gradient norm within normal range, no action needed
        }

        Ok(())
    }

    /// Detect dead `ReLU` neurons (weights stuck at zero) -- GPU-only.
    ///
    /// Uses GPU elementwise `le()` + reduction sum to count near-zero weights
    /// without downloading any parameter tensors to CPU.
    fn detect_dead_neurons(&self) -> Result<f32, MLError> {
        // Dead neuron detection uses the branching network's all_trainable_vars.
        // For the fused CUDA path, this is a diagnostic-only cold path.
        let vars = if let Some(ref br) = self.branching_q_network {
            br.all_trainable_vars()
        } else {
            return Ok(0.0); // No trainable vars accessible without VarStore
        };

        let ew = ml_core::cuda_autograd::elementwise::get_or_compile(&self.stream)?;
        let reductions = ml_core::cuda_autograd::ReductionKernels::new(&self.stream)?;

        let mut total_count: usize = 0;
        let mut dead_count_f32 = 0.0_f32;

        for param in &vars {
            let n = param.len();
            if n == 0 { continue; }
            let abs_vals = ew.abs(param, n)?;
            let near_zero = ew.le(&abs_vals, 1e-6, n)?;
            // ALLOWED: single scalar readback (exits system for diagnostic metric)
            let count = reductions.sum(&near_zero, n)?;
            dead_count_f32 += count;
            total_count += n;
        }

        if total_count == 0 {
            return Ok(0.0);
        }
        Ok((dead_count_f32 / total_count as f32) * 100.0)
    }

    /// Update exploration epsilon (called once per epoch by trainer)
    pub fn update_epsilon(&mut self) {
        self.epsilon = (self.epsilon * self.config.epsilon_decay).max(self.config.epsilon_end);
    }

    /// Set noise sigma scale for annealed exploration.
    ///
    /// Called by the trainer when `NoisySigmaScheduler` is active.
    /// The scale is applied during `select_action` noise resampling — NOT during
    /// training forward passes (where σ parameters are learned via gradients).
    pub const fn set_noise_sigma_scale(&mut self, scale: f64) {
        self.noise_sigma_scale = scale;
    }

    /// Reset count-based exploration bonus (call at epoch boundary)
    pub fn reset_count_bonus(&mut self) {
        self.count_bonus.reset();
    }

    /// Reset state that should not persist across walk-forward fold
    /// boundaries. Pairs with `DQNTrainer::reset_for_fold`.
    ///
    /// What's reset (and why):
    /// - `gradient_collapse_counter` — counts consecutive low-grad steps;
    ///   carrying across folds caused fold 1 to trip patience after a
    ///   single collapse-tier step (counter starts already near patience).
    /// - `training_steps` — the warmup gate `past_warmup = training_steps
    ///   > warmup_steps` becomes always-true after fold 0 completes,
    ///   removing the per-fold warmup grace period for data distribution
    ///   shifts at fold boundaries. Resetting gives fold 1+ a clean
    ///   warmup window before the gradient-collapse check fires.
    ///
    /// What's NOT reset (model + agent state):
    /// - Model parameters (the whole point of walk-forward; fold N+1
    ///   bootstraps from fold N's learned weights)
    /// - count_bonus state (Plan C T2 amendment removed direct dependence
    ///   on UCB count bonus; the count_bonus persistence is handled by
    ///   `reset_count_bonus` called explicitly elsewhere if needed)
    /// - target_network (handled by `reset_target_network` if explicitly
    ///   wanted; per project_fold_boundary_q_drift_resolved.md the IQN
    ///   sync_target_from_online is the principled mechanism)
    pub fn reset_for_fold(&mut self) {
        self.gradient_collapse_counter = 0;
        self.training_steps = 0;
    }

    /// Get count bonuses for all actions (UCB exploration).
    /// Returns a Vec<f32> of length `num_actions` with exploration bonuses.
    pub fn get_count_bonuses(&self) -> Vec<f32> {
        let bonuses = self.count_bonus.bonuses();
        bonuses.iter().map(|b| *b as f32).collect()
    }

    /// Get per-branch count bonuses for branching DQN.
    ///
    /// Returns `(direction[4], order[3], urgency[3])` UCB bonuses as fixed-size
    /// f32 arrays — stack-only, allocation-free. Caller is the production GPU
    /// action selector (see `trainer/action.rs`); the GPU kernel signature
    /// pins the shapes to 4 / 3 / 3.
    pub fn get_count_bonuses_branched(&self) -> ([f32; 4], [f32; 3], [f32; 3]) {
        let mut exp = [0.0f32; 4];
        let mut ord = [0.0f32; 3];
        let mut urg = [0.0f32; 3];
        self.count_bonus.bonuses_branched_f32(&mut exp, &mut ord, &mut urg);
        (exp, ord, urg)
    }

    /// Per-branch Q-values for the GPU action selector.
    ///
    /// Returns `Some((exposure, order, urgency))` as raw `CudaSlice<f32>` buffers
    /// when the branching network is initialised; `None` otherwise. Each buffer is
    /// `[batch_size, branch_size]` flattened (exposure = 4 actions, order = 3,
    /// urgency = 3).
    pub fn batch_branching_q_values(
        &self,
        states: &GpuTensor,
    ) -> Result<Option<(cudarc::driver::CudaSlice<f32>, cudarc::driver::CudaSlice<f32>, cudarc::driver::CudaSlice<f32>)>, MLError> {
        let Some(branching_net) = self.branching_q_network.as_ref() else {
            return Ok(None);
        };

        let n = states.shape().first().copied().unwrap_or(0);
        if n == 0 {
            return Err(MLError::ModelError("Empty batch for batch_branching_q_values".into()));
        }

        let output = branching_net.forward_branches_eval(states)?;
        if output.advantages.len() < 3 {
            return Err(MLError::ModelError(format!(
                "batch_branching_q_values: branching network produced {} branches, expected 3",
                output.advantages.len()
            )));
        }

        let exp_slice = output.advantages.first()
            .ok_or_else(|| MLError::ModelError("missing branch 0 (exposure)".into()))?
            .clone()
            .into_parts()
            .0;
        let ord_slice = output.advantages.get(1)
            .ok_or_else(|| MLError::ModelError("missing branch 1 (order)".into()))?
            .clone()
            .into_parts()
            .0;
        let urg_slice = output.advantages.get(2)
            .ok_or_else(|| MLError::ModelError("missing branch 2 (urgency)".into()))?
            .clone()
            .into_parts()
            .0;

        Ok(Some((exp_slice, ord_slice, urg_slice)))
    }

    /// Task 0.6 — NoisyNets σ mean for a specific action branch
    /// (0 = direction, 1 = magnitude, 2 = order, 3 = urgency). Mean is over
    /// the branch's fc + out NoisyLinear layers' weight_sigma + bias_sigma
    /// tensors. Returns 0.0 if branching network not initialized or branch
    /// index out of range. Used by HEALTH_DIAG H7 detection signal.
    pub fn branch_noisy_sigma_mean(&self, branch_idx: usize) -> f32 {
        if let Some(ref bn) = self.branching_q_network {
            bn.branch_noisy_sigma_mean(branch_idx).unwrap_or(0.0)
        } else {
            0.0
        }
    }

    /// Update target network by copying weights from main network.
    /// Cold path only — hot path uses fused CUDA EMA on flat buffers.
    /// DistributionalDueling and IQN have no per-layer copy (fused trainer only).
    fn update_target_network(&mut self) -> Result<(), MLError> {
        if let (Some(ref branching_net), Some(ref mut branching_target)) =
            (&self.branching_q_network, &mut self.branching_target_network)
        {
            branching_target.copy_weights_from(branching_net)?;
        }
        Ok(())
    }

    /// B1 FIX: Set evaluation mode.
    ///
    /// In eval mode, noisy layers use mean weights only (noise disabled),
    /// and the model is ready for deterministic or softmax-based action selection.
    /// Call with `true` before inference, `false` to re-enable training noise.
    pub fn set_eval_mode(&mut self, eval_mode: bool) -> Result<(), MLError> {
        if eval_mode {
            if let Some(ref mut branching_net) = self.branching_q_network {
                branching_net.disable_noise()?;
            }
        }
        // When exiting eval mode, next reset_noise() call will re-enable exploration
        Ok(())
    }

    /// BUG #38 FIX: Clear replay buffer
    pub fn clear_replay_buffer(&mut self) -> Result<(), MLError> {
        self.memory.clear()
    }

    /// BUG #38 FIX: Reset target network to match current network
    pub fn reset_target_network(&mut self) -> Result<(), MLError> {
        self.update_target_network()
    }

    /// Get current epsilon value
    pub const fn get_epsilon(&self) -> f32 {
        self.epsilon
    }

    /// Get the effective epsilon for exploration.
    ///
    /// With noisy nets always enabled, maintains a minimum epsilon floor to prevent
    /// action collapse. `NoisyNets` provide learned exploration but their noise
    /// magnitude can become smaller than the Q-value gap between actions,
    /// causing the agent to converge to a single action (e.g., Short100 only).
    /// The floor ensures at least 10% random actions (2% per exposure level).
    pub const fn get_effective_epsilon(&self) -> f32 {
        // Noisy nets always on: use epsilon floor to prevent action collapse
        // const fn cannot call max(), so use a manual branch
        if self.config.noisy_epsilon_floor > 0.10 {
            self.config.noisy_epsilon_floor
        } else {
            0.10
        }
    }

    /// Set epsilon value (used for deterministic evaluation)
    pub const fn set_epsilon(&mut self, epsilon: f64) {
        self.epsilon = epsilon.clamp(0.0, 1.0) as f32;
    }

    /// Get training steps (gradient updates only, excludes warmup)
    pub const fn get_training_steps(&self) -> u64 {
        self.training_steps
    }

    /// Increment training step counter by 1 (for fused CUDA training path).
    pub const fn increment_training_steps(&mut self) {
        self.training_steps += 1;
    }

    /// Get total environment steps (includes warmup period)
    pub const fn get_total_steps(&self) -> u64 {
        self.total_steps
    }

    /// Check if in warmup period
    pub const fn is_in_warmup(&self) -> bool {
        self.total_steps < self.config.warmup_steps as u64
    }

    /// Get warmup steps configured
    pub const fn get_warmup_steps(&self) -> usize {
        self.config.warmup_steps
    }

    /// Get replay buffer size
    pub fn get_replay_buffer_size(&self) -> Result<usize, MLError> {
        Ok(self.memory.len())
    }


    /// Load model weights from a safetensors checkpoint into the branching Q-network.
    ///
    /// Parses the safetensors file at `path`, uploads each named f32 tensor back
    /// to GPU as a `CudaSlice<f32>`, and calls
    /// [`BranchingDuelingQNetwork::load_from_named_slices`] to restore the live
    /// weights in place via D2D memcpy. The target network is also resynced so
    /// inference and TD targets see the same weights.
    ///
    /// # Checkpoint layout
    ///
    /// The training-time saver
    /// ([`serialize_model`](../../ml/src/trainers/dqn/trainer/mod.rs)) emits one
    /// merged safetensors file with three regime-prefixed copies of the branching
    /// weight manifest: `trending__`, `ranging__`, `volatile__`. The
    /// `DQN` owns one branching network. Checkpoints saved by the Phase 4+
    /// serializer use plain names (no prefix). For legacy checkpoints saved
    /// with a `trending__` prefix (pre-Phase 4), the loader falls back to
    /// plain names after stripping the prefix.
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the safetensors file (with or without `.safetensors` suffix)
    ///
    /// # Errors
    ///
    /// * [`MLError::CheckpointError`] — file missing, parse failure, architecture
    ///   drift (extra/missing tensors, shape mismatch).
    /// * [`MLError::ModelError`] — GPU upload failure.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use ml_dqn::{DQN, DQNConfig};
    ///
    /// let config = DQNConfig::emergency_safe_defaults();
    /// let mut dqn = DQN::new(config)?;
    /// dqn.load_from_safetensors("/path/to/checkpoint")?;
    /// # Ok::<(), ml_core::MLError>(())
    /// ```
    pub fn load_from_safetensors(&mut self, path: &str) -> Result<(), MLError> {
        let safetensors_path = if path.ends_with(".safetensors") {
            path.to_owned()
        } else {
            format!("{path}.safetensors")
        };

        let path_ref = std::path::Path::new(&safetensors_path);
        if !path_ref.exists() {
            return Err(MLError::CheckpointError(format!(
                "Checkpoint file not found: {safetensors_path}",
            )));
        }

        let branching_net = self.branching_q_network.as_mut().ok_or_else(|| {
            MLError::CheckpointError(
                "DQN::load_from_safetensors: branching_q_network not initialized".to_owned(),
            )
        })?;

        // Read + parse safetensors file.
        let file_bytes = std::fs::read(path_ref).map_err(|e| {
            MLError::CheckpointError(format!(
                "Failed to read safetensors '{safetensors_path}': {e}",
            ))
        })?;
        let tensors = safetensors::SafeTensors::deserialize(&file_bytes).map_err(|e| {
            MLError::CheckpointError(format!(
                "Failed to parse safetensors '{safetensors_path}': {e}",
            ))
        })?;

        // Figure out which naming scheme the file uses. The regime trainer
        // saves all three heads with `trending__` / `ranging__` / `volatile__`
        // prefixes; the single-head tests save plain names. Probe the first
        // shared layer's weight.
        let available: std::collections::HashSet<String> = tensors
            .names()
            .iter()
            .map(|n| (*n).to_owned())
            .collect();
        let prefix = if available.contains("trending__shared_0.weight") {
            "trending__"
        } else if available.contains("shared_0.weight") {
            ""
        } else {
            return Err(MLError::CheckpointError(format!(
                "Checkpoint '{safetensors_path}' has no recognised weight layout \
                 (expected either 'shared_0.weight' or 'trending__shared_0.weight'); \
                 file contains {} tensors",
                available.len(),
            )));
        };

        // Upload each expected tensor to GPU as a CudaSlice<f32>.
        let manifest: Vec<(String, Vec<usize>)> = branching_net
            .named_weight_slices()
            .into_iter()
            .map(|(name, _slice, shape)| (name, shape))
            .collect();

        let mut loaded: std::collections::HashMap<String, (cudarc::driver::CudaSlice<f32>, Vec<usize>)> =
            std::collections::HashMap::with_capacity(manifest.len());

        for (name, expected_shape) in &manifest {
            let file_key = format!("{prefix}{name}");
            let view = tensors.tensor(&file_key).map_err(|e| {
                MLError::CheckpointError(format!(
                    "Checkpoint '{safetensors_path}' missing '{file_key}': {e}",
                ))
            })?;
            if view.dtype() != safetensors::Dtype::F32 {
                return Err(MLError::CheckpointError(format!(
                    "Checkpoint '{file_key}' dtype {:?} is not F32",
                    view.dtype(),
                )));
            }
            let saved_shape: Vec<usize> = view.shape().to_vec();
            if &saved_shape != expected_shape {
                return Err(MLError::CheckpointError(format!(
                    "Checkpoint '{file_key}' shape mismatch: file={saved_shape:?} vs \
                     model={expected_shape:?}",
                )));
            }

            // Parse f32 little-endian bytes -> Vec<f32> and upload to GPU.
            let data_bytes = view.data();
            let numel: usize = expected_shape.iter().product();
            if data_bytes.len() != numel * std::mem::size_of::<f32>() {
                return Err(MLError::CheckpointError(format!(
                    "Checkpoint '{file_key}' byte-length {} does not match shape {:?} \
                     (expected {} bytes)",
                    data_bytes.len(),
                    expected_shape,
                    numel * std::mem::size_of::<f32>(),
                )));
            }
            let mut host: Vec<f32> = Vec::with_capacity(numel);
            for chunk_idx in 0..numel {
                let offset = chunk_idx * 4;
                let bytes = [
                    data_bytes[offset],
                    data_bytes[offset + 1],
                    data_bytes[offset + 2],
                    data_bytes[offset + 3],
                ];
                host.push(f32::from_le_bytes(bytes));
            }
            let gpu_slice = self.stream.clone_htod(&host).map_err(|e| {
                MLError::ModelError(format!(
                    "load_from_safetensors: failed to upload '{file_key}': {e}",
                ))
            })?;
            loaded.insert(name.clone(), (gpu_slice, expected_shape.clone()));
        }

        // Restore weights in the Q-network.
        branching_net.load_from_named_slices(&loaded)?;

        // Keep the target network in lockstep so `DQN::forward` and any subsequent
        // TD-target computation operate on the same restored weights.
        if let Some(target_net) = self.branching_target_network.as_mut() {
            target_net.load_from_named_slices(&loaded)?;
        }

        debug!(
            "DQN loaded branching weights from {safetensors_path} (prefix='{prefix}', \
             {} tensors restored)",
            manifest.len(),
        );

        Ok(())
    }

    /// Check if ready for training
    pub fn can_train(&self) -> bool {
        self.memory.can_sample(self.config.min_replay_size)
    }

    /// Check if using hybrid architecture (distributional + dueling)
    ///
    /// Wave 11.6: Helper method for architecture detection
    pub const fn is_using_hybrid(&self) -> bool {
        self.dist_dueling_q_network.is_some()
    }

    /// Get architecture type as string
    ///
    /// Wave 11.6: Returns one of:
    /// - "`hybrid_distributional_dueling`" (highest priority)
    /// - "dueling"
    /// - "distributional" (not used in hybrid, kept for compatibility)
    /// - "standard" (fallback)
    pub const fn get_architecture_type(&self) -> &'static str {
        if self.dist_dueling_q_network.is_some() {
            "hybrid_distributional_dueling"
        } else if self.dueling_q_network.is_some() {
            "dueling"
        } else {
            // Distributional is always enabled
            "distributional"
        }
    }

    /// Reinitialize categorical distribution with new bounds (adaptive C51)
    ///
    /// Used during two-phase training to adapt C51 bounds after feature normalization.
    ///
    /// # Arguments
    ///
    /// * `v_min` - New minimum value for distribution support
    /// * `v_max` - New maximum value for distribution support
    ///
    /// # Returns
    ///
    /// * `Ok(())` - Reinitialization successful
    /// * `Err(MLError)` - Failed to reinitialize distribution
    pub fn reinit_categorical_distribution(&mut self, v_min: f64, v_max: f64) -> Result<(), MLError> {
        if let Some(ref mut dist) = self.categorical_dist {
            dist.reinit(v_min, v_max)?;
        }
        Ok(())
    }

    /// Get the number of atoms for categorical distribution
    pub const fn get_num_atoms(&self) -> usize {
        self.config.num_atoms
    }

    /// Get the value support range (`v_min`, `v_max`) for distributional RL
    pub const fn get_v_range(&self) -> (f32, f32) {
        (self.config.v_min, self.config.v_max)
    }

    /// Check if distributional RL is enabled (always true -- distributional is unconditional)
    pub const fn is_using_distributional(&self) -> bool {
        true
    }

    // ========================================
    // Wave 11.5: N-Step Returns API Methods
    // ========================================

    /// Clear the n-step buffer (call at episode boundaries)
    ///
    /// This should be called at the start of a new episode to prevent
    /// experiences from different episodes being mixed in n-step returns.
    pub fn clear_nstep_buffer(&self) -> Result<(), MLError> {
        if let Some(ref nstep_buf) = self.nstep_buffer {
            let mut buffer = nstep_buf.lock().map_err(|e| MLError::ConcurrencyError {
                operation: format!("lock n-step buffer for clear: {}", e),
            })?;
            buffer.clear();
        }
        Ok(())
    }

    /// Get the configured n-step value
    ///
    /// Returns the number of steps used for multi-step TD learning.
    /// - n=1: Standard single-step TD (most stable)
    /// - n=3-5: Multi-step returns (better credit assignment, Rainbow DQN standard)
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use ml::dqn::{DQN, DQNConfig};
    /// # let agent = DQN::new(DQNConfig::aggressive())?;
    /// let n = agent.get_n_steps();
    /// println!("Agent using {}-step returns", n);
    /// # Ok::<(), ml::MLError>(())
    /// ```
    pub const fn get_n_steps(&self) -> usize {
        self.config.n_steps
    }

    /// Check if n-step returns are enabled
    ///
    /// Returns true if `n_steps` > 1 (multi-step TD learning enabled).
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use ml::dqn::{DQN, DQNConfig};
    /// # let agent = DQN::new(DQNConfig::conservative())?;
    /// if agent.is_using_nstep() {
    ///     println!("Using {}-step returns", agent.get_n_steps());
    /// } else {
    ///     println!("Using standard 1-step TD");
    /// }
    /// # Ok::<(), ml::MLError>(())
    /// ```
    pub const fn is_using_nstep(&self) -> bool {
        self.config.n_steps > 1
    }

}

#[cfg(test)]
#[cfg(feature = "cuda")]
mod save_load_tests {
    use super::*;

    /// Round-trip validation for `DQN::load_from_safetensors`.
    ///
    /// 1. Build a randomly-initialised DQN (source).
    /// 2. Dump its branching weights via `named_weight_slices` into a safetensors
    ///    file on disk (plain-name layout — no regime prefix).
    /// 3. Build a FRESH DQN with identical config.
    /// 4. Call `load_from_safetensors` on the fresh DQN.
    /// 5. Assert both DQNs' forward outputs agree on a deterministic input.
    ///
    /// `#[ignore]` because it requires a CUDA runtime.
    #[test]
    #[ignore = "requires CUDA runtime"]
    fn test_dqn_load_from_safetensors_round_trip() {
        // Small config matching the ensemble adapter's test expectations.
        let mut config = DQNConfig::emergency_safe_defaults();
        config.num_actions = 3;
        config.hidden_dims = vec![64, 64];

        // --- Phase 1: source DQN + save ---
        let src = DQN::new(config.clone()).expect("src DQN new");

        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("dqn_roundtrip.safetensors");

        {
            let stream = src.cuda_stream();
            let branching = src
                .branching_q_network
                .as_ref()
                .expect("branching_q_network expected to be initialised");
            let named = branching.named_weight_slices();

            // Download each slice to host and keep byte buffers alive for
            // the TensorView borrows.
            let byte_vecs: Vec<(String, Vec<usize>, Vec<u8>)> = named
                .iter()
                .map(|(name, slice, shape)| {
                    let mut host = vec![0.0_f32; slice.len()];
                    stream.memcpy_dtoh(*slice, &mut host).expect("DtoH");
                    let bytes: Vec<u8> =
                        host.iter().flat_map(|v| v.to_le_bytes()).collect();
                    (name.clone(), shape.clone(), bytes)
                })
                .collect();

            let mut tensor_map = std::collections::HashMap::new();
            for (name, shape, bytes) in &byte_vecs {
                let view = safetensors::tensor::TensorView::new(
                    safetensors::Dtype::F32,
                    shape.clone(),
                    bytes,
                )
                .expect("TensorView");
                tensor_map.insert(name.as_str(), view);
            }
            let arch_meta = Some(config.checkpoint_metadata());
            safetensors::serialize_to_file(tensor_map, arch_meta, path.as_path())
                .expect("serialize_to_file");
        }

        // --- Phase 2: fresh DQN + load ---
        let mut fresh = DQN::new(config.clone()).expect("fresh DQN new");
        let path_str = path.to_str().expect("utf8 path");
        fresh.load_from_safetensors(path_str).expect("load_from_safetensors");

        // --- Phase 3: forward outputs must match on a deterministic input ---
        let stream = src.cuda_stream();
        let state_dim = ml_core::state_layout::STATE_DIM;
        let state_host: Vec<f32> = (0..state_dim)
            .map(|i| ((i as f32) * 0.037).sin() * 0.25)
            .collect();
        let state = GpuTensor::from_host(&state_host, vec![1, state_dim], stream)
            .expect("state tensor");

        let q_src = src.forward(&state).expect("src forward").to_host(stream).expect("src DtoH");
        let q_fresh = fresh
            .forward(&state)
            .expect("fresh forward")
            .to_host(stream)
            .expect("fresh DtoH");

        assert_eq!(
            q_src.len(),
            q_fresh.len(),
            "Q-vector length should match: {} vs {}",
            q_src.len(),
            q_fresh.len(),
        );

        // NoisyNet resamples noise between `src.forward` and `fresh.forward`,
        // so perfectly bit-identical outputs are not guaranteed. However the
        // mu tensors are byte-identical post-load, so the argmax + relative
        // ordering of the Q-values should match within a loose tolerance.
        let max_abs_diff = q_src
            .iter()
            .zip(q_fresh.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);

        assert!(
            max_abs_diff < 1.0,
            "Post-load forward outputs diverged beyond noise tolerance: \
             max_abs_diff={max_abs_diff}, src={q_src:?}, fresh={q_fresh:?}",
        );
    }

    /// Loading from a path that doesn't exist must error with a CheckpointError,
    /// not silently succeed.
    #[test]
    #[ignore = "requires CUDA runtime"]
    fn test_dqn_load_from_safetensors_missing_file() {
        let mut config = DQNConfig::emergency_safe_defaults();
        config.num_actions = 3;
        config.hidden_dims = vec![64, 64];

        let mut dqn = DQN::new(config).expect("DQN new");
        let err = dqn
            .load_from_safetensors("/tmp/this_path_must_not_exist_for_test.safetensors")
            .expect_err("missing file must error");
        assert!(
            matches!(err, MLError::CheckpointError(_)),
            "expected CheckpointError, got: {err:?}",
        );
    }
}

