//! DQN Trainer — Training metrics, Q-value diagnostics, and validation

use anyhow::Result;
use super::DQNTrainer;
use crate::TrainingMetrics;
use super::super::config::DQNAgentType;

impl DQNTrainer {
    /// Full training loop with existing logic (Wave 12 Group 3)
    /// Calculate average metrics for an epoch
    pub(crate) fn calculate_epoch_metrics(
        epoch_loss: f64,
        epoch_q_value: f64,
        epoch_gradient_norm: f64,
        samples_processed: usize,
    ) -> (f64, f64, f64) {
        if samples_processed > 0 {
            let count = samples_processed as f64;
            (
                epoch_loss / count,
                epoch_q_value / count,
                epoch_gradient_norm / count,
            )
        } else {
            (0.0, 0.0, 0.0)
        }
    }

    /// Compute validation loss on held-out data
    /// WAVE 10.6: Batched validation for 5-10x speedup
    
    // collect_qvalue_statistics removed: zero callers, was a Candle forward path.
    // Q-value stats are now computed exclusively via fused_ctx.compute_q_stats() (cuBLAS).

    /// Create final training metrics
    pub(crate) async fn create_final_metrics(
        &self,
        total_loss: f64,
        total_q_value: f64,
        total_gradient_norm: f64,
        total_reward: f64,
        num_epochs: usize,
        training_duration: std::time::Duration,
        early_stopped: bool,
        total_action_counts: [usize; 9], // 5 exposure levels
        total_factored_action_counts: [usize; 81], // 81 factored actions
    ) -> Result<TrainingMetrics> {
        let final_loss = total_loss / num_epochs as f64;
        let avg_q_value_final = total_q_value / num_epochs as f64;
        let avg_grad_norm_final = total_gradient_norm / num_epochs as f64;
        let avg_episode_reward = total_reward / num_epochs as f64;

        let mut metrics = TrainingMetrics {
            loss: final_loss,
            accuracy: 0.0,
            precision: 0.0,
            recall: 0.0,
            f1_score: 0.0,
            training_time_seconds: training_duration.as_secs_f64(),
            epochs_trained: num_epochs as u32,
            convergence_achieved: final_loss < 1.0,
            additional_metrics: std::collections::HashMap::new(),
        };

        metrics.add_metric("avg_q_value", avg_q_value_final);
        metrics.add_metric("avg_gradient_norm", avg_grad_norm_final);
        metrics.add_metric("final_epsilon", self.get_epsilon().await.unwrap_or(0.1));
        metrics.add_metric("avg_episode_reward", avg_episode_reward);

        // Action metrics: use factored 81-action space if branching, else 5 exposure levels
        let total_factored: usize = total_factored_action_counts.iter().sum();
        let total_exposure: usize = total_action_counts.iter().sum();
        let total_actions = total_factored.max(total_exposure);
        if total_actions > 0 {
            // Factored 81-action diversity (primary metric when branching)
            if total_factored > 0 {
                let unique_factored = total_factored_action_counts.iter()
                    .filter(|&&count| count > 0).count();
                let factored_diversity = (unique_factored as f64 / 81.0) * 100.0;
                metrics.add_metric("action_diversity", factored_diversity);
                metrics.add_metric("factored_unique_actions", unique_factored as f64);
                metrics.add_metric("action_space_size", 81.0);

                // Active factored actions (used >0.5% of the time)
                let active_threshold = (total_factored as f64 * 0.005).max(1.0);
                let active_count = total_factored_action_counts.iter()
                    .filter(|&&count| count as f64 >= active_threshold)
                    .count();
                let active_diversity_pct = (active_count as f64 / 81.0) * 100.0;
                metrics.add_metric("active_actions_count", active_count as f64);
                metrics.add_metric("active_diversity_pct", active_diversity_pct);

                // Top factored actions
                let mut sorted_actions: Vec<(usize, usize)> = total_factored_action_counts.iter()
                    .enumerate()
                    .map(|(idx, &count)| (idx, count))
                    .collect();
                sorted_actions.sort_by(|a, b| b.1.cmp(&a.1));

                if let Some((top1_idx, top1_count)) = sorted_actions.first() {
                    let top1_pct = (*top1_count as f64 / total_factored as f64) * 100.0;
                    metrics.add_metric("top1_action_idx", *top1_idx as f64);
                    metrics.add_metric("top1_action_count", *top1_count as f64);
                    metrics.add_metric("top1_action_pct", top1_pct);
                }

                let top5_count: usize = sorted_actions.iter().take(5).map(|(_, c)| c).sum();
                let top5_coverage_pct = (top5_count as f64 / total_factored as f64) * 100.0;
                metrics.add_metric("top5_coverage_pct", top5_coverage_pct);
            } else {
                // Fallback: 5 exposure levels
                let unique_actions = total_action_counts.iter()
                    .filter(|&&count| count > 0).count();
                let action_diversity = (unique_actions as f64 / 5.0) * 100.0;
                metrics.add_metric("action_diversity", action_diversity);
                metrics.add_metric("action_space_size", 5.0);

                let active_threshold = (total_exposure as f64 * 0.005).max(1.0);
                let active_count = total_action_counts.iter()
                    .filter(|&&count| count as f64 >= active_threshold)
                    .count();
                let active_diversity_pct = (active_count as f64 / 5.0) * 100.0;
                metrics.add_metric("active_actions_count", active_count as f64);
                metrics.add_metric("active_diversity_pct", active_diversity_pct);

                let mut sorted_actions: Vec<(usize, usize)> = total_action_counts.iter()
                    .enumerate()
                    .map(|(idx, &count)| (idx, count))
                    .collect();
                sorted_actions.sort_by(|a, b| b.1.cmp(&a.1));

                if let Some((top1_idx, top1_count)) = sorted_actions.first() {
                    let top1_pct = (*top1_count as f64 / total_exposure as f64) * 100.0;
                    metrics.add_metric("top1_action_idx", *top1_idx as f64);
                    metrics.add_metric("top1_action_count", *top1_count as f64);
                    metrics.add_metric("top1_action_pct", top1_pct);
                }

                let top5_count: usize = sorted_actions.iter().take(5).map(|(_, c)| c).sum();
                let top5_coverage_pct = (top5_count as f64 / total_exposure as f64) * 100.0;
                metrics.add_metric("top5_coverage_pct", top5_coverage_pct);
            }

            metrics.add_metric("total_actions", total_actions as f64);

            // Buy/sell/hold always from 5-exposure space (meaningful for P&L)
            // 0=Short100, 1=Short50, 2=Flat, 3=Long50, 4=Long100
            let sell_count: usize = total_action_counts.get(0).copied().unwrap_or(0)
                + total_action_counts.get(1).copied().unwrap_or(0);
            let hold_count: usize = total_action_counts.get(2).copied().unwrap_or(0);
            let buy_count: usize = total_action_counts.get(3).copied().unwrap_or(0)
                + total_action_counts.get(4).copied().unwrap_or(0);
            metrics.add_metric("buy_count", buy_count as f64);
            metrics.add_metric("sell_count", sell_count as f64);
            metrics.add_metric("hold_count", hold_count as f64);
        }

        // Compute Q-value standard deviation across epochs for hyperopt stability penalty.
        // self.q_value_history stores per-epoch average Q-values; their std measures
        // how much Q-values fluctuate during training (volatility indicator).
        if self.q_value_history.len() >= 2 {
            let n = self.q_value_history.len() as f64;
            let mean = self.q_value_history.iter().sum::<f64>() / n;
            let variance = self
                .q_value_history
                .iter()
                .map(|&q| (q - mean).powi(2))
                .sum::<f64>()
                / n;
            let std_dev = variance.sqrt();
            metrics.add_metric("q_value_std", std_dev);
        } else {
            // Not enough data points to compute std; default to 0.0 (no volatility signal)
            metrics.add_metric("q_value_std", 0.0);
        }

        if early_stopped {
            metrics.add_metric("early_stopped", 1.0);
        }

        // Aggregate trade stats across all epochs
        let agg_total_trades: usize = self.trade_stats_history.iter()
            .map(|ts| ts.total_trades)
            .sum();
        let agg_winning: usize = self.trade_stats_history.iter()
            .map(|ts| ts.winning_trades)
            .sum();
        let overall_win_rate = if agg_total_trades > 0 {
            agg_winning as f64 / agg_total_trades as f64
        } else {
            0.0
        };
        let agg_sum_wins: f64 = self.trade_stats_history.iter()
            .map(|ts| ts.sum_wins)
            .sum();
        let agg_sum_losses: f64 = self.trade_stats_history.iter()
            .map(|ts| ts.sum_losses)
            .sum();
        let overall_pf = if agg_sum_losses > 1e-10 {
            agg_sum_wins / agg_sum_losses
        } else if agg_sum_wins > 0.0 {
            f64::INFINITY
        } else {
            0.0
        };

        metrics.add_metric("total_trades", agg_total_trades as f64);
        metrics.add_metric("overall_win_rate", overall_win_rate);
        metrics.add_metric("overall_profit_factor", overall_pf);

        // Also add last-epoch Sharpe for hyperopt consumption
        if let Some(last_ts) = self.trade_stats_history.last() {
            let financials = super::super::financials::compute_epoch_financials(
                last_ts,
                &total_action_counts,
                100_000.0,
                self.hyperparams.bars_per_day,
            );
            metrics.add_metric("sharpe_ratio", financials.sharpe);
            metrics.add_metric("max_drawdown", financials.max_drawdown);
        }

        Ok(metrics)
    }

