#![allow(unsafe_code)] // Required for CudaSlice reinterpret-cast in GPU PER insert path
//! DQN main training loop — `train_with_data_full_loop`.
//!
//! The main loop delegates to focused helper methods on `DQNTrainer`:
//! - `log_training_config`: one-time Rainbow/target-update logging
//! - `init_gpu_data`: upload training data to GPU (Phase 1)
//! - `init_gpu_raw_buffers`: upload raw cudarc targets/features + portfolio sim (Phase 1b)
//! - `init_gpu_experience_collector`: build the zero-roundtrip CUDA collector (Phase 1c)
//! - `collect_gpu_experiences`: run the GPU experience collection kernel (Phase 3)
//! - `run_training_steps`: batched training from replay buffer with guard kernel
//! - `process_epoch_boundary`: single readback + safety checks at epoch end
//! - `sync_gpu_weights`: push updated weights to GPU collector
//! - `refresh_stale_per_priorities`: M2 PER staleness refresh
//! - `log_epoch_metrics_and_financials`: logging, Prometheus, QuestDB, VaR/CVaR
//! - `handle_epoch_checkpoints_and_early_stopping`: best-checkpoint + early stopping

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::CudaSlice;
use common::metrics::{questdb_sink, training_metrics};
use tracing::{debug, info, warn};

use crate::cuda_pipeline::DqnGpuData;
use crate::dqn::logging::{log_epoch_start, log_epoch_end, log_training_progress};
use crate::dqn::target_update::convergence_half_life;
use crate::evaluation::metrics::calculate_var_cvar;
use crate::trainers::TargetUpdateMode;
use crate::TrainingMetrics;
use crate::features::extraction::FeatureVector;
use super::super::config::DQNAgentType;
use super::super::financials::compute_epoch_financials;
use super::super::monitoring::TrainingMonitor;
use super::DQNTrainer;

/// Metrics returned from `process_epoch_boundary` for downstream logging.
pub(crate) struct EpochBoundaryMetrics {
    pub avg_loss: f32,
    pub avg_grad: f32,
    pub avg_q: f64,
}

/// Metrics returned from `log_epoch_metrics_and_financials` for checkpointing.
pub(crate) struct EpochLogOutput {
    pub avg_loss: f64,
    pub avg_q_value: f64,
    pub avg_grad_norm: f64,
    pub epoch_sharpe: f64,
    pub val_loss: f64,
    pub q_min: f64,
    pub q_max: f64,
    pub q_mean: f64,
}

impl DQNTrainer {
    // ═══════════════════════════════════════════════════════════════════════
    // Main training loop — orchestrates helpers
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn train_with_data_full_loop<F>(
        &mut self,
        training_data: &[(FeatureVector, Vec<f64>)],
        mut checkpoint_callback: F,
    ) -> Result<TrainingMetrics>
    where
        F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
    {
        let start_time = std::time::Instant::now();
        let mut total_loss = 0.0;
        let mut total_q_value = 0.0;
        let mut total_gradient_norm = 0.0;
        let mut total_reward = 0.0;
        let mut total_action_counts = [0_usize; 9];
        let mut total_factored_action_counts = [0_usize; 81];

        self.log_training_config().await;

        // Decision Transformer pre-training on offline trajectories
        if self.hyperparams.dt_pretrain_epochs > 0 {
            info!(
                epochs = self.hyperparams.dt_pretrain_epochs,
                context_len = self.hyperparams.dt_context_len,
                embed_dim = self.hyperparams.dt_embed_dim,
                "Starting Decision Transformer pre-training"
            );

            // Ensure GPU data is uploaded before DT pre-training
            self.init_gpu_raw_buffers(training_data).await?;

            if let Some(ref stream) = self.cuda_stream {
                let stream = Arc::clone(stream);

                // Build DT config from hyperparameters.
                // DT state_dim = 42 (raw market features), NOT agent.get_state_dim()
                // which includes portfolio dims appended by the experience collector.
                // The raw features_raw_cuda buffer is [num_bars, 42].
                let dt_state_dim: usize = 42;

                let dt_config = crate::cuda_pipeline::decision_transformer::DecisionTransformerConfig {
                    state_dim: dt_state_dim,
                    num_actions: 9, // DT uses branch_0 exposure actions only
                    embed_dim: self.hyperparams.dt_embed_dim,
                    num_layers: self.hyperparams.dt_num_layers,
                    num_heads: 4,   // standard default
                    context_len: self.hyperparams.dt_context_len,
                    dropout: 0.1,
                    batch_size: self.hyperparams.batch_size.min(256),
                };

                let mut dt = crate::cuda_pipeline::decision_transformer::DecisionTransformer::new(
                    stream, dt_config,
                ).map_err(|e| anyhow::anyhow!("DT init: {e}"))?;

                // Build trajectories from GPU-resident features and targets
                let num_bars = training_data.len();
                if let (Some(ref features_gpu), Some(ref targets_gpu)) =
                    (&self.features_raw_cuda, &self.targets_raw_cuda)
                {
                    let (trajectories, target_actions, num_batches) = dt
                        .build_dt_trajectories(
                            features_gpu,
                            targets_gpu,
                            num_bars,
                            self.hyperparams.gamma,
                        )
                        .map_err(|e| anyhow::anyhow!("DT trajectory build: {e}"))?;

                    if num_batches == 0 {
                        warn!(
                            num_bars,
                            context_len = self.hyperparams.dt_context_len,
                            batch_size = dt.config().batch_size,
                            "DT pre-training: not enough data for a full batch, skipping"
                        );
                    } else {
                        let batch_size = dt.config().batch_size;
                        let context_len = dt.config().context_len;
                        let input_dim = dt.config().state_dim + 2;

                        for epoch in 0..self.hyperparams.dt_pretrain_epochs {
                            let mut epoch_loss = 0.0_f32;
                            let mut batch_count = 0_usize;

                            for batch_idx in 0..num_batches {
                                // Offset into the trajectory/target buffers for this batch.
                                // trajectories: [num_episodes, T, input_dim]
                                // Each batch is batch_size consecutive episodes.
                                let ep_offset = batch_idx * batch_size;
                                let traj_elem_offset = ep_offset * context_len * input_dim;
                                let act_elem_offset = ep_offset * context_len;

                                let loss = dt.pretrain_step(
                                    &trajectories,
                                    &target_actions,
                                    batch_size,
                                    traj_elem_offset,
                                    act_elem_offset,
                                ).map_err(|e| anyhow::anyhow!("DT pretrain_step: {e}"))?;

                                epoch_loss += loss;
                                batch_count += 1;
                            }

                            let avg_loss = if batch_count > 0 {
                                epoch_loss / batch_count as f32
                            } else {
                                0.0
                            };

                            info!(
                                epoch = epoch + 1,
                                total_epochs = self.hyperparams.dt_pretrain_epochs,
                                avg_loss = format!("{avg_loss:.4}"),
                                batches = batch_count,
                                "DT pre-training"
                            );

                            training_metrics::set_epoch(
                                "dt_pretrain",
                                "loss",
                                avg_loss as f64,
                            );
                        }

                        info!(
                            epochs = self.hyperparams.dt_pretrain_epochs,
                            "DT pre-training complete"
                        );
                    }
                } else {
                    warn!("DT pre-training: GPU features/targets not available, skipping");
                }
            } else {
                warn!("DT pre-training: no CUDA stream available, skipping");
            }
        }

        // Training loop
        for epoch in 0..self.hyperparams.epochs {
            self.current_epoch = epoch;
            self.reset_epoch_state(epoch);

            // C51 warmup: linearly ramp c51_alpha from 0→1 over warmup epochs.
            // MSE provides 10-100x stronger gradient when per-bar returns are ~0.001
            // and C51 cross-entropy between nearly-identical distributions produces
            // near-zero gradient for the value mean. Smooth blending avoids the
            // gradient discontinuity of a hard MSE→C51 switch.
            {
                let alpha = if self.hyperparams.c51_warmup_epochs > 0 {
                    (epoch as f32 / self.hyperparams.c51_warmup_epochs as f32).min(1.0)
                } else {
                    1.0
                };
                if let Some(ref mut fused) = self.fused_ctx {
                    fused.set_c51_alpha(alpha);
                }
            }

            // Q-gap warmup: ramp from 0 to configured threshold over first 5 epochs.
            // In early training, Q-values are near-random → Q-gaps are tiny → static
            // threshold forces ALL greedy actions to Flat, creating a "only learn Flat" loop.
            // Linear ramp lets exploration happen in early epochs.
            if self.hyperparams.q_gap_threshold > 0.0 {
                let warmup_epochs = 5.0_f32;
                let ramp = (epoch as f32 / warmup_epochs).min(1.0);
                let ramped_threshold = self.hyperparams.q_gap_threshold as f32 * ramp;
                if let Some(ref mut selector) = self.gpu_action_selector {
                    selector.set_q_gap_threshold(ramped_threshold);
                }
            }

            log_epoch_start(epoch + 1, self.hyperparams.epochs, self.hyperparams.learning_rate);
            training_metrics::set_epoch("dqn", "current", (epoch + 1) as f64);

            let mut monitor = TrainingMonitor::new(epoch + 1);

            // BUG #15 FIX: Portfolio compounds across epochs (see detailed comment in reset_epoch_state)
            self.portfolio_tracker.reset_drawdown_tracking();

            let epoch_start = std::time::Instant::now();

            // ── Phase 1: GPU data upload + experience collector init ──
            let phase1_start = std::time::Instant::now();
            self.init_gpu_data(training_data).await?;
            self.init_gpu_raw_buffers(training_data).await?;
            self.init_gpu_experience_collector().await?;
            let phase1_ms = phase1_start.elapsed().as_secs_f64() * 1000.0;

            // ── Phase 2: GPU experience collection ──
            let phase2_start = std::time::Instant::now();
            let gpu_experiences_collected = self.collect_gpu_experiences(
                training_data,
            ).await?;
            let phase2_ms = phase2_start.elapsed().as_secs_f64() * 1000.0;

            // CUDA builds: GPU experience collector is MANDATORY
            if !gpu_experiences_collected {
                return Err(anyhow::anyhow!(
                    "GPU experience collector MUST be active for CUDA training. \
                     No CPU fallback path exists in CUDA builds. \
                     or check GPU collector initialization errors above."
                ));
            }

            // ── Phase 3: Batched training from replay buffer ──
            let phase3_start = std::time::Instant::now();
            let train_step_count = self.run_training_steps(training_data).await?;
            let phase3_ms = phase3_start.elapsed().as_secs_f64() * 1000.0;

            // ── Phase 4: Epoch boundary + validation ──
            let phase4_start = std::time::Instant::now();
            let boundary = if train_step_count > 0 {
                let b = self.process_epoch_boundary(epoch, train_step_count, &mut monitor).await?;
                Some(b)
            } else {
                None
            };

            // Sync GPU weights after training
            self.sync_gpu_weights().await?;

            // ── Three-Phase Training Pipeline ──────────────────────────────
            //
            // Phase 1 (Behavioral Cloning): epochs 0..c51_warmup_epochs
            //   C51 alpha ramps 0→1 (MSE→C51 blend). Expert demos active.
            //   Already implemented via c51_alpha ramp + expert_demo_ratio above.
            //
            // Phase 2 (Full-Stack Online RL): c51_warmup_epochs..phase3_start
            //   All features active. Expert ratio decays to 0 via ExpertDemoGenerator.
            //   Already implemented via expert_demo_decay_epochs linear decay.
            //
            // Phase 3 (Refinement): last 20% of epochs
            //   Pure C51. No expert demos. Shrink-and-perturb for plasticity.
            //   Forces the network to consolidate learned policy without expert crutch.

            let total_epochs = self.hyperparams.epochs;
            let phase3_start = total_epochs * 80 / 100;
            let in_phase3 = epoch >= phase3_start && total_epochs > 4;

            if in_phase3 {
                // Phase 3: Force expert demos off (pure RL refinement)
                if let Some(ref mut c) = self.gpu_experience_collector {
                    c.set_expert_ratio(0.0);
                }

                // Phase 3: Shrink-and-perturb at the phase boundary for plasticity
                if epoch == phase3_start {
                    if let Some(ref mut fused) = self.fused_ctx {
                        let alpha = 0.9_f32; // 90% old weights, 10% noise
                        let sigma = 0.01_f32; // noise scale
                        if let Err(e) = fused.shrink_and_perturb(alpha, sigma) {
                            tracing::warn!("Shrink-and-Perturb failed at Phase 3 start (non-fatal): {e}");
                        } else {
                            tracing::info!(
                                epoch,
                                phase3_start,
                                alpha,
                                sigma,
                                "Phase 3 (Refinement): Shrink-and-Perturb applied at phase boundary"
                            );
                        }
                    }
                }
            }

            // Shrink-and-Perturb for plasticity maintenance (Phases 1-2).
            // Applied every 25% of total epochs (at epoch boundaries only).
            // Prevents the network from losing its ability to learn new patterns
            // across market regime shifts during long training runs.
            // In Phase 3, the boundary perturb above handles it — skip periodic.
            if !in_phase3 && self.hyperparams.epochs > 4 {
                let interval = (self.hyperparams.epochs / 4).max(1);
                if epoch > 0 && epoch % interval == 0 {
                    if let Some(ref mut fused) = self.fused_ctx {
                        let alpha = 0.9_f32; // 90% old weights, 10% noise
                        let sigma = 0.01_f32; // noise scale
                        if let Err(e) = fused.shrink_and_perturb(alpha, sigma) {
                            tracing::warn!("Shrink-and-Perturb failed (non-fatal): {e}");
                        } else {
                            tracing::info!(
                                epoch,
                                alpha,
                                sigma,
                                "Shrink-and-Perturb applied (plasticity maintenance)"
                            );
                        }
                    }
                }
            }

            // Compute CVaR scales from IQN head for next epoch's experience collection.
            // C51 picks direction, IQN sizes the risk via CVaR at 5th percentile.
            if let Some(ref mut fused) = self.fused_ctx {
                let bs = fused.batch_size();
                // Clone the actions CudaSlice reference to avoid borrow conflict
                // (actions_buf borrows immutably, compute_cvar borrows mutably)
                let actions_clone = {
                    let actions = fused.actions_buf();
                    actions.clone()
                };
                let cvar_ptr = fused.compute_cvar_device_ptr(&actions_clone, bs, 0.05)
                    .unwrap_or(0);
                if let Some(ref mut collector) = self.gpu_experience_collector {
                    collector.set_cvar_scales(cvar_ptr);
                }
            }

            // Flush GPU-accumulated max priority
            if train_step_count > 0 {
                let agent = self.agent.read().await;
                if let Err(e) = agent.flush_max_priority() {
                    debug!("GPU max_priority flush failed (non-fatal): {}", e);
                }
            }
            let phase4_ms = phase4_start.elapsed().as_secs_f64() * 1000.0;

            let epoch_duration = epoch_start.elapsed();

            // ── Epoch phase breakdown ──
            info!(
                "Epoch {}/{} phase breakdown: init={:.0}ms experience={:.0}ms training={:.0}ms validation={:.0}ms total={:.0}ms",
                epoch + 1, self.hyperparams.epochs,
                phase1_ms, phase2_ms, phase3_ms, phase4_ms,
                epoch_duration.as_secs_f64() * 1000.0,
            );

            // Calculate and log epoch metrics
            let log_output = self.log_epoch_metrics_and_financials(
                epoch,
                train_step_count,
                &boundary,
                &mut monitor,
                epoch_duration,
                &mut total_action_counts,
                &mut total_factored_action_counts,
            ).await?;

            total_loss += log_output.avg_loss;
            total_q_value += log_output.avg_q_value;
            total_gradient_norm += log_output.avg_grad_norm;

            let epoch_avg_reward = if !monitor.reward_history.is_empty() {
                monitor.reward_history.iter().sum::<f32>() / monitor.reward_history.len() as f32
            } else {
                0.0
            };
            total_reward += epoch_avg_reward as f64;

            // Checkpoints + early stopping (returns Err to signal hyperopt on stop)
            self.handle_epoch_checkpoints_and_early_stopping(
                epoch,
                train_step_count,
                &log_output,
                &mut checkpoint_callback,
            ).await?;

            // WAVE 13-A2: Save periodic checkpoint every N epochs
            if (epoch + 1) % self.hyperparams.checkpoint_frequency == 0 {
                info!(
                    "Saving periodic checkpoint at epoch {}/{}",
                    epoch + 1, self.hyperparams.epochs
                );
                let checkpoint_data = self.serialize_model().await?;
                let checkpoint_size = checkpoint_data.len();
                let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false)
                    .context("Failed to save periodic checkpoint")?;
                info!(
                    "Periodic checkpoint saved: {} ({} bytes)",
                    checkpoint_path, checkpoint_size
                );
            }
        }