    /// Epoch-end Q-value diagnostics: gap analysis + per-action averages.
    ///
    /// Uses cuBLAS forward (`fused_ctx.compute_q_values`) — zero Candle involvement.
    /// Downloads the Q-value matrix to host for gap + per-action computation
    /// (cold path, epoch boundary only, ~8KB for 200 samples).
    ///
    /// Returns (gap_stats, per_action_avgs) where:
    /// - gap_stats: (mean_gap, min_gap, max_gap) of Q_best - Q_second_best
    /// - per_action_avgs: `[f64; 9]` averages (one per exposure action, capped at 5)
    pub(crate) async fn compute_epoch_q_diagnostics(&mut self) -> Option<(
        (f64, f64, f64),
        [f64; 9],
    )> {
        let agent = self.agent.read().await;
        let buffer = agent.memory();

        if buffer.len() < 10 {
            return None;
        }

        // Cap sample_size at fused_ctx batch_size to avoid exceeding GPU buffer
        let fused_batch = self.fused_ctx.as_ref().map(|f| f.batch_size()).unwrap_or(0);
        if fused_batch == 0 {
            return None;
        }
        let sample_size = buffer.len().min(200).min(fused_batch);
        let batch_sample = match buffer.sample(sample_size) {
            Ok(s) => s,
            Err(_) => return None,
        };

        // GPU PER path: states are already a CudaSlice<f32>
        let gpu_batch = batch_sample.gpu_batch.as_ref()?;
        let states_data = gpu_batch.states.data();

        // cuBLAS forward → Q-values in q_out_buf [sample_size, total_actions]
        let fused = self.fused_ctx.as_mut()?;
        let total_actions = fused.total_actions();
        let q_out = match fused.compute_q_values(states_data, sample_size) {
            Ok(q) => q,
            Err(_) => return None,
        };

        // Download Q-values to host (cold path, epoch boundary only).
        // q_out is [config.batch_size, total_actions] but we only need [sample_size, total_actions].
        // Allocate host buffer to match GPU buffer size, then truncate.
        let stream = self.cuda_stream.as_ref()?;
        let mut host_q = vec![0.0_f32; q_out.len()];
        stream.memcpy_dtoh(q_out, &mut host_q).ok()?;
        host_q.truncate(sample_size * total_actions);

        // Compute gap (best - second_best) and per-action averages on host
        let cols = total_actions;
        let rows = sample_size;
        let mut gaps = Vec::with_capacity(rows);
        let mut col_sums = [0.0_f64; 9];
        for r in 0..rows {
            let row_offset = r * cols;
            let mut best_val = f32::NEG_INFINITY;
            let mut second_best = f32::NEG_INFINITY;
            for c in 0..cols.min(9) {
                let v = host_q.get(row_offset + c).copied().unwrap_or(0.0);
                if let Some(s) = col_sums.get_mut(c) {
                    *s += v as f64;
                }
                if v > best_val {
                    second_best = best_val;
                    best_val = v;
                } else if v > second_best {
                    second_best = v;
                } else {
                    // v <= second_best: no update needed
                }
            }
            // Also check remaining columns (9..total_actions) for best/second_best
            // (order_type + urgency branches, if total_actions > 9)
            for c in 9..cols {
                let v = host_q.get(row_offset + c).copied().unwrap_or(0.0);
                if v > best_val {
                    second_best = best_val;
                    best_val = v;
                } else if v > second_best {
                    second_best = v;
                }
            }
            gaps.push(best_val - second_best);
        }

        let rows_f64 = rows as f64;
        let gap_mean = gaps.iter().sum::<f32>() / rows as f32;
        let gap_min = gaps.iter().copied().fold(f32::INFINITY, f32::min);
        let gap_max = gaps.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        let mut per_action_avgs = [0.0_f64; 9];
        for (i, avg) in per_action_avgs.iter_mut().enumerate() {
            *avg = col_sums.get(i).copied().unwrap_or(0.0) / rows_f64;
        }

        Some((
            (gap_mean as f64, gap_min as f64, gap_max as f64),
            per_action_avgs,
        ))
    }

    /// Get current training metrics
    pub async fn get_metrics(&self) -> TrainingMetrics {
        self.metrics.read().await.clone()
    }

    /// Check if early stopping criteria are met
    pub(crate) fn check_early_stopping(&self, avg_q_value: f64, epoch: usize) -> Option<String> {
        if !self.hyperparams.early_stopping_enabled
            || epoch + 1 < self.hyperparams.min_epochs_before_stopping
        {
            return None;
        }

        // Criterion 1: Q-value floor check
        if avg_q_value < self.hyperparams.q_value_floor {
            return Some(format!(
                "Q-value {:.4} below floor threshold {:.4}",
                avg_q_value, self.hyperparams.q_value_floor
            ));
        }

        // C4 FIX: Criterion 2 — Sharpe plateau check (was val-loss plateau).
        // Sharpe directly measures trading quality. A plateau means the model
        // has stopped improving its trading strategy, even if TD-loss still moves.
        if self.sharpe_history.len() >= self.hyperparams.plateau_window {
            let window = self.hyperparams.plateau_window;
            let recent_sharpes: Vec<f64> = self
                .sharpe_history
                .iter()
                .rev()
                .take(window)
                .copied()
                .collect();

            if let (Some(&newest), Some(&oldest)) = (recent_sharpes.first(), recent_sharpes.last()) {
                // For Sharpe, improvement = newest - oldest (higher is better)
                let improvement = newest - oldest;

                if improvement < 0.01 {
                    let msg = if improvement < -0.01 {
                        format!(
                            "Sharpe worsening detected (delta: {:.4}, window: {})",
                            improvement,
                            window
                        )
                    } else {
                        format!(
                            "Sharpe plateau detected (improvement: {:.4}, window: {})",
                            improvement,
                            window
                        )
                    };
                    return Some(msg);
                }
            }
        }

        None
    }