        let training_duration = start_time.elapsed();

        let metrics = self.create_final_metrics(
            total_loss, total_q_value, total_gradient_norm, total_reward,
            self.hyperparams.epochs, training_duration, false,
            total_action_counts, total_factored_action_counts,
        ).await?;

        {
            let mut stored_metrics = self.metrics.write().await;
            *stored_metrics = metrics.clone();
        }

        info!(
            "Training completed in {:.2}s: final_loss={:.6}, avg_q_value={:.4}",
            training_duration.as_secs_f64(),
            metrics.loss,
            metrics.additional_metrics.get("avg_q_value").unwrap_or(&0.0)
        );

        info!("Best model summary:");
        info!(
            "  Best Sharpe: {:.4} at epoch {} (val_loss={:.6})",
            self.best_sharpe, self.best_epoch, self.best_val_loss
        );
        info!("  Best model checkpoint: best_model.safetensors");

        Ok(metrics)
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: one-time training configuration logging
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn log_training_config(&mut self) {
        match self.hyperparams.target_update_mode {
            TargetUpdateMode::Soft => {
                let half_life = convergence_half_life(self.hyperparams.tau);
                info!("WAVE 16: Using soft target updates (Polyak averaging)");
                info!("  Tau: {}", self.hyperparams.tau);
                info!("  Convergence half-life: {} steps", half_life as usize);
                info!("  Strategy: Smooth Q-value tracking (50-70% variance reduction)");
            },
            TargetUpdateMode::Hard => {
                info!("WAVE 16: Using hard target updates (legacy mode)");
                info!("  Update frequency: every 1000 steps");
                info!("  Warning: Sudden Q-value shifts may cause instability");
            },
        }

        info!("Rainbow DQN Components:");
        info!("  Double DQN (always enabled)");

        if true {
            info!("  Dueling Networks (value/advantage streams, hidden_dim={})", self.hyperparams.dueling_hidden_dim);
        } else {
            info!("  Dueling Networks: disabled");
        }

        if true {
            info!("  Prioritized Experience Replay (a={}, b={}->1.0)",
                  self.hyperparams.per_alpha, self.hyperparams.per_beta_start);
        } else {
            info!("  Prioritized Experience Replay: disabled");
        }

        if self.hyperparams.n_steps > 1 {
            info!("  N-Step Returns (n={})", self.hyperparams.n_steps);
        } else {
            info!("  N-Step Returns (n=1, standard TD)");
        }

        if true {
            info!("  Categorical DQN (atoms={}, V=[{}, {}])",
                  self.hyperparams.num_atoms, self.hyperparams.v_min, self.hyperparams.v_max);
        } else {
            info!("  Categorical DQN / C51: disabled");
        }

        if true {
            info!("  Noisy Networks (sigma_init={})", self.hyperparams.noisy_sigma_init);
        } else {
            info!("  Noisy Networks: disabled");
        }

        if true {
            let floor = self.hyperparams.noisy_epsilon_floor.unwrap_or(0.05);
            let mut agent = self.agent.write().await;
            agent.set_epsilon(floor);
            info!("  Noisy nets active: epsilon set to noisy_epsilon_floor={:.4}", floor);
        }
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: reset per-epoch state
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) fn reset_epoch_state(&mut self, _epoch: usize) {
        self.metrics_aggregator.reset();
        self.reward_fn.reset_epoch_state();
        self.portfolio_tracker.reset();

        // GPU-persistent epoch state: set reset flags instead of CPU state mutation.
        if let Some(ref mut collector) = self.gpu_experience_collector {
            let flags: u32 = 1 | 2; // reset portfolio + DSR normalizer (always enabled)
            collector.set_reset_flags(flags);
        }
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Phase 1 — GPU data upload (once)
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn init_gpu_data(
        &mut self,
        training_data: &[(FeatureVector, Vec<f64>)],
    ) -> Result<()> {
        if self.gpu_data.is_some() {
            return Ok(());
        }

        // Check double-buffer: skip upload if active slot already populated
        let skip_upload = self.double_buffer.as_ref().is_some_and(|db| db.active().is_some());

        if skip_upload {
            info!("DoubleBuffer: active slot populated, skipping re-upload");
            return Ok(());
        }

        let upload_result = if let Some(ref mut pool) = self.buffer_pool {
            info!("GpuBufferPool: reusing pre-allocated staging buffers for {} bars", training_data.len());
            pool.upload_dqn(training_data, self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream required for upload"))?)
        } else {
            DqnGpuData::upload(training_data, self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream required for DqnGpuData upload"))?)
        };

        match upload_result {
            Ok(mut gpu_data) => {
                info!("GPU data pre-uploaded: {} bars x {} features ({:.1} MB)",
                    gpu_data.num_bars,
                    gpu_data.feature_dim,
                    (gpu_data.num_bars * (42 + 4) * 4) as f64 / 1_048_576.0
                );
                // Upload OFI features to GPU if available
                if let Some(ref ofi) = self.ofi_features {
                    match gpu_data.upload_ofi(ofi, self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream for OFI upload"))?) {
                        Ok(()) => info!("GPU OFI features uploaded: {} bars x 8 dims", ofi.len()),
                        Err(e) => {
                            return Err(anyhow::anyhow!("GPU OFI upload FAILED (no CPU fallback): {e}"));
                        }
                    }
                }
                // Set tensor core alignment so build_*_states pads output
                let ofi_enabled = self.ofi_features.is_some();
                let raw_dim = if ofi_enabled { 53 } else { 45 };
                let aligned_dim = (raw_dim + 7) & !7;
                gpu_data.set_aligned_state_dim(aligned_dim);
                self.gpu_data = Some(gpu_data);
            }
            Err(e) => {
                return Err(anyhow::anyhow!("GPU data pre-upload FAILED (no CPU fallback): {e}"));
            }
        }
        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Phase 1b — raw cudarc targets + features + portfolio sim (once)
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn init_gpu_raw_buffers(
        &mut self,
        training_data: &[(FeatureVector, Vec<f64>)],
    ) -> Result<()> {
        if self.targets_raw_cuda.is_some() {
            return Ok(());
        }

        let stream = match self.cuda_stream {
            Some(ref s) => Arc::clone(s),
            None => return Ok(()),
        };

        let target_dim = 4;
        let feature_dim = 42;
        let num_bars = training_data.len();

        // Build flat targets (same as DqnGpuData::upload but for cudarc)
        let mut flat_targets = Vec::with_capacity(num_bars * target_dim);
        for (_, targets) in training_data {
            for i in 0..target_dim {
                flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32);
            }
        }

        match stream.clone_htod(&flat_targets) {
            Ok(buf) => {
                info!("CUDA targets_raw uploaded: {} bars x 4 ({:.1} KB)",
                    num_bars, (num_bars * target_dim * 4) as f64 / 1024.0);
                self.targets_raw_cuda = Some(buf);
            }
            Err(e) => {
                return Err(anyhow::anyhow!("CUDA targets_raw upload FAILED (no CPU fallback): {e}"));
            }
        }

        // Build flat features [num_bars * 42] for GPU experience kernel
        if self.features_raw_cuda.is_none() {
            let mut flat_features = Vec::with_capacity(num_bars * feature_dim);
            for (features, _) in training_data {
                for &v in features.iter() {
                    flat_features.push(v as f32);
                }
            }
            match stream.clone_htod(&flat_features) {
                Ok(buf) => {
                    info!("CUDA features_raw uploaded: {} bars x {} ({:.1} KB)",
                        num_bars, feature_dim, (num_bars * feature_dim * 4) as f64 / 1024.0);
                    self.features_raw_cuda = Some(buf);
                }
                Err(e) => {
                    return Err(anyhow::anyhow!("CUDA features_raw upload FAILED (no CPU fallback): {e}"));
                }
            }
        }

        // GpuPortfolioSimulator removed — backtest evaluation uses GpuBacktestEvaluator.

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Phase 1c — GPU experience collector init (once)
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn init_gpu_experience_collector(&mut self) -> Result<()> {
        if self.gpu_experience_collector.is_some() {
            return Ok(());
        }

        let stream = match self.cuda_stream {
            Some(ref s) => Arc::clone(s),
            None => return Ok(()),
        };

        use crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector;

        // Read-lock agent to extract dueling network weights for GPU collector.
        let agent = self.agent.read().await;
        let dqn_ref: Option<&crate::dqn::DQN> = match &*agent {
            DQNAgentType::Standard(ref dqn) => Some(dqn),
            DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()),
        };

        let init_result = if let Some(dqn) = dqn_ref {
            let curiosity_vars = self.curiosity_module.as_ref()
                .map(|c| c.forward_model_vars());

            let state_dim = agent.get_state_dim();
            let market_dim: usize = 42;
            let num_atoms_max = self.hyperparams.num_atoms as usize;
            let kernel_dims = (state_dim, market_dim, num_atoms_max);

            // Priority: branching > plain dueling > hybrid (distributional+dueling)
            if let (Some(online_br), Some(target_br)) = (
                dqn.branching_q_network.as_ref(),
                dqn.branching_target_network.as_ref(),
            ) {
                let cfg = online_br.config();
                let dims = (
                    *cfg.shared_hidden_dims.first().unwrap_or(&256),
                    *cfg.shared_hidden_dims.get(1).unwrap_or(&256),
                    cfg.value_hidden_dim,
                    cfg.branch_hidden_dim,
                );
                let alloc_episodes = self.compute_alloc_episodes(state_dim);
                Some(GpuExperienceCollector::new(
                    stream,
                    online_br.vars(),
                    target_br.vars(),
                    curiosity_vars,
                    self.hyperparams.initial_capital as f32,
                    self.hyperparams.avg_spread as f32,
                    self.hyperparams.cash_reserve_percent as f32,
                    dims,
                    kernel_dims,
                    alloc_episodes,
                    self.hyperparams.gpu_timesteps_per_episode,
                ))
            } else if let (Some(online), Some(target)) = (
                dqn.dueling_q_network.as_ref(),
                dqn.dueling_target_network.as_ref(),
            ) {
                let cfg = online.config();
                let dims = (
                    *cfg.shared_hidden_dims.first().unwrap_or(&256),
                    *cfg.shared_hidden_dims.get(1).unwrap_or(&256),
                    cfg.value_hidden_dim,
                    cfg.advantage_hidden_dim,
                );
                let alloc_episodes = self.compute_alloc_episodes(state_dim);
                Some(GpuExperienceCollector::new(
                    stream,
                    online.store(),
                    target.store(),
                    curiosity_vars,
                    self.hyperparams.initial_capital as f32,
                    self.hyperparams.avg_spread as f32,
                    self.hyperparams.cash_reserve_percent as f32,
                    dims,
                    kernel_dims,
                    alloc_episodes,
                    self.hyperparams.gpu_timesteps_per_episode,
                ))
            } else if let (Some(online), Some(target)) = (
                dqn.dist_dueling_q_network.as_ref(),
                dqn.dist_dueling_target_network.as_ref(),
            ) {
                let cfg = online.config();
                let dims = (
                    *cfg.shared_hidden_dims.first().unwrap_or(&256),
                    *cfg.shared_hidden_dims.get(1).unwrap_or(&256),
                    cfg.value_hidden_dim,
                    cfg.advantage_hidden_dim,
                );
                let alloc_episodes = self.compute_alloc_episodes(state_dim);
                Some(GpuExperienceCollector::new(
                    stream,
                    online.vars(),
                    target.vars(),
                    curiosity_vars,
                    self.hyperparams.initial_capital as f32,
                    self.hyperparams.avg_spread as f32,
                    self.hyperparams.cash_reserve_percent as f32,
                    dims,
                    kernel_dims,
                    alloc_episodes,
                    self.hyperparams.gpu_timesteps_per_episode,
                ))
            } else {
                return Err(anyhow::anyhow!(
                    "GPU experience collector FAILED: no dueling, hybrid, or branching networks — \
                     DQN agent has no Q-network. This is a configuration bug."
                ));
            }
        } else {
            None
        };
        drop(agent);

        if let Some(result) = init_result {
            match result {
                Ok(mut collector) => {
                    // Lazy-init GPU monitoring reducer on same stream as collector
                    if self.gpu_monitoring.is_none() {
                        match crate::cuda_pipeline::gpu_monitoring::GpuMonitoringReducer::new(collector.stream()) {
                            Ok(mon) => {
                                info!("GPU monitoring reducer initialized");
                                self.gpu_monitoring = Some(mon);
                            }
                            Err(e) => {
                                return Err(anyhow::anyhow!(
                                    "GPU monitoring reducer init FAILED: {e} — metrics pipeline must be GPU-resident"
                                ));
                            }
                        }
                    }
                    // Lazy-init reduction kernels for fused Q-value stats/diagnostics
                    if self.reduction_kernels.is_none() {
                        match ml_core::cuda_autograd::ReductionKernels::new(collector.stream()) {
                            Ok(rk) => {
                                info!("GPU reduction kernels initialized");
                                self.reduction_kernels = Some(rk);
                            }
                            Err(e) => {
                                return Err(anyhow::anyhow!(
                                    "GPU reduction kernels init FAILED: {e} — diagnostics pipeline must be GPU-resident"
                                ));
                            }
                        }
                    }
                    // Upload OFI features to GPU if available.
                    if let Some(ref ofi) = self.ofi_features {
                        let flat: Vec<f32> = ofi.iter()
                            .flat_map(|f| f.iter().map(|&v| v as f32))
                            .collect();
                        collector.upload_ofi_features(&flat)?;
                    }
                    self.gpu_experience_collector = Some(collector);
                }
                Err(e) => {
                    return Err(anyhow::anyhow!(
                        "GPU experience collector init FAILED (no CPU fallback): {e}"
                    ));
                }
            }
        }
        Ok(())
    }

    /// Compute allocation episode count for GPU buffers.
    /// Auto-scales when configured >= 128 (production), respects smaller test overrides.
    fn compute_alloc_episodes(&self, state_dim: usize) -> usize {
        use ml_core::memory_optimization::detect_gpu_hardware;
        let configured = self.hyperparams.gpu_n_episodes;
        if configured >= 128 {
            match detect_gpu_hardware() {
                Ok(hw) => configured.max(hw.optimal_n_episodes(
                    state_dim,
                    self.hyperparams.gpu_timesteps_per_episode,
                )).min(0x8000),
                Err(_) => configured,
            }
        } else {
            configured
        }
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Phase 3 — GPU experience collection
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn collect_gpu_experiences(
        &mut self,
        training_data: &[(FeatureVector, Vec<f64>)],
    ) -> Result<bool> {
        // Feature normalization stats calculation epoch (kept for API compat)
        let _stats_collection_epochs = {
            let ratio_based = (self.hyperparams.epochs as f32 * self.hyperparams.feature_stats_collection_ratio) as usize;
            let capped = match self.hyperparams.max_feature_stats_epochs {
                Some(max_epochs) => ratio_based.min(max_epochs),
                None => ratio_based,
            };
            capped.max(1)
        };

        // Set expert demo ratio BEFORE borrowing collector (avoids borrow conflict)
        if self.hyperparams.expert_demo_ratio > 0.0 {
            use crate::trainers::dqn::expert_demos::ExpertDemoGenerator;
            let effective = ExpertDemoGenerator::effective_ratio(
                self.hyperparams.expert_demo_ratio,
                self.hyperparams.expert_demo_decay_epochs,
                self.current_epoch,
            ) as f32;
            if let Some(ref mut c) = self.gpu_experience_collector {
                c.set_expert_ratio(effective);
            }
        }

        let (
            Some(ref mut collector),
            Some(ref features_buf),
            Some(ref targets_buf),
        ) = (
            &mut self.gpu_experience_collector,
            &self.features_raw_cuda,
            &self.targets_raw_cuda,
        ) else {
            return Ok(false);
        };

        use crate::cuda_pipeline::gpu_experience_collector::ExperienceCollectorConfig;

        let raw_sd = if self.hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 };
        let aligned_sd = (raw_sd + 7) & !7;

        // Cache n_episodes on first epoch
        let n_episodes = if let Some(cached) = self.cached_n_episodes {
            cached
        } else {
            use ml_core::memory_optimization::detect_gpu_hardware;
            let configured = self.hyperparams.gpu_n_episodes;
            let computed = if configured >= 128 {
                match detect_gpu_hardware() {
                    Ok(hw) => {
                        let optimal = hw.optimal_n_episodes(
                            aligned_sd,
                            self.hyperparams.gpu_timesteps_per_episode,
                        );
                        let chosen = configured.max(optimal).min(4096);
                        if chosen != configured {
                            info!(
                                "GPU auto-scaled n_episodes: {} -> {} (SMs={}, VRAM={:.0}MB)",
                                configured, chosen, hw.sm_count, hw.free_memory_mb
                            );
                        }
                        chosen as i32
                    }
                    Err(_) => configured as i32,
                }
            } else {
                configured as i32
            };
            self.cached_n_episodes = Some(computed);
            computed
        };

        let timesteps = self.hyperparams.gpu_timesteps_per_episode.min(1000) as i32;
        let total_bars = training_data.len() as i32;
        let usable_bars = (total_bars - timesteps).max(1);
        let stride = (usable_bars / n_episodes).max(1);
        let mut episode_starts: Vec<i32> = (0..n_episodes)
            .map(|i| (i * stride).rem_euclid(usable_bars))
            .collect();

        // ── Curriculum learning: filter episode starts by difficulty ──────
        // Phase 0 (Easy): only start episodes at trending bars (ADX > 30)
        // Phase 1 (Mixed): all starts, but duplicate trending ones (2x weight)
        // Phase 2 (Full): all starts, equal weight (default)
        let curriculum = self.hyperparams.curriculum_phase;
        if curriculum < 2 && !training_data.is_empty() {
            let adx_idx = 40; // ADX at market feature index 40
            let adx_threshold = 30.0_f64;
            let original = episode_starts.clone();
            episode_starts.clear();

            for &start in &original {
                let bar = (start as usize).min(training_data.len().saturating_sub(1));
                let features = &training_data[bar].0;
                let adx = if features.len() > adx_idx { features[adx_idx] } else { 0.0 };

                match curriculum {
                    0 => {
                        // Easy: only trending bars
                        if adx > adx_threshold {
                            episode_starts.push(start);
                        }
                    }
                    1 => {
                        // Mixed: all bars, trending bars duplicated
                        episode_starts.push(start);
                        if adx > adx_threshold {
                            episode_starts.push(start);
                        }
                    }
                    _ => episode_starts.push(start),
                }
            }

            // Ensure at least 16 episodes (fall back to Full if too few)
            if episode_starts.len() < 16 {
                tracing::warn!(
                    curriculum,
                    filtered = episode_starts.len(),
                    original = original.len(),
                    "Curriculum phase filtered too aggressively, falling back to Full"
                );
                episode_starts = original;
            } else {
                tracing::debug!(
                    curriculum,
                    filtered = episode_starts.len(),
                    original = original.len(),
                    "Curriculum learning: episode starts filtered by ADX difficulty"
                );
            }
        }

        // Reset per-episode state before each epoch
        if let Err(e) = collector.reset_episodes(
            self.hyperparams.initial_capital as f32,
            self.hyperparams.avg_spread as f32,
            self.hyperparams.cash_reserve_percent as f32,
        ) {
            return Err(anyhow::anyhow!("GPU episode reset FAILED (no CPU fallback): {e}"));
        }

        let agent = self.agent.read().await;
        let epsilon = agent.get_epsilon();
        drop(agent);

        // Expert ratio already set above (before collector borrow).

        let config = ExperienceCollectorConfig {
            n_episodes,
            timesteps_per_episode: timesteps,
            total_bars,
            epsilon,
            gamma: self.hyperparams.gamma as f32,
            max_position: self.max_position as f32,
            enable_action_masking: self.enable_action_masking,
            curiosity_scale: if self.curiosity_module.is_some() { 1.0 } else { 0.0 },
            loss_aversion: self.hyperparams.loss_aversion as f32,
            tx_cost_multiplier: self.hyperparams.transaction_cost_multiplier as f32,
            count_bonus_coefficient: self.hyperparams.count_bonus_coefficient
                .unwrap_or(0.0) as f32,
            q_clip_min: self.hyperparams.q_clip_min as f32,
            q_clip_max: self.hyperparams.q_clip_max as f32,
            huber_kappa: if true {
                self.hyperparams.huber_delta as f32
            } else {
                0.0
            },
            use_noisy_nets: true,
            noisy_sigma_init: self.hyperparams.noisy_sigma_init as f32,
            use_distributional: true,
            num_atoms: self.hyperparams.num_atoms as i32,
            v_min: self.hyperparams.v_min as f32,
            v_max: self.hyperparams.v_max as f32,
            fill_median_spread: self.hyperparams.avg_spread as f32,
            fill_median_vol: self.median_vol as f32,
            fill_ioc_fill_prob: self.hyperparams.fill_ioc_fill_prob as f32,
            fill_limit_fill_min: self.hyperparams.fill_limit_fill_min as f32,
            fill_limit_fill_max: self.hyperparams.fill_limit_fill_max as f32,
            fill_spread_cost_frac: self.hyperparams.fill_spread_cost_frac as f32,
            fill_spread_capture_frac: self.hyperparams.fill_spread_capture_frac as f32,
            fill_simulation_enabled: self.median_vol > 0.0,
            // Q-gap cold-start fix: ramp threshold from 0 to target over 10 epochs.
            // At epoch 0, Q-values are random — forcing Q-gap > 0.1 prevents ALL
            // trades (except 10% epsilon), starving the model of learning signal.
            // Linear ramp: epoch 0→0.0, epoch 5→0.05, epoch 10→full threshold.
            q_gap_threshold: {
                let target = self.hyperparams.q_gap_threshold as f32;
                let epoch = self.current_epoch.min(10) as f32;
                target * (epoch / 10.0)
            },
            dsr_eta: self.hyperparams.dsr_eta as f32,
            n_steps: self.hyperparams.n_steps as i32,
            ..Default::default()
        };

        // Zero-roundtrip GPU path — GPU PER is always active in CUDA builds
        let gpu_batch = collector.collect_experiences_gpu(
            features_buf, targets_buf, &episode_starts, &config,
        ).map_err(|e| anyhow::anyhow!(
            "GPU zero-roundtrip collection FAILED (no CPU fallback): {e}"
        ))?;

        let count = gpu_batch.n_episodes * gpu_batch.timesteps;

        // Synchronize forked stream before cross-stream operations.
        // GpuReplayBuffer::insert_batch_tensors uses Candle GpuTensor ops (slice_scatter,
        // cumsum, narrow, etc.) which execute on Candle's default stream (stream 0).
        // The experience collector outputs data on the forked CudaStream. Without this
        // barrier, the default stream may read partially-written data, causing a
        // cross-stream deadlock or data corruption.
        if let Some(ref stream) = self.cuda_stream {
            stream.synchronize()
                .map_err(|e| anyhow::anyhow!("Stream sync before PER insert: {e}"))?;
        }

        if let Some(ref mut mon) = self.gpu_monitoring {
            if let Err(e) = mon.reduce(collector.rewards_gpu(), collector.actions_gpu(), count) {
                debug!("GPU monitoring reduce failed: {e}");
            }
            // Catch async monitoring kernel crashes before they poison the stream
            if let Some(ref stream) = self.cuda_stream {
                if stream.synchronize().is_err() {
                    warn!("GPU monitoring kernel crashed — disabling monitoring");
                    self.gpu_monitoring = None;
                }
            }
        }

        if count > 0 {
            if let Err(e) = collector.train_curiosity_gpu(
                gpu_batch.n_episodes, gpu_batch.timesteps,
            ) {
                debug!("GPU curiosity training failed (non-fatal): {e}");
            }
            // Sync after curiosity to detect async kernel crashes immediately.
            // Without this guard, a silent CUDA error from the curiosity kernel
            // poisons the stream, and the NEXT stream.synchronize() (in the PER
            // insert path) deadlocks forever instead of returning the error.
            if let Some(ref stream) = self.cuda_stream {
                if let Err(e) = stream.synchronize() {
                    warn!("Curiosity kernel async error detected: {e} — disabling curiosity for this run");
                    self.curiosity_module = None;
                    collector.disable_curiosity();
                }
            }
        }

        if count > 0 {
            // Direct CudaSlice insertion — zero Candle, zero CPU roundtrip.
            let total = gpu_batch.n_episodes * gpu_batch.timesteps;

            // actions: CudaSlice<i32> → CudaSlice<u32> (safe reinterpret for values [0..44])
            let actions_u32: &CudaSlice<u32> = unsafe {
                &*(&gpu_batch.actions as *const CudaSlice<i32> as *const CudaSlice<u32>)
            };

            // dones: already CudaSlice<f32> (kernel outputs float 0.0/1.0)
            let agent = self.agent.read().await;
            let mut gpu_buf = agent.memory().as_gpu_buffer()
                .ok_or_else(|| anyhow::anyhow!("GPU PER buffer required"))?;
            gpu_buf.gpu.insert_batch(
                &gpu_batch.states,
                &gpu_batch.next_states,
                actions_u32,
                &gpu_batch.rewards,
                &gpu_batch.dones,
                total,
            ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?;
        }

        // Sync GPU DSR EMA state back to CPU for logging/checkpointing.
        // The GPU kernel's dsr_step() maintains its own accumulators in epoch_state[5..6].
        // After experience collection, read them back so the CPU RewardFunction stays current.
        match collector.read_epoch_dsr_state() {
            Ok((dsr_a, dsr_b)) => {
                self.reward_fn.sync_dsr_from_gpu(dsr_a as f64, dsr_b as f64);
                debug!(
                    "DSR GPU->CPU sync: ema_return={:.6}, ema_return_sq={:.6}",
                    dsr_a, dsr_b
                );
            }
            Err(e) => {
                debug!("DSR epoch_state readback failed (non-fatal): {e}");
            }
        }

        Ok(true)
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Batched training steps with guard kernel
    // ═══════════════════════════════════════════════════════════════════════

    /// Runs all training steps for one epoch, returning the step count.
    pub(crate) async fn run_training_steps(
        &mut self,
        training_data: &[(FeatureVector, Vec<f64>)],
    ) -> Result<usize> {
        let batch_size = self.hyperparams.batch_size;
        let num_training_steps = if self.can_train().await? {
            let full_steps = (training_data.len() / batch_size).max(1);
            let cap = self.hyperparams.max_training_steps_per_epoch;
            if cap > 0 { full_steps.min(cap) } else { full_steps }
        } else {
            0
        };

        let mut train_step_count = 0;

        // Lazy-init fused CUDA training context (Standard DQN only).
        // Recreate if batch_size changed (OOM recovery).
        if num_training_steps > 0 {
            self.ensure_fused_ctx().await?;
        }

        // Lazy-init + reset guard accumulators for this epoch
        {
            if self.training_guard.is_none() && self.device.is_cuda() {
                match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::from_stream(Arc::clone(self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream for training guard"))?)) {
                    Ok(g) => {
                        info!("GPU training guard initialized (epoch loop)");
                        self.training_guard = Some(g);
                    }
                    Err(e) => return Err(anyhow::anyhow!("GPU training guard init: {e}")),
                }
            }
            if let Some(ref mut guard) = self.training_guard {
                guard.reset_accumulators()
                    .map_err(|e| anyhow::anyhow!("guard reset: {e}"))?;
            }
        }

        let guard_collapse_thresh =
            self.hyperparams.learning_rate as f32
                * self.hyperparams.gradient_collapse_multiplier as f32;
        let guard_past_warmup = {
            let ws = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64;
            self.gradient_logging_step as u64 > ws
        };

        // Batch pre-sampling: sample K batches under one READ lock
        const PREFETCH_K: usize = 32;
        let mut sample_total_us = 0_u64;
        let mut fused_total_us = 0_u64;
        let mut guard_total_us = 0_u64;
        for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) {
            let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps);

            let sample_start = std::time::Instant::now();
            let batches = {
                let agent = self.agent.read().await;
                let buffer = agent.memory();
                let can = buffer.can_sample(self.current_batch_size);
                let buf_len = buffer.len();
                tracing::debug!(
                    can_sample = can,
                    buffer_len = buf_len,
                    batch_size = self.current_batch_size,
                    chunk = chunk_start,
                    "Pre-sampling replay buffer"
                );
                let mut b = Vec::with_capacity(chunk_end - chunk_start);
                for _ in chunk_start..chunk_end {
                    b.push(buffer.can_sample(self.current_batch_size).then(|| {
                        buffer
                            .sample(self.current_batch_size)
                            .map_err(|e| anyhow::anyhow!("Pre-sample: {e}"))
                    }).transpose()?);
                }
                b
            };
            sample_total_us += sample_start.elapsed().as_micros() as u64;

            // GPU train steps (single WRITE lock)
            {
                let mut agent = self.agent.write().await;
                let accum_steps = self.hyperparams.gradient_accumulation_steps;

                if accum_steps <= 1 {
                    for explicit_batch in batches {
                        let fused_start = std::time::Instant::now();
                        let _gpu_result = if let Some(ref mut fused) = self.fused_ctx {
                            let batch_data = explicit_batch.as_ref().ok_or_else(|| {
                                anyhow::anyhow!("No batch data for fused training step")
                            })?;
                            fused.run_full_step(batch_data, &mut *agent, &self.device)
                                .context("Fused CUDA training step failed")?
                        } else {
                            agent.train_step(explicit_batch)
                                .map_err(|e| anyhow::anyhow!("Train step FAILED: {e}"))?
                        };

                        fused_total_us += fused_start.elapsed().as_micros() as u64;

                        let guard_start = std::time::Instant::now();
                        if let Some(ref mut guard) = self.training_guard {
                            let loss_slice = _gpu_result.loss_cuda_slice()
                                .map_err(|e| anyhow::anyhow!("guard loss CudaSlice: {e}"))?;
                            let grad_slice = _gpu_result.grad_norm_cuda_slice()
                                .map_err(|e| anyhow::anyhow!("guard grad CudaSlice: {e}"))?;
                            let gr = guard.check_and_accumulate(
                                &loss_slice,
                                &grad_slice,
                                1e6_f32,
                                guard_collapse_thresh,
                                !guard_past_warmup,
                            ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?;
                            if gr.halt_nan {
                                return Err(anyhow::anyhow!(
                                    "NaN/Inf at step {}: loss={}, grad={}",
                                    train_step_count, gr.raw_loss, gr.raw_grad_norm
                                ));
                            }
                            if gr.halt_grad_collapse {
                                agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| {
                                    tracing::info!("Early stopping (gradient collapse): {}", e);
                                    anyhow::anyhow!("Early stopping: {}", e)
                                })?;
                            }
                        }
                        guard_total_us += guard_start.elapsed().as_micros() as u64;
                        train_step_count += 1;
                        self.gradient_logging_step += 1;
                    }
                } else {
                    // Gradient accumulation path (gradient_accumulation_steps > 1).
                    //
                    // GpuTensor elementwise kernels (add, broadcast_mul, cat) cache
                    // CudaFunction handles on the shared CUDA primary context, causing
                    // sequential test failures. We replace them with:
                    //   - CPU-side gradient accumulation (download, add, re-upload via
                    //     GpuTensor::new — no elementwise kernel invoked)
                    //   - CPU-side gradient scaling (multiply in place on Vec<f32>)
                    //   - CPU-side PER priority updates via update_priorities() instead of
                    //     GpuTensor::cat + update_priorities_gpu()
                    let mut batch_iter = batches.into_iter().peekable();

                    while batch_iter.peek().is_some() {
                        // CPU-side gradient accumulator: name → flat Vec<f32> + shape
                        let mut cpu_acc_grads: Option<std::collections::BTreeMap<String, (Vec<f32>, Vec<usize>)>> = None;
                        // CPU-side PER priority accumulators
                        let mut group_td_errors: Vec<f32> = Vec::new();
                        let mut group_indices: Vec<usize> = Vec::new();
                        let mut accum_count: usize = 0;

                        for _ in 0..accum_steps {
                            let batch = match batch_iter.next() {
                                Some(b) => b,
                                None => break,
                            };

                            let result = match agent.compute_gradients(batch) {
                                Ok(r) => r,
                                Err(e) => {
                                    let msg = e.to_string();
                                    if msg.contains("Early stopping") || msg.contains("Gradient collapse") {
                                        return Err(anyhow::anyhow!("{}", msg));
                                    }
                                    return Err(anyhow::anyhow!("GPU gradient compute FAILED: {e}"));
                                }
                            };

                            // Extract CudaSlice handles BEFORE destructuring result.
                            let loss_sl_opt = result.loss_cuda_slice();
                            let grad_sl_opt = result.grad_norm_cuda_slice();
                            let r_td_errors = result.td_errors;
                            let r_indices = result.indices;
                            let r_grads = result.grads;

                            // Accumulate gradients on CPU (download → add → store).
                            // No elementwise CUDA kernels — avoids CudaFunction handle caching.
                            let stream_ref = self.cuda_stream.as_ref()
                                .ok_or_else(|| anyhow::anyhow!("CUDA stream needed for grad accum"))?;
                            match &mut cpu_acc_grads {
                                Some(ref mut acc) => {
                                    for (name, grad) in &r_grads {
                                        let n = grad.shape().iter().product::<usize>();
                                        let mut host = vec![0.0_f32; n];
                                        stream_ref.memcpy_dtoh(grad.cuda_data(), &mut host)
                                            .map_err(|e| anyhow::anyhow!("grad DtoH '{name}': {e}"))?;
                                        if let Some(entry) = acc.get_mut(name) {
                                            for (a, b) in entry.0.iter_mut().zip(host.iter()) {
                                                *a += b;
                                            }
                                        } else {
                                            acc.insert(name.clone(), (host, grad.shape().to_vec()));
                                        }
                                    }
                                }
                                None => {
                                    let mut map = std::collections::BTreeMap::new();
                                    for (name, grad) in &r_grads {
                                        let n = grad.shape().iter().product::<usize>();
                                        let mut host = vec![0.0_f32; n];
                                        stream_ref.memcpy_dtoh(grad.cuda_data(), &mut host)
                                            .map_err(|e| anyhow::anyhow!("grad DtoH '{name}' (init): {e}"))?;
                                        map.insert(name.clone(), (host, grad.shape().to_vec()));
                                    }
                                    cpu_acc_grads = Some(map);
                                }
                            }

                            // Guard checks use CudaSlice directly — no GpuTensor elementwise ops.
                            if let (Some(loss_sl_r), Some(grad_sl_r)) =
                                (loss_sl_opt, grad_sl_opt)
                            {
                                if let Some(ref mut guard) = self.training_guard {
                                    let loss_sl = loss_sl_r
                                        .map_err(|e| anyhow::anyhow!("accum loss CudaSlice: {e}"))?;
                                    let grad_sl = grad_sl_r
                                        .map_err(|e| anyhow::anyhow!("accum grad CudaSlice: {e}"))?;
                                    let gr = guard.check_and_accumulate(
                                        &loss_sl, &grad_sl, 1e6_f32,
                                        guard_collapse_thresh, !guard_past_warmup,
                                    ).map_err(|e| anyhow::anyhow!("guard accum step: {e}"))?;
                                    if gr.halt_nan {
                                        return Err(anyhow::anyhow!(
                                            "NaN/Inf in accum sub-step: loss={}, grad={}",
                                            gr.raw_loss, gr.raw_grad_norm
                                        ));
                                    }
                                    if gr.halt_grad_collapse {
                                        agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| {
                                            tracing::info!("Early stopping (gradient collapse): {}", e);
                                            anyhow::anyhow!("Early stopping: {}", e)
                                        })?;
                                    }
                                }
                            }

                            // Collect CPU-side PER priorities (avoid GpuTensor::cat).
                            group_td_errors.extend_from_slice(&r_td_errors);
                            group_indices.extend_from_slice(&r_indices);

                            accum_count += 1;
                        }

                        if accum_count == 0 { break; }

                        // Scale accumulated gradients by 1/N on CPU, re-upload as GpuTensor.
                        // GpuTensor::new() wraps an existing CudaSlice — no elementwise kernel.
                        if let Some(cpu_grads) = cpu_acc_grads {
                            let stream_ref = self.cuda_stream.as_ref()
                                .ok_or_else(|| anyhow::anyhow!("CUDA stream needed for grad scale"))?;
                            let scale = 1.0 / accum_count as f32;
                            let mut gpu_grads = std::collections::BTreeMap::new();
                            for (name, (mut host, shape)) in cpu_grads {
                                for v in host.iter_mut() { *v *= scale; }
                                let gpu_buf = stream_ref.clone_htod(&host)
                                    .map_err(|e| anyhow::anyhow!("grad HtoD '{name}': {e}"))?;
                                let tensor = ml_core::cuda_autograd::GpuTensor::new(gpu_buf, shape)
                                    .map_err(|e| anyhow::anyhow!("GpuTensor::new '{name}': {e}"))?;
                                gpu_grads.insert(name, tensor);
                            }
                            agent.apply_accumulated_gradients(&gpu_grads)
                                .map_err(|e| anyhow::anyhow!("apply grads: {e}"))?;
                        }

                        // Update PER priorities via CPU path — avoids GpuTensor::cat.
                        if !group_indices.is_empty() {
                            agent.update_priorities(&group_indices, &group_td_errors)
                                .map_err(|e| anyhow::anyhow!("PER update: {e}"))?;
                        }
                        agent.step_replay_buffer();

                        train_step_count += accum_count;
                        self.gradient_logging_step += accum_count;
                    }
                }
            }
        }

        if train_step_count > 0 {
            info!(
                "Training step breakdown ({} steps): sample={:.0}ms fused={:.0}ms guard={:.0}ms (per-step: sample={:.1}ms fused={:.1}ms guard={:.1}ms)",
                train_step_count,
                sample_total_us as f64 / 1000.0,
                fused_total_us as f64 / 1000.0,
                guard_total_us as f64 / 1000.0,
                sample_total_us as f64 / 1000.0 / train_step_count as f64,
                fused_total_us as f64 / 1000.0 / train_step_count as f64,
                guard_total_us as f64 / 1000.0 / train_step_count as f64,
            );
        }
        Ok(train_step_count)
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Epoch boundary readback + safety checks
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn process_epoch_boundary(
        &mut self,
        epoch: usize,
        train_step_count: usize,
        monitor: &mut TrainingMonitor,
    ) -> Result<EpochBoundaryMetrics> {
        let (avg_loss, avg_grad) = if let Some(ref mut guard) = self.training_guard {
            let (al, ag) = guard.read_accumulators()
                .map_err(|e| anyhow::anyhow!("guard read_accumulators: {e}"))?;
            (al as f32, ag as f32)
        } else {
            (0.0_f32, 0.0_f32)
        };

        let avg_q = {
            let agent_clone = Arc::clone(&self.agent);
            let mut agent = agent_clone.write().await;
            self.estimate_avg_q_value_with_early_stopping(&mut agent).await?
        };

        // Epoch-level safety checks
        if !avg_loss.is_finite() || !avg_grad.is_finite() {
            let msg = format!(
                "SAFETY: NaN/Inf in epoch {} avg — loss={:.6}, grad={:.6}",
                epoch, avg_loss, avg_grad
            );
            match self.safety_level {
                crate::safety::SafetyLevel::Strict => {
                    return Err(anyhow::anyhow!("{} (stopping training)", msg));
                },
                crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => {
                    debug!("{} (continuing training)", msg);
                },
            }
        }

        // Loss history (epoch-level average)
        self.safety_loss_history.push_back(avg_loss);
        if self.safety_loss_history.len() > 30 {
            self.safety_loss_history.pop_front();
        }

        let n = train_step_count as f64;

        // Metrics aggregator
        let current_avg_reward = if !monitor.reward_history.is_empty() {
            monitor.reward_history.iter().sum::<f32>() / monitor.reward_history.len() as f32
        } else {
            0.0
        };
        self.metrics_aggregator.record(avg_loss, avg_q as f32, current_avg_reward);
        self.metrics_aggregator.record_gradient(avg_grad);

        if self.metrics_aggregator.should_log(&self.logging_config) {
            let aggregated = self.metrics_aggregator.aggregate_and_clear();
            log_training_progress(&aggregated);
        }

        // Gradient diagnostics at epoch boundary
        {
            let mut agent = self.agent.write().await;
            agent.log_diagnostics(avg_grad)
                .map_err(|e| {
                    tracing::info!("Early stopping (gradient collapse): {}", e);
                    anyhow::anyhow!("Early stopping: {}", e)
                })?;
        }

        monitor.track_q_value_range(avg_q);

        // Return raw accumulated totals for epoch metric computation
        // The caller computes per-step averages from epoch_loss/n, etc.
        Ok(EpochBoundaryMetrics {
            avg_loss: (avg_loss as f64 * n) as f32,
            avg_grad: (avg_grad as f64 * n) as f32,
            avg_q: avg_q * n,
        })
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Sync GPU weight copies after training updates
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn sync_gpu_weights(&mut self) -> Result<()> {
        let Some(ref mut collector) = self.gpu_experience_collector else {
            return Ok(());
        };

        let agent = self.agent.read().await;
        let dqn_ref: Option<&crate::dqn::DQN> = match &*agent {
            DQNAgentType::Standard(ref dqn) => Some(dqn),
            DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()),
        };

        if let Some(dqn) = dqn_ref {
            // Sync online weights: branching > plain dueling > hybrid
            let online_synced = if let Some(ref bn) = dqn.branching_q_network {
                collector.sync_online_weights(bn.vars()).is_ok()
            } else if let Some(ref online) = dqn.dueling_q_network {
                collector.sync_online_weights(online.store()).is_ok()
            } else if let Some(ref online) = dqn.dist_dueling_q_network {
                collector.sync_online_weights(online.vars()).is_ok()
            } else {
                false
            };
            if !online_synced {
                return Err(anyhow::anyhow!("GPU online weight sync FAILED -- stale Q-values would corrupt training"));
            }

            // Sync target weights: branching > plain dueling > hybrid
            let target_synced = if let Some(ref bn) = dqn.branching_target_network {
                collector.sync_target_weights(bn.vars()).is_ok()
            } else if let Some(ref target) = dqn.dueling_target_network {
                collector.sync_target_weights(target.store()).is_ok()
            } else if let Some(ref target) = dqn.dist_dueling_target_network {
                collector.sync_target_weights(target.vars()).is_ok()
            } else {
                false
            };
            if !target_synced {
                return Err(anyhow::anyhow!("GPU target weight sync FAILED -- stale target Q-values would corrupt training"));
            }

            // Sync branching DQN extra heads (order + urgency, branches 1+2)
            if let Some(ref bn) = dqn.branching_q_network {
                collector.sync_online_branching(bn.vars())
                    .map_err(|e| anyhow::anyhow!("GPU online branching weight sync FAILED: {e}"))?;
            }
            if let Some(ref bn) = dqn.branching_target_network {
                collector.sync_target_branching(bn.vars())
                    .map_err(|e| anyhow::anyhow!("GPU target branching weight sync FAILED: {e}"))?;
            }

            // Sync RMSNorm weights for distributional dueling networks (D6)
            if let Some(ref online) = dqn.dist_dueling_q_network {
                collector.sync_online_rmsnorm(online.vars())
                    .map_err(|e| anyhow::anyhow!("GPU online RMSNorm weight sync FAILED: {e}"))?;
            }
            if let Some(ref target) = dqn.dist_dueling_target_network {
                collector.sync_target_rmsnorm(target.vars())
                    .map_err(|e| anyhow::anyhow!("GPU target RMSNorm weight sync FAILED: {e}"))?;
            }
        }
        drop(agent);

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: M2 — Refresh stale PER priorities
    // ═══════════════════════════════════════════════════════════════════════

    /// Refresh stale PER priorities using cuBLAS forward — zero Candle.
    ///
    /// Samples stale experiences, runs cuBLAS forward via fused_ctx to compute
    /// max Q-values (TD error proxy), then updates priorities via GPU PER.
    pub(crate) async fn refresh_stale_per_priorities(&mut self, epoch: usize) -> Result<()> {
        if !true || self.hyperparams.n_steps <= 1 {
            return Ok(());
        }

        let agent = self.agent.read().await;
        let buffer = agent.memory();
        let buf_len = buffer.len();
        let refresh_limit = (buf_len / 20).max(32).min(256);
        let max_age = 500;
        let stale_indices = buffer.get_stale_indices(max_age, refresh_limit);

        if stale_indices.is_empty() {
            return Ok(());
        }

        let experiences = buffer.get_experiences_at(&stale_indices);

        let mut flat_states: Vec<f32> = Vec::with_capacity(stale_indices.len() * 64);
        let mut valid_indices = Vec::with_capacity(stale_indices.len());
        let mut state_dim: Option<usize> = None;

        for (idx, exp_opt) in stale_indices.iter().zip(experiences.iter()) {
            if let Some(Some(exp)) = exp_opt.as_ref().map(Some) {
                if let Some(sd) = state_dim {
                    if exp.state.len() != sd {
                        continue;
                    }
                } else {
                    state_dim = Some(exp.state.len());
                }
                flat_states.extend_from_slice(&exp.state);
                valid_indices.push(*idx);
            }
        }

        if let Some(_sd) = state_dim {
            if !valid_indices.is_empty() {
                let refresh_stream = match self.cuda_stream.as_ref() {
                    Some(s) => s,
                    None => return Ok(()),
                };

                // Cap at fused_ctx batch_size to avoid exceeding cuBLAS buffer
                let fused = match self.fused_ctx.as_mut() {
                    Some(f) => f,
                    None => return Ok(()),
                };
                let fused_batch = fused.batch_size();
                let batch_size_refresh = valid_indices.len().min(fused_batch);
                // Truncate flat_states if needed
                let state_dim_val = flat_states.len() / valid_indices.len();
                let truncated_states = &flat_states[..batch_size_refresh * state_dim_val];

                // Upload states to GPU as CudaSlice<f32>
                let mut states_gpu = refresh_stream.alloc_zeros::<f32>(batch_size_refresh * state_dim_val)
                    .map_err(|e| anyhow::anyhow!("PER refresh alloc: {e}"))?;
                refresh_stream.memcpy_htod(truncated_states, &mut states_gpu)
                    .map_err(|e| anyhow::anyhow!("PER refresh HtoD: {e}"))?;

                // cuBLAS forward → Q-values in q_out_buf
                let total_actions = fused.total_actions();
                let q_result = fused.compute_q_values(&states_gpu, batch_size_refresh);

                if let Ok(q_out) = q_result {
                    // Download Q-values (q_out may be larger than needed — match GPU buffer size)
                    let mut host_q = vec![0.0_f32; q_out.len()];
                    if refresh_stream.memcpy_dtoh(q_out, &mut host_q).is_ok() {
                        host_q.truncate(batch_size_refresh * total_actions);
                        let mut max_vals = Vec::with_capacity(batch_size_refresh);
                        for row in 0..batch_size_refresh {
                            let offset = row * total_actions;
                            let mut max_v = 0.01_f32; // floor
                            for col in 0..total_actions {
                                let v = host_q.get(offset + col).copied().unwrap_or(0.0).abs();
                                if v > max_v && v.is_finite() { max_v = v; }
                            }
                            max_vals.push(max_v);
                        }

                        // Upload TD errors and indices as raw CudaSlice (zero GpuTensor)
                        if let Ok(td_slice) = refresh_stream.clone_htod(&max_vals) {
                            let idx_u32: Vec<u32> = valid_indices.iter()
                                .take(batch_size_refresh)
                                .map(|&i| i as u32)
                                .collect();
                            if let Ok(idx_slice) = refresh_stream.clone_htod(&idx_u32) {
                                let mut gpu_buf = agent.memory().as_gpu_buffer();
                                if let Some(ref mut buf) = gpu_buf {
                                    if let Err(e) = buf.gpu.update_priorities_gpu(
                                        &idx_slice, &td_slice, batch_size_refresh,
                                    ) {
                                        debug!("M2: GPU priority refresh failed (non-fatal): {}", e);
                                    } else {
                                        debug!(
                                            "Epoch {}: Refreshed {} stale PER priorities (max_age={}, cuBLAS)",
                                            epoch + 1, batch_size_refresh, max_age
                                        );
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Epoch metrics logging + financials + Prometheus + QuestDB
    // ═══════════════════════════════════════════════════════════════════════

    #[allow(clippy::too_many_arguments)]
    pub(crate) async fn log_epoch_metrics_and_financials(
        &mut self,
        epoch: usize,
        train_step_count: usize,
        boundary: &Option<EpochBoundaryMetrics>,
        monitor: &mut TrainingMonitor,
        epoch_duration: std::time::Duration,
        total_action_counts: &mut [usize; 9],
        total_factored_action_counts: &mut [usize; 81],
    ) -> Result<EpochLogOutput> {
        // Calculate epoch metrics (average over training steps)
        let (epoch_loss, epoch_q_value, epoch_gradient_norm) = match boundary {
            Some(b) => (b.avg_loss as f64, b.avg_q, b.avg_grad as f64),
            None => (0.0, 0.0, 0.0),
        };

        let (avg_loss, avg_q_value, avg_grad_norm) = if train_step_count > 0 {
            (
                epoch_loss / train_step_count as f64,
                epoch_q_value / train_step_count as f64,
                epoch_gradient_norm / train_step_count as f64,
            )
        } else {
            (0.0, 0.0, 0.0)
        };

        // WAVE 30: Log epoch end
        use crate::dqn::logging::AggregatedMetrics;
        let epoch_avg_reward = if !monitor.reward_history.is_empty() {
            monitor.reward_history.iter().sum::<f32>() / monitor.reward_history.len() as f32
        } else {
            0.0
        };
        let epoch_metrics = AggregatedMetrics {
            mean_loss: avg_loss as f32,
            std_loss: 0.0,
            mean_q_value: avg_q_value as f32,
            std_q_value: 0.0,
            mean_reward: epoch_avg_reward,
            mean_gradient_norm: avg_grad_norm as f32,
            batch_count: train_step_count,
        };
        log_epoch_end(epoch + 1, &epoch_metrics, epoch_duration.as_secs_f64());

        // Download GPU monitoring summary
        if let Some(ref mon) = self.gpu_monitoring {
            if let Ok(summary) = mon.download_summary() {
                if summary.total_experiences > 0 {
                    info!(
                        "GPU epoch summary: mean_reward={:.6}, std={:.6}, sharpe={:.3}, actions={:?}",
                        summary.mean_reward, summary.reward_std, summary.sharpe_estimate,
                        summary.action_counts
                    );
                    monitor.track_reward(summary.mean_reward);
                    for (idx, &count) in summary.action_counts.iter().enumerate() {
                        if count > 0 {
                            if let Ok(exp_level) = crate::dqn::action_space::ExposureLevel::from_index(idx) {
                                let factored = self.route_action(exp_level, self.hyperparams.avg_spread as f32);
                                monitor.track_action(&factored);
                            } else {
                                monitor.track_action_by_exposure(idx);
                            }
                        }
                    }
                }
            }
        }

        // Collect GPU trade stats (cold path, once per epoch)
        let trade_stats = if let Some(ref mut collector) = self.gpu_experience_collector {
            match collector.collect_trade_stats() {
                Ok(ts) => {
                    info!(
                        "GPU trades={} wins={} losses={} PF={:.3}",
                        ts.total_trades,
                        ts.winning_trades,
                        ts.losing_trades,
                        if ts.sum_losses > 1e-10 { ts.sum_wins / ts.sum_losses } else { 0.0 },
                    );
                    ts
                }
                Err(e) => {
                    warn!("collect_trade_stats failed (using defaults): {e}");
                    crate::cuda_pipeline::gpu_experience_collector::TradeStats::default()
                }
            }
        } else {
            crate::cuda_pipeline::gpu_experience_collector::TradeStats::default()
        };
        self.trade_stats_history.push(trade_stats.clone());

        // Reward statistics every 10 epochs
        if (epoch + 1) % 10 == 0 && !monitor.reward_history.is_empty() {
            let rewards = &monitor.reward_history;
            let reward_mean = rewards.iter().sum::<f32>() / rewards.len() as f32;
            let reward_variance = rewards.iter()
                .map(|r| (r - reward_mean).powi(2))
                .sum::<f32>() / rewards.len() as f32;
            let reward_std = reward_variance.sqrt();
            let reward_min = rewards.iter().copied().fold(f32::INFINITY, f32::min);
            let reward_max = rewards.iter().copied().fold(f32::NEG_INFINITY, f32::max);
            let non_zero_count = rewards.iter().filter(|&&r| r.abs() > 1e-9).count();
            let non_zero_pct = (non_zero_count as f32 / rewards.len() as f32) * 100.0;

            info!(
                "REWARD_STATS: epoch={}, mean={:.6}, std={:.6}, min={:.6}, max={:.6}, non_zero={}/{} ({:.1}%)",
                epoch + 1, reward_mean, reward_std, reward_min, reward_max,
                non_zero_count, rewards.len(), non_zero_pct
            );
        }

        // Safety plateau detection
        if self.safety_loss_history.len() >= 10 {
            let recent_losses: Vec<f32> = self.safety_loss_history
                .iter().rev().take(10).copied().collect();

            let mean: f32 = recent_losses.iter().sum::<f32>() / 10.0;
            let variance: f32 = recent_losses.iter()
                .map(|&x| (x - mean).powi(2))
                .sum::<f32>() / 10.0;
            let std_dev = variance.sqrt();

            if std_dev < mean * 0.01 && mean > 1e-6 {
                self.safety_loss_plateau_counter += 1;
                if self.safety_loss_plateau_counter >= 10 {
                    debug!(
                        "SAFETY: Training stuck for {} epochs (loss variance: {:.6}, mean: {:.6})",
                        self.safety_loss_plateau_counter, std_dev, mean
                    );
                }
            } else {
                self.safety_loss_plateau_counter = 0;
            }
        }

        // Accumulate action counts
        for (i, count) in monitor.action_counts.iter().enumerate() {
            total_action_counts[i] += count;
        }
        for (i, count) in monitor.factored_action_counts.iter().enumerate() {
            total_factored_action_counts[i] += count;
        }

        // Monitoring validation
        if let Err(e) = monitor.validate_all() {
            return Err(e);
        }

        let current_epsilon = self.get_epsilon().await?;
        let (q_min, q_max, q_mean) = monitor.get_q_value_stats();

        info!(
            "Epoch {}/{}: train_loss={:.6}, Q-value={:.4}, grad_norm={:.6}, train_steps={}, epsilon={:.4}, duration={:.2}s",
            epoch + 1, self.hyperparams.epochs,
            avg_loss, avg_q_value, avg_grad_norm, train_step_count,
            current_epsilon, epoch_duration.as_secs_f64()
        );

        // Epsilon decay (skip when noisy nets)
        if !true {
            let mut agent = self.agent.write().await;
            agent.update_epsilon();
        }

        // Noisy sigma schedule
        if let Some(ref mut scheduler) = self.noisy_sigma_scheduler {
            for _ in 0..train_step_count {
                scheduler.step();
            }
            let sigma_scale = scheduler.get_sigma();
            let mut agent = self.agent.write().await;
            agent.set_noise_sigma_scale(sigma_scale);
        }

        // Reset count-based exploration bonus
        {
            let mut agent = self.agent.write().await;
            agent.reset_count_bonus();
        }

        // M2: Refresh stale PER priorities
        self.refresh_stale_per_priorities(epoch).await?;

        // LR scheduler
        self.lr_scheduler.step();
        let current_lr = self.lr_scheduler.get_lr();
        if epoch % 10 == 0 {
            info!(
                "Learning rate scheduled update: epoch={}, lr={:.2e} (initial={:.2e})",
                epoch + 1, current_lr, self.lr_scheduler.get_initial_lr()
            );
        }
        let initial_lr = self.lr_scheduler.get_initial_lr();
        if initial_lr > 0.0 {
            let decay_factor = current_lr / initial_lr;
            self.agent.write().await.update_learning_rate(decay_factor)?;
        }
        training_metrics::set_learning_rate("dqn", "current", current_lr);

        // Q-value diagnostics
        if train_step_count > 0 {
            info!(
                "Epoch {}/{}: Q-value range=[{:.2}, {:.2}], mean={:.2}",
                epoch + 1, self.hyperparams.epochs, q_min, q_max, q_mean
            );

            const Q_VALUE_WARNING_THRESHOLD: f64 = 500_000.0;
            if q_max > Q_VALUE_WARNING_THRESHOLD {
                warn!(
                    "Q-value explosion detected at epoch {}: max Q-value {:.2e} exceeds threshold {:.2e}",
                    epoch + 1, q_max, Q_VALUE_WARNING_THRESHOLD
                );
                warn!("Consider:");
                warn!("  Reducing learning rate (current: {:.2e})", self.hyperparams.learning_rate);
                warn!("  Enabling target network soft updates (Polyak averaging, tau=0.005)");
                warn!("  Adjusting reward scaling");
                training_metrics::record_gradient_explosion("dqn", "current");
            }

            // M3: Combined Q-value diagnostics
            if let Some(((gap_mean, gap_min, gap_max), per_action_avgs)) =
                self.compute_epoch_q_diagnostics().await
            {
                info!(
                    "Epoch {}/{}: Q-value gap (best-2nd): mean={:.4}, min={:.4}, max={:.4}",
                    epoch + 1, self.hyperparams.epochs,
                    gap_mean, gap_min, gap_max
                );

                let names = ["S100", "S75", "S50", "S25", "Flat", "L25", "L50", "L75", "L100"];
                let parts: Vec<String> = names.iter()
                    .zip(per_action_avgs.iter())
                    .map(|(n, q)| format!("{}={:.4}", n, q))
                    .collect();
                info!(
                    "Epoch {}/{}: Per-action Q: {}",
                    epoch + 1, self.hyperparams.epochs, parts.join(", ")
                );
            }
        }

        // Validation loss
        let val_loss = self.compute_validation_loss().await?;
        info!(
            "Epoch {}/{}: val_loss={} (backtest Sharpe proxy)",
            epoch + 1, self.hyperparams.epochs,
            if val_loss.abs() < 1e-10 { "N/A".to_owned() } else { format!("{val_loss:.6}") }
        );

        // Prometheus metrics
        training_metrics::set_epoch("dqn", "current", (epoch + 1) as f64);
        training_metrics::set_epoch_loss("dqn", "current", avg_loss);
        training_metrics::set_validation_loss("dqn", "current", val_loss);
        if epoch_duration.as_secs_f64() > 0.0 {
            training_metrics::set_batches_per_second(
                "dqn", "current",
                train_step_count as f64 / epoch_duration.as_secs_f64(),
            );
        }
        training_metrics::set_q_value_stats("dqn", "current", q_mean, q_max);
        training_metrics::set_gradient_norm("dqn", "current", avg_grad_norm);
        training_metrics::set_epoch_duration("dqn", "current", epoch_duration.as_secs_f64());
        {
            let agent = self.agent.read().await;
            if let Ok(buf_size) = agent.get_replay_buffer_size() {
                training_metrics::set_replay_buffer_size("dqn", "current", buf_size as f64);
            }
        }

        // Adaptive tau
        if epoch > 0 {
            let q_mean_growth = q_mean - self.prev_epoch_q_mean;
            if q_mean_growth > 0.005 {
                self.adaptive_tau = (self.adaptive_tau * 2.0).min(0.05);
                warn!(
                    "Q-value drift detected (delta={:.3}), increasing tau to {:.4}",
                    q_mean_growth, self.adaptive_tau
                );
            } else if q_mean_growth < 0.001 {
                self.adaptive_tau = (self.adaptive_tau * 0.9).max(self.hyperparams.tau);
            } else {
                // Q-value growth in [0.001, 0.005] — tau unchanged
            }
            {
                let mut agent = self.agent.write().await;
                match &mut *agent {
                    DQNAgentType::Standard(dqn) => dqn.config.tau = self.adaptive_tau,
                    DQNAgentType::RegimeConditional(regime) => {
                        regime.primary_head_mut().config.tau = self.adaptive_tau;
                    }
                }
            }
        }
        self.prev_epoch_q_mean = q_mean;

        // Action diversity
        let epoch_total_factored: usize = monitor.factored_action_counts.iter().sum();
        let epoch_total_exposure: usize = monitor.action_counts.iter().sum();
        let (active_actions_count, action_space_size, diversity_percentage) =
            if epoch_total_factored > 0 {
                let active_threshold = (epoch_total_factored as f64 * 0.005).max(1.0);
                let active = monitor.factored_action_counts.iter()
                    .filter(|&&count| count as f64 >= active_threshold)
                    .count();
                (active, 45_usize, (active as f64 / 45.0) * 100.0)
            } else {
                let active_threshold = (epoch_total_exposure as f64 * 0.005).max(1.0);
                let active = monitor.action_counts.iter()
                    .filter(|&&count| count as f64 >= active_threshold)
                    .count();
                (active, 5_usize, (active as f64 / 5.0) * 100.0)
            };

        info!(
            "Epoch {}/{}: Action diversity={}/{} ({:.1}%)",
            epoch + 1, self.hyperparams.epochs,
            active_actions_count, action_space_size, diversity_percentage
        );

        // Epoch entropy
        let (epoch_entropy, entropy_total, entropy_size) = if epoch_total_factored > 0 {
            let mut entropy_raw = 0.0_f64;
            for &count in monitor.factored_action_counts.iter() {
                if count > 0 {
                    let p = count as f64 / epoch_total_factored as f64;
                    entropy_raw -= p * p.ln();
                }
            }
            (entropy_raw / 81.0_f64.ln(), epoch_total_factored, 81_usize)
        } else if epoch_total_exposure > 0 {
            let mut entropy_raw = 0.0_f64;
            for &count in monitor.action_counts.iter() {
                if count > 0 {
                    let p = count as f64 / epoch_total_exposure as f64;
                    entropy_raw -= p * p.ln();
                }
            }
            (entropy_raw / 5.0_f64.ln(), epoch_total_exposure, 5_usize)
        } else {
            (0.0, 0, 5)
        };
        let _ = (entropy_total, entropy_size);

        info!(
            "  Exploration: entropy={:.3} (1.0=uniform), epsilon={:.4}, noisy_nets={}, count_bonus={}",
            epoch_entropy,
            self.get_epsilon().await.unwrap_or(0.0),
            true,
            self.hyperparams.count_bonus_coefficient.unwrap_or(0.0) > 0.0,
        );

        training_metrics::set_action_entropy("dqn", "current", epoch_entropy);
        training_metrics::set_action_diversity("dqn", "current", diversity_percentage / 100.0);

        const DIVERSITY_THRESHOLD: usize = 2;
        if active_actions_count < DIVERSITY_THRESHOLD {
            warn!(
                "LOW ACTION DIVERSITY: {}/5 exposure levels (<40%), entropy={:.3}",
                active_actions_count, epoch_entropy,
            );
        }

        // Episode stats
        let (mean_len, std_len, min_len, max_len, exit_counts) = monitor.get_episode_stats();
        if !monitor.episode_lengths.is_empty() {
            let total_episodes = monitor.episode_lengths.len();

            info!(
                "WAVE P2 Episode Stats [Epoch {}]: {} episodes, length: mean={:.1}+/-{:.1}, min={}, max={}",
                epoch + 1, total_episodes, mean_len, std_len, min_len, max_len
            );

            info!(
                "  Exit breakdown: profit={}({:.1}%), stop={}({:.1}%), time={}({:.1}%), boundary={}({:.1}%)",
                exit_counts[0], (exit_counts[0] as f64 / total_episodes as f64) * 100.0,
                exit_counts[1], (exit_counts[1] as f64 / total_episodes as f64) * 100.0,
                exit_counts[2], (exit_counts[2] as f64 / total_episodes as f64) * 100.0,
                exit_counts[3], (exit_counts[3] as f64 / total_episodes as f64) * 100.0,
            );
        }

        // VaR/CVaR from GPU trade stats step_returns
        if trade_stats.step_returns.len() > 20 {
            let returns = &trade_stats.step_returns;

            let (var_95, cvar_95) = calculate_var_cvar(returns, 0.05);
            let (var_99, cvar_99) = calculate_var_cvar(returns, 0.01);

            info!(
                "Epoch {}/{}: Risk Metrics - VaR(95%)={:.4}%, CVaR(95%)={:.4}%, VaR(99%)={:.4}%, CVaR(99%)={:.4}% (from {} step returns)",
                epoch + 1, self.hyperparams.epochs,
                var_95 * 100.0, cvar_95 * 100.0,
                var_99 * 100.0, cvar_99 * 100.0,
                returns.len()
            );
        }

        // Financial metrics from GPU trade stats
        let epoch_sharpe = {
            let financials = compute_epoch_financials(
                &trade_stats,
                &monitor.action_counts,
                100_000.0,
                self.hyperparams.bars_per_day,
            );
            training_metrics::set_epoch_financial_metrics(
                "dqn", "current",
                financials.sharpe, financials.sortino, financials.win_rate,
                financials.max_drawdown, financials.profit_factor,
                financials.total_return, financials.avg_return,
                financials.total_trades as f64,
            );
            training_metrics::set_epoch_action_distribution(
                "dqn", "current",
                financials.buy_pct, financials.sell_pct, financials.hold_pct,
            );
            info!(
                "Epoch {}/{}: Sharpe={:.2} WinRate={:.1}% MaxDD={:.3}% PF={:.2} Return={:+.2}% Trades={}",
                epoch + 1, self.hyperparams.epochs,
                financials.sharpe, financials.win_rate * 100.0,
                financials.max_drawdown * 100.0, financials.profit_factor,
                financials.total_return * 100.0, financials.total_trades,
            );
            // QuestDB epoch row
            {
                let bps = if epoch_duration.as_secs_f64() > 0.0 {
                    train_step_count as f64 / epoch_duration.as_secs_f64()
                } else {
                    0.0
                };
                let buf_size = {
                    let agent = self.agent.read().await;
                    agent.get_replay_buffer_size().unwrap_or(0) as f64
                };
                let run_id = std::env::var("HOSTNAME").unwrap_or_else(|_| "local".to_owned());
                questdb_sink::record_training_epoch(&questdb_sink::EpochRecord {
                    model: "dqn",
                    fold: "current",
                    run_id: &run_id,
                    symbol: "",
                    epoch: (epoch + 1) as u32,
                    loss: avg_loss,
                    val_loss,
                    batches_per_second: bps,
                    epoch_duration_secs: epoch_duration.as_secs_f64(),
                    q_mean,
                    q_max,
                    gradient_norm: avg_grad_norm,
                    learning_rate: current_lr,
                    replay_buffer_size: buf_size,
                    action_entropy: epoch_entropy,
                    action_diversity: diversity_percentage / 100.0,
                    sharpe: financials.sharpe,
                    sortino: financials.sortino,
                    win_rate: financials.win_rate,
                    max_drawdown: financials.max_drawdown,
                    profit_factor: financials.profit_factor,
                    total_return: financials.total_return,
                    avg_return: financials.avg_return,
                    total_trades: financials.total_trades as f64,
                });
            }

            financials.sharpe
        };

        // Track metrics for early stopping
        self.loss_history.push(avg_loss);
        self.q_value_history.push(avg_q_value);
        self.val_loss_history.push(val_loss);
        self.sharpe_history.push(epoch_sharpe);

        // Limit history vectors to prevent unbounded growth
        const MAX_HISTORY_LEN: usize = 100;
        if self.loss_history.len() > MAX_HISTORY_LEN {
            self.loss_history.drain(0..50);
        }
        if self.q_value_history.len() > MAX_HISTORY_LEN {
            self.q_value_history.drain(0..50);
        }
        if self.val_loss_history.len() > MAX_HISTORY_LEN {
            self.val_loss_history.drain(0..50);
        }
        if self.sharpe_history.len() > MAX_HISTORY_LEN {
            self.sharpe_history.drain(0..50);
        }

        Ok(EpochLogOutput {
            avg_loss,
            avg_q_value,
            avg_grad_norm,
            epoch_sharpe,
            val_loss,
            q_min,
            q_max,
            q_mean,
        })
    }

    // ═══════════════════════════════════════════════════════════════════════
    // Helper: Checkpoint saving + early stopping checks
    // Returns Err for early stopping (caller propagates), Ok(()) to continue
    // ═══════════════════════════════════════════════════════════════════════

    pub(crate) async fn handle_epoch_checkpoints_and_early_stopping<F>(
        &mut self,
        epoch: usize,
        train_step_count: usize,
        log_output: &EpochLogOutput,
        checkpoint_callback: &mut F,
    ) -> Result<()>
    where
        F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
    {
        // C4 FIX: Save best model checkpoint when Sharpe improves
        if train_step_count > 0 && log_output.epoch_sharpe > self.best_sharpe {
            self.best_sharpe = log_output.epoch_sharpe;
            self.best_val_loss = log_output.val_loss;
            self.best_epoch = epoch + 1;

            info!(
                "New best Sharpe: {:.4} at epoch {} (val_loss={:.6})",
                log_output.epoch_sharpe, epoch + 1, log_output.val_loss,
            );

            let checkpoint_data = self.serialize_model().await?;

            // Verify checkpoint integrity
            if self.safety_level != crate::safety::SafetyLevel::Permissive {
                if checkpoint_data.is_empty() {
                    let msg = "SAFETY: Checkpoint verification failed - empty checkpoint data";
                    match self.safety_level {
                        crate::safety::SafetyLevel::Strict => {
                            return Err(anyhow::anyhow!("{}", msg));
                        },
                        crate::safety::SafetyLevel::Normal => {
                            debug!("{} (continuing anyway)", msg);
                        },
                        _ => {},
                    }
                } else {
                    debug!("SAFETY: Checkpoint verification passed ({} bytes)", checkpoint_data.len());
                }
            }

            let ckpt_size = checkpoint_data.len() as f64;
            let ckpt_start = std::time::Instant::now();
            let best_checkpoint_path = checkpoint_callback(
                epoch + 1, checkpoint_data, true,
            ).context("Failed to save best checkpoint")?;
            training_metrics::record_checkpoint_save("dqn", "current", ckpt_start.elapsed().as_secs_f64(), ckpt_size);

            info!("Best model saved to: {}", best_checkpoint_path);
        }

        // Early stopping checks (skip if no training occurred)
        if train_step_count > 0 {
            let old_should_stop = self.check_early_stopping(log_output.avg_q_value, epoch);
            let patience_should_stop = if self.hyperparams.early_stopping_enabled
                && epoch + 1 >= self.hyperparams.min_epochs_before_stopping {
                self.early_stopping.should_stop(-log_output.epoch_sharpe)
            } else {
                false
            };

            if let Some(stop_reason) = old_should_stop {
                warn!(
                    "Early stopping triggered at epoch {}/{}: {}",
                    epoch + 1, self.hyperparams.epochs, stop_reason
                );
                info!(
                    "Final metrics: loss={:.6}, Q-value={:.4}",
                    log_output.avg_loss, log_output.avg_q_value
                );

                let checkpoint_data = self.serialize_model().await
                    .context("Failed to serialize model for early stopping checkpoint")?;
                let checkpoint_size = checkpoint_data.len();
                let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false)
                    .context("Failed to save early stopping checkpoint")?;
                info!(
                    "Early stopping checkpoint saved to: {} ({} bytes)",
                    checkpoint_path, checkpoint_size
                );

                return Err(anyhow::anyhow!(
                    "Training terminated by early stopping at epoch {}/{}: {}",
                    epoch + 1, self.hyperparams.epochs, stop_reason
                ));
            }

            if patience_should_stop {
                warn!(
                    "WAVE 24 Early stopping (patience) triggered at epoch {}/{}! No improvement for {} epochs",
                    epoch + 1, self.hyperparams.epochs, self.hyperparams.gradient_collapse_patience
                );
                info!("Best Sharpe: {:.4} at epoch {} (best val_loss proxy: {:.6})",
                    self.best_sharpe, self.best_epoch, self.early_stopping.get_best_val_loss());

                let checkpoint_data = self.serialize_model().await?;
                let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false)?;
                info!("Patience-based early stopping checkpoint saved to: {}", checkpoint_path);

                return Err(anyhow::anyhow!("Training terminated by patience-based early stopping at epoch {}", epoch + 1));
            }
        }

        Ok(())
    }
}