    /// Compute validation loss (negative Sharpe proxy) using cuBLAS forward — zero GpuTensor.
    ///
    /// Builds the state tensor on CPU (features + portfolio broadcast + OFI + alignment
    /// padding), uploads it once as a raw `CudaSlice<f32>`, runs cuBLAS forward via
    /// `fused_ctx`, then downloads Q-values and performs all action selection and Sharpe
    /// computation on CPU.  This avoids GpuTensor elementwise kernels (add, broadcast_mul,
    /// cat, etc.) that cache `CudaFunction` handles on the shared CUDA primary context,
    /// which causes sequential test failures.
    ///
    /// Epoch-boundary cold path — simplicity over GPU utilisation.
    pub(crate) async fn compute_validation_loss(&mut self) -> Result<f64> {
        if self.val_data.is_empty() {
            return Ok(0.0);
        }

        // Cap sample_size at fused_ctx batch_size to avoid exceeding cuBLAS buffer
        let fused_batch = self.fused_ctx.as_ref().map(|f| f.batch_size()).unwrap_or(0);
        if fused_batch == 0 {
            return Ok(0.0); // fused_ctx not initialized — skip validation
        }

        // Save current epsilon and force to 0 for deterministic evaluation
        let original_epsilon = self.get_epsilon().await?;
        self.set_epsilon(0.0).await?; // Pure greedy selection

        let sample_size = self.val_data.len().min(1000).min(fused_batch);

        let aligned_state_dim = {
            let agent = self.agent.read().await;
            agent.get_state_dim()
        };

        let stream = self.cuda_stream.as_ref()
            .ok_or_else(|| anyhow::anyhow!("cuda_stream required for validation"))?;

        // ── CPU-cached validation data (lazy init, built once per training run) ──
        // val_features_gpu stores flat features as Vec<f32> [sample_size * 42] on CPU.
        // val_closes_gpu stores (current_closes, next_closes) as Vec<f32> on CPU.
        // val_ofi_gpu stores flat OFI features as Vec<f32> [sample_size * 8] on CPU.
        if self.val_features_gpu.is_none() {
            let mut flat_features: Vec<f32> = Vec::with_capacity(sample_size * 42);
            let mut current_closes: Vec<f32> = Vec::with_capacity(sample_size);
            let mut next_closes: Vec<f32> = Vec::with_capacity(sample_size);

            for (feature_vec, target) in self.val_data.iter().take(sample_size) {
                for &v in feature_vec.iter() {
                    flat_features.push(v as f32);
                }
                let cur = if target.len() >= 2 { target[0] } else { feature_vec[3] };
                let nxt = if target.len() >= 2 { target[1] } else { cur };
                current_closes.push(cur as f32);
                next_closes.push(nxt as f32);
            }

            self.val_features_gpu = Some(flat_features);
            self.val_closes_gpu = Some((current_closes, next_closes));

            if let Some(ref ofi) = self.ofi_features {
                let mut flat_ofi: Vec<f32> = Vec::with_capacity(sample_size * 8);
                for i in 0..sample_size {
                    let idx = self.ofi_val_offset + i;
                    if let Some(row) = ofi.get(idx) {
                        for &v in row.iter() {
                            flat_ofi.push(v as f32);
                        }
                    } else {
                        flat_ofi.extend_from_slice(&[0.0_f32; 8]);
                    }
                }
                self.val_ofi_gpu = Some(flat_ofi);
            }
        }

        let flat_features = self.val_features_gpu.as_ref()
            .ok_or_else(|| anyhow::anyhow!("val_features_gpu must be initialized"))?;
        let (current_closes, next_closes) = self.val_closes_gpu.as_ref()
            .ok_or_else(|| anyhow::anyhow!("val_closes_gpu must be initialized"))?;

        // ── Build complete state tensor on CPU per validation call ──
        // Portfolio features change every epoch (position, P&L), so we rebuild each call.
        // Feature layout per row: [42 market features | 3 portfolio | 8 OFI (opt) | padding]
        let val_price = self.val_data.first()
            .map(|(fv, tgt)| if tgt.len() >= 2 { tgt[0] as f32 } else { fv[3] as f32 })
            .unwrap_or(0.0);
        let portfolio_f = self.portfolio_tracker.get_portfolio_features(val_price);

        let raw_state_dim = if self.hyperparams.mbp10_data_dir.is_some() { 74 } else { 66 };
        let pad_cols = if aligned_state_dim > raw_state_dim {
            aligned_state_dim - raw_state_dim
        } else {
            0
        };

        let mut state_flat: Vec<f32> = Vec::with_capacity(sample_size * aligned_state_dim);
        let feat_cols = flat_features.len() / sample_size; // 42
        for i in 0..sample_size {
            // Market features
            let feat_start = i * feat_cols;
            let feat_end = feat_start + feat_cols;
            let feat_slice = flat_features.get(feat_start..feat_end)
                .ok_or_else(|| anyhow::anyhow!("val features OOB at row {i}"))?;
            state_flat.extend_from_slice(feat_slice);

            // Portfolio features (same for all rows — broadcast)
            state_flat.extend_from_slice(&portfolio_f[..3]);

            // OFI features (optional, 8 floats per row)
            if let Some(ref flat_ofi) = self.val_ofi_gpu {
                let ofi_start = i * 8;
                let ofi_end = ofi_start + 8;
                let ofi_slice = flat_ofi.get(ofi_start..ofi_end)
                    .ok_or_else(|| anyhow::anyhow!("val OFI OOB at row {i}"))?;
                state_flat.extend_from_slice(ofi_slice);
            }

            // Alignment padding (zeros)
            state_flat.extend(std::iter::repeat(0.0_f32).take(pad_cols));
        }

        // Upload state to GPU as raw CudaSlice<f32> — no GpuTensor, no elementwise kernels.
        let state_gpu = stream.clone_htod(&state_flat)
            .map_err(|e| anyhow::anyhow!("val state HtoD: {e}"))?;

        // cuBLAS forward via fused_ctx — zero GpuTensor involvement
        let fused = self.fused_ctx.as_mut()
            .ok_or_else(|| anyhow::anyhow!("fused_ctx required for validation forward"))?;
        let total_actions = fused.total_actions();
        let q_out = fused.compute_q_values(&state_gpu, sample_size)
            .map_err(|e| anyhow::anyhow!("cuBLAS validation forward: {e}"))?;

        // Download Q-values to CPU for action selection + Sharpe computation.
        // q_out is [config.batch_size, total_actions]; allocate to match, then truncate.
        let mut host_q = vec![0.0_f32; q_out.len()];
        stream.memcpy_dtoh(q_out, &mut host_q)
            .map_err(|e| anyhow::anyhow!("Q-val DtoH: {e}"))?;
        host_q.truncate(sample_size * total_actions);

        // ── CPU greedy action selection ──
        // Direction LUT maps exposure head index → position size direction.
        // 0=Short100(-1.0), 1=Short50(-0.5), 2=Flat(0.0), 3=Long50(0.5), 4=Long100(1.0)
        const DIRECTION_LUT: [f32; 9] = [-1.0, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0];
        let use_branching = true && total_actions >= 15;

        let rewards: Vec<f32> = (0..sample_size).map(|i| {
            let row_start = i * total_actions;
            let row = &host_q[row_start..row_start + total_actions];

            // Greedy exposure action selection (argmax over exposure head, 5 values)
            let exposure_idx = if use_branching {
                // Branching DQN: exposure head is first 9 Q-values
                row.get(..9).unwrap_or(row)
                    .iter()
                    .enumerate()
                    .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
                    .map(|(j, _)| j)
                    .unwrap_or(2) // default: Flat
            } else {
                // Non-branching: argmax over all Q-values maps to exposure
                row.iter()
                    .enumerate()
                    .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
                    .map(|(j, _)| j.min(4))
                    .unwrap_or(2) // default: Flat
            };
            let direction = DIRECTION_LUT[exposure_idx.min(8)];

            // PnL-based reward: direction × (next_close - cur_close) / cur_close
            let cur = current_closes.get(i).copied().unwrap_or(0.0);
            let nxt = next_closes.get(i).copied().unwrap_or(cur);
            if cur.abs() > 1e-10 {
                direction * (nxt - cur) / cur
            } else {
                0.0
            }
        }).collect();

        // ── CPU Sharpe computation ──
        let n = rewards.len() as f64;
        let mean_scalar: f64 = rewards.iter().map(|&v| v as f64).sum::<f64>() / n;
        let var_scalar: f64 = if n > 1.0 {
            rewards.iter().map(|&v| {
                let d = v as f64 - mean_scalar;
                d * d
            }).sum::<f64>() / n
        } else {
            0.0
        };
        let std_val = var_scalar.sqrt();
        let val_sharpe = if std_val > 1e-10 {
            (mean_scalar / std_val) * (252.0_f64).sqrt()
        } else {
            0.0
        };

        self.set_epsilon(original_epsilon).await?;

        Ok(-val_sharpe)
    }

    /// Estimate average Q-value from replay buffer samples for monitoring.
    ///
    /// Uses the cuBLAS forward path (`fused_ctx.compute_q_values`) — zero Candle involvement.
    /// Eliminates the 27ms/epoch overhead from `agent.forward()` (Candle dispatch chain).
    ///
    /// WAVE 23 P0: Includes Q-value divergence check (early stopping) via
    /// `agent.log_q_values_from_stats()` — computed from the host readback.
    pub(crate) async fn estimate_avg_q_value_with_early_stopping(&mut self, agent: &mut DQNAgentType) -> Result<f64> {
        let buffer = agent.memory();

        if buffer.len() == 0 {
            return Ok(0.0);
        }

        // Sample up to 10 experiences for Q-value estimation
        let sample_size = buffer.len().min(10);
        let batch_sample = buffer
            .sample(sample_size)
            .map_err(|e| anyhow::anyhow!("Failed to sample experiences: {}", e))?;

        // GPU PER path: use gpu_batch.states directly (F32 CudaSlice)
        let gpu_batch = batch_sample.gpu_batch.as_ref()
            .ok_or_else(|| anyhow::anyhow!("GPU PER must be active -- gpu_batch is None"))?;

        // cuBLAS forward + GPU reduction → 5 scalars, 20-byte readback
        let fused = self.fused_ctx.as_mut()
            .ok_or_else(|| anyhow::anyhow!("fused_ctx required for Q-value estimation"))?;

        let states_data = gpu_batch.states.data();
        let stats = fused.compute_q_stats(states_data, sample_size)
            .map_err(|e| anyhow::anyhow!("cuBLAS Q-stats: {}", e))?;

        // Q-value divergence check (early stopping)
        let num_actions = fused.total_actions();
        agent.log_q_values_from_stats(stats.q_min, stats.q_max, stats.q_mean, stats.q_variance, num_actions)
            .map_err(|e| {
                tracing::info!("Early stopping triggered (Q-value divergence): {}", e);
                anyhow::anyhow!("Early stopping: {}", e)
            })?;

        Ok(stats.avg_max_q)
    }

}
