//! DQN training step methods — single-batch and gradient-accumulation paths.

use anyhow::Result;
use ml_core::cuda_autograd::GpuTensor;
use tracing::{debug, info, warn};

use super::DQNTrainer;

impl DQNTrainer {
    /// Perform one training step using real DQN algorithm
    ///
    /// This method implements the core Deep Q-Learning algorithm:
    /// 1. Sample batch from experience replay buffer
    /// 2. Compute current Q-values: Q(s, a)
    /// 3. Compute target Q-values: r + γ * max_a' Q_target(s', a')
    /// 4. Calculate TD-error and MSE loss
    /// 5. Backpropagate gradients and update Q-network
    /// 6. Periodically update target network
    ///
    /// WAVE 26 P2.2: Now supports gradient accumulation for larger effective batch sizes
    ///
    /// Returns: (loss, avg_q_value, grad_norm)
    async fn train_step(&mut self) -> Result<(f64, f64, f64)> {
        let accumulation_steps = self.hyperparams.gradient_accumulation_steps;

        // OOM recovery loop: retry up to 3 times with halved batch size
        const MAX_OOM_RETRIES: usize = 3;

        for retry in 0..=MAX_OOM_RETRIES {
            let result = if accumulation_steps > 1 {
                self.train_step_with_accumulation().await
            } else {
                self.train_step_single_batch().await
            };

            match result {
                Ok(metrics) => return Ok(metrics),
                Err(e) => {
                    // Check if this is an OOM error by inspecting the error chain
                    let err_str = format!("{:?}", e).to_lowercase();
                    let is_oom = err_str.contains("out of memory")
                        || err_str.contains("oom")
                        || err_str.contains("cuda error 2")
                        || err_str.contains("cudamalloc")
                        || err_str.contains("failed to allocate");

                    if is_oom && retry < MAX_OOM_RETRIES {
                        let old_batch = self.current_batch_size;
                        self.current_batch_size = (old_batch / 2).max(1);
                        warn!(
                            "OOM detected (retry {}/{}): reducing batch size {} -> {}",
                            retry + 1,
                            MAX_OOM_RETRIES,
                            old_batch,
                            self.current_batch_size
                        );
                        // Continue to next retry
                    } else {
                        return Err(e);
                    }
                }
            }
        }

        Err(anyhow::anyhow!(
            "Training failed after {} OOM retries",
            MAX_OOM_RETRIES
        ))
    }

    /// Standard single-batch training step (no gradient accumulation)
    ///
    /// Pre-samples from PER buffer using READ lock before acquiring WRITE lock
    /// for GPU training. This preserves PER IS-weights and indices for proper
    /// importance sampling correction and priority updates.
    async fn train_step_single_batch(&mut self) -> Result<(f64, f64, f64)> {
        // Pre-sample batch OUTSIDE the write lock using read-only access to the buffer.
        // PER IS-weights and indices are preserved for correct importance sampling.
        // The write lock is only held during GPU forward/backward + optimizer step.
        let explicit_batch = {
            let agent = self.agent.read().await;
            let buffer = agent.memory();
            let sample_size = self.current_batch_size;
            buffer.can_sample(sample_size).then(|| {
                buffer
                    .sample(sample_size)
                    .map_err(|e| anyhow::anyhow!("Failed to pre-sample batch: {}", e))
            }).transpose()?
        }; // READ lock released here

        let mut agent = self.agent.write().await;

        // train_step returns GpuTrainResult with GPU-resident scalar tensors.
        #[allow(unused_variables)]
        let gpu_result = agent
            .train_step(explicit_batch)
            .map_err(|e| anyhow::anyhow!("Training step failed: {}", e))?;

        // GPU training guard: on-device NaN/loss-clip/grad-collapse checks.
        // Zero cudaStreamSynchronize — kernel writes halt flags to pinned host memory.
        let (loss_clipped, grad_norm) = {
            // Lazy-init training guard on first call
            if self.training_guard.is_none() && self.device.is_cuda() {
                match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::from_stream(std::sync::Arc::clone(self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream for training guard"))?)) {
                    Ok(guard) => {
                        info!("GPU training guard initialized");
                        self.training_guard = Some(guard);
                    }
                    Err(e) => {
                        return Err(anyhow::anyhow!("GPU training guard init FAILED (no CPU fallback): {e}"));
                    }
                }
            }

            if let Some(ref mut guard) = self.training_guard {
                let grad_collapse_threshold =
                    self.hyperparams.learning_rate as f32
                        * self.hyperparams.gradient_collapse_multiplier as f32;
                // Use original buffer_size (before AutoReplaySizer) for warmup guard
                let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64;
                let past_warmup = self.gradient_logging_step as u64 > warmup_steps;

                let loss_slice = gpu_result.loss_cuda_slice()
                    .map_err(|e| anyhow::anyhow!("GPU guard loss CudaSlice: {e}"))?;
                let grad_slice = gpu_result.grad_norm_cuda_slice()
                    .map_err(|e| anyhow::anyhow!("GPU guard grad CudaSlice: {e}"))?;
                let result = guard
                    .check_and_accumulate(
                        &loss_slice,
                        &grad_slice,
                        1e6_f32, // loss clip threshold
                        grad_collapse_threshold,
                        !past_warmup,
                    )
                    .map_err(|e| anyhow::anyhow!("GPU guard check: {e}"))?;

                // Handle halt conditions
                if result.halt_nan {
                    return Err(anyhow::anyhow!(
                        "NaN/Inf detected in loss ({}) or grad_norm ({})",
                        result.raw_loss,
                        result.raw_grad_norm
                    ));
                }
                if result.halt_loss_clip {
                    warn!(
                        "Loss clipped from {:.2e} to 1.0e6 (TD error explosion, epoch {})",
                        result.raw_loss,
                        self.loss_history.len() + 1
                    );
                }

                // GPU guard path: collapse check only (no detect_dead_neurons GPU->CPU sync).
                // Dead neuron detection runs at epoch boundary via log_diagnostics().
                agent.check_gradient_collapse(result.raw_grad_norm).map_err(|e| {
                    tracing::info!("Early stopping triggered (gradient collapse): {}", e);
                    anyhow::anyhow!("Early stopping: {}", e)
                })?;

                (result.clipped_loss as f64, result.raw_grad_norm as f64)
            } else {
                return Err(anyhow::anyhow!(
                    "GPU training guard not initialized — CUDA device required for DQN training"
                ));
            }
        };
        // Unreachable in non-cuda mode (return above), but Rust still name-checks.

        // Q-value estimation: periodic (every 50 steps)
        // cuBLAS forward + GPU reduction → 5 scalars, 20-byte readback (zero Candle).
        self.q_estimation_counter += 1;
        if self.q_estimation_counter % 50 == 1 {
            let buffer = agent.memory();
            if buffer.len() > 0 {
                let sample_size = buffer.len().min(10);
                let batch_sample = buffer
                    .sample(sample_size)
                    .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?;

                // GPU PER path: states are already a CudaSlice<f32>
                let gpu_batch = batch_sample.gpu_batch.as_ref()
                    .ok_or_else(|| anyhow::anyhow!("GPU PER must be active — gpu_batch is None"))?;
                let states_data = gpu_batch.states.data();

                // cuBLAS forward + GPU reduction → QValueStatsResult (20-byte readback)
                let fused = self.fused_ctx.as_mut()
                    .ok_or_else(|| anyhow::anyhow!("fused_ctx required for Q-value estimation"))?;
                let stats = fused.compute_q_stats(states_data, sample_size)
                    .map_err(|e| anyhow::anyhow!("cuBLAS Q-stats: {e}"))?;
                let num_actions = fused.total_actions();

                // Q-value divergence check (early stopping)
                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 (Q-value divergence): {}", e);
                    anyhow::anyhow!("Early stopping: {}", e)
                })?;

                self.cached_avg_q = stats.q_mean as f64;

                // Accumulate Q-value mean via Welford running mean (CPU-side f32)
                if let Some(ref mut guard) = self.training_guard {
                    guard.accumulate_q_value(stats.q_mean);
                }
            }
        }
        let avg_q_value = self.cached_avg_q;

        debug!("Gradient norm after clip (actual): {:.4}", grad_norm);

        self.gradient_logging_step += 1;
        if self.gradient_logging_step % 10 == 0 {
            debug!(
                "Step {}: grad={:.4}, loss={:.4}",
                self.gradient_logging_step, grad_norm, loss_clipped
            );
        }

        Ok((loss_clipped, avg_q_value, grad_norm))
    }

    /// Training step with true gradient accumulation across N mini-batches.
    ///
    /// Unlike the previous implementation which ran N independent optimizer
    /// steps, this version computes gradients for each mini-batch, accumulates
    /// them, averages, and then applies a **single** optimizer step. This
    /// simulates training with an effective batch size of
    /// `accumulation_steps * batch_size` while keeping memory usage at
    /// `batch_size`.
    ///
    /// Returns: (avg_loss, avg_q_value, final_grad_norm)
    async fn train_step_with_accumulation(&mut self) -> Result<(f64, f64, f64)> {
        let accumulation_steps = self.hyperparams.gradient_accumulation_steps;

        debug!(
            "Starting true gradient accumulation with {} steps (effective batch: {})",
            accumulation_steps,
            self.current_batch_size * accumulation_steps
        );

        // Pre-sample ALL mini-batches using READ lock (no GPU contention).
        let pre_sampled: Vec<Option<crate::dqn::replay_buffer_type::BatchSample>> = {
            let agent = self.agent.read().await;
            let buffer = agent.memory();
            let sample_size = self.current_batch_size;
            let mut batches = Vec::with_capacity(accumulation_steps);
            for step_idx in 0..accumulation_steps {
                batches.push(
                    buffer.can_sample(sample_size).then(|| {
                        buffer.sample(sample_size).map_err(|e| {
                            anyhow::anyhow!(
                                "Failed to pre-sample batch (accum step {}): {}",
                                step_idx,
                                e
                            )
                        })
                    }).transpose()?,
                );
            }
            batches
        }; // READ lock released

        let mut agent = self.agent.write().await;

        // === Phase 1: Accumulate gradients across N mini-batches ===
        let mut accumulated_grads: Option<std::collections::BTreeMap<String, GpuTensor>> = None;
        // Used by non-CUDA fallback and CUDA empty-tensor fallback paths.
        #[allow(unused_mut, unused_assignments, unused_variables)]
        let mut total_loss = 0.0_f64;
        let mut all_td_errors = Vec::new();
        let mut all_indices = Vec::new();
        #[allow(unused_mut, unused_assignments, unused_variables)]
        let mut final_grad_norm = 0.0_f32;
        let mut gpu_td_errors: Vec<GpuTensor> = Vec::new();
        let mut gpu_indices: Vec<GpuTensor> = Vec::new();
        let mut gpu_loss_tensors: Vec<GpuTensor> = Vec::new();
        let mut gpu_grad_tensors: Vec<GpuTensor> = Vec::new();

        for (step, batch) in pre_sampled.into_iter().enumerate() {
            // Compute forward pass + backward WITHOUT optimizer step
            let result = agent
                .compute_gradients(batch)
                .map_err(|e| anyhow::anyhow!("Gradient computation step {} failed: {}", step, e))?;

            // Extract CudaSlice handles and destructure BEFORE moving fields out.
            let loss_sl_opt = result.loss_cuda_slice();
            let grad_sl_opt = result.grad_norm_cuda_slice();
            let r_grads = result.grads;
            let r_td_errors = result.td_errors;
            let r_indices = result.indices.clone();
            let r_td_gpu = result.td_errors_gpu;
            let r_idx_gpu = result.indices_gpu;
            let r_loss_gpu = result.loss_tensor_gpu;
            let r_gn_gpu = result.grad_norm_gpu;

            // Get vars for accumulation. Var is an Arc wrapper so cloning is cheap.
            let _vars: Vec<GpuTensor> = agent
                .optimizer_vars()
                .map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?;

            // Accumulate gradients: add r_grads into accumulated_grads elementwise
            match &mut accumulated_grads {
                Some(ref mut acc) => {
                    for (name, grad) in r_grads {
                        if let Some(existing) = acc.get(&name) {
                            let stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream needed for grad accum"))?;
                            let summed = existing.add(&grad, stream)
                                .map_err(|e| anyhow::anyhow!("Gradient accumulation step {} add '{name}': {e}", step))?;
                            acc.insert(name, summed);
                        } else {
                            acc.insert(name, grad);
                        }
                    }
                }
                None => {
                    accumulated_grads = Some(r_grads);
                }
            }

            all_td_errors.extend(r_td_errors);
            all_indices.extend(r_indices);
            // GPU guard: check + accumulate loss/grad for this sub-step.
            // Uses CudaSlice accessors on GradientResult to avoid tensor_to_cuda_slice_f32.
            {
                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 grad_collapse_threshold =
                            self.hyperparams.learning_rate as f32
                            * self.hyperparams.gradient_collapse_multiplier as f32;
                        // Use original buffer_size (before AutoReplaySizer) for warmup guard
                        let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64;
                        let past_warmup = self.gradient_logging_step as u64 > warmup_steps;

                        let loss_slice = loss_sl_r
                            .map_err(|e| anyhow::anyhow!("GPU guard accum loss CudaSlice: {e}"))?;
                        let grad_slice = grad_sl_r
                            .map_err(|e| anyhow::anyhow!("GPU guard accum grad CudaSlice: {e}"))?;
                        let guard_result = guard.check_and_accumulate(
                            &loss_slice,
                            &grad_slice,
                            1e6_f32,
                            grad_collapse_threshold,
                            !past_warmup,
                        ).map_err(|e| anyhow::anyhow!("GPU guard sub-step {}: {e}", step))?;

                        if guard_result.halt_nan {
                            return Err(anyhow::anyhow!(
                                "NaN/Inf at accumulation sub-step {}: loss={}, grad={}",
                                step, guard_result.raw_loss, guard_result.raw_grad_norm
                            ));
                        }
                        if guard_result.halt_grad_collapse {
                            agent.check_gradient_collapse(guard_result.raw_grad_norm).map_err(|e| {
                                tracing::info!("Early stopping (gradient collapse): {}", e);
                                anyhow::anyhow!("Early stopping: {}", e)
                            })?;
                        }
                    }
                }
            }
            {
                if let Some(td_gpu) = r_td_gpu {
                    gpu_td_errors.push(td_gpu);
                }
                if let Some(idx_gpu) = r_idx_gpu {
                    gpu_indices.push(idx_gpu);
                }
                if let Some(loss_gpu) = r_loss_gpu {
                    gpu_loss_tensors.push(loss_gpu);
                }
                if let Some(gn_gpu) = r_gn_gpu {
                    gpu_grad_tensors.push(gn_gpu);
                }
            }
        }

        // === Phase 2: Average and apply gradients (single optimizer step) ===
        if let Some(ref mut grads) = accumulated_grads {
            let _vars: Vec<GpuTensor> = agent
                .optimizer_vars()
                .map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?;

            // Scale accumulated gradients by 1/N
            let scale = 1.0 / accumulation_steps as f32;
            let stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream needed for grad scale"))?;
            let scale_tensor = GpuTensor::scalar(scale, stream)
                .map_err(|e| anyhow::anyhow!("Gradient scale tensor: {e}"))?;
            for (_name, grad) in grads.iter_mut() {
                *grad = grad.broadcast_mul(&scale_tensor, stream)
                    .map_err(|e| anyhow::anyhow!("Gradient scaling failed: {e}"))?;
            }

            let guard_active = self.training_guard.is_some();
            // Gradient finiteness is checked by GpuTrainingGuard per sub-step.
            // No additional check needed here when guard is active.
            let _ = guard_active;

            agent
                .apply_accumulated_gradients(grads)
                .map_err(|e| anyhow::anyhow!("Apply accumulated gradients failed: {}", e))?;
        }

        // === Phase 3: Bookkeeping ===
        {
            // GPU PER path: concatenate GPU tensors and update in one shot
            if !gpu_td_errors.is_empty() && !gpu_indices.is_empty() {
                let cat_stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream for PER concat"))?;
                let td_refs: Vec<&GpuTensor> = gpu_td_errors.iter().collect();
                let idx_refs: Vec<&GpuTensor> = gpu_indices.iter().collect();
                let td_cat = GpuTensor::cat(&td_refs, 0, cat_stream)
                    .map_err(|e| anyhow::anyhow!("GPU TD error concat failed: {}", e))?;
                let idx_cat = GpuTensor::cat(&idx_refs, 0, cat_stream)
                    .map_err(|e| anyhow::anyhow!("GPU index concat failed: {}", e))?;
                agent
                    .update_priorities_gpu(&idx_cat, &td_cat)
                    .map_err(|e| anyhow::anyhow!("GPU PER priority update failed: {}", e))?;
            } else if !all_indices.is_empty() {
                agent
                    .update_priorities(&all_indices, &all_td_errors)
                    .map_err(|e| anyhow::anyhow!("PER priority update failed: {}", e))?;
            } else {
                // No priority updates needed (uniform buffer or empty batch)
            }
        }
        agent.step_replay_buffer();

        // Single readback at accumulation boundary — prefer GPU guard accumulators
        // (zero extra sync), fall back to cat+mean+to_scalar if guard absent.
        let (avg_loss, final_grad_norm_f64) = {
            if let Some(ref mut guard) = self.training_guard {
                let (avg_l, avg_gn) = guard.read_accumulators()
                    .map_err(|e| anyhow::anyhow!("GPU guard read_accumulators: {e}"))?;
                guard.reset_accumulators()
                    .map_err(|e| anyhow::anyhow!("GPU guard reset: {e}"))?;
                (avg_l, avg_gn)
            } else {
                return Err(anyhow::anyhow!(
                    "GPU training guard not initialized — CUDA device required"
                ));
            }
        };

        // Gradient collapse detection (early stopping) -- no dead neuron GPU->CPU sync.
        // Dead neuron detection runs at epoch boundary via log_diagnostics().
        agent
            .check_gradient_collapse(final_grad_norm_f64 as f32)
            .map_err(|e| {
                tracing::info!("Early stopping triggered (gradient collapse): {}", e);
                anyhow::anyhow!("Early stopping: {}", e)
            })?;

        // Clip averaged loss
        let loss_clipped = if avg_loss > 1e6 {
            warn!("Averaged loss clipped from {:.2e} to 1.0e6", avg_loss);
            1e6
        } else {
            avg_loss
        };

        // Q-value estimation: periodic (every 50 steps)
        // cuBLAS forward + GPU reduction → 5 scalars, 20-byte readback (zero Candle).
        self.q_estimation_counter += 1;
        if self.q_estimation_counter % 50 == 1 {
            let buffer = agent.memory();
            if buffer.len() > 0 {
                let sample_size = buffer.len().min(10);
                let batch_sample = buffer
                    .sample(sample_size)
                    .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?;

                // GPU PER path: states are already a CudaSlice<f32>
                let gpu_batch = batch_sample.gpu_batch.as_ref()
                    .ok_or_else(|| anyhow::anyhow!("GPU PER must be active — gpu_batch is None"))?;
                let states_data = gpu_batch.states.data();

                // cuBLAS forward + GPU reduction → QValueStatsResult (20-byte readback)
                let fused = self.fused_ctx.as_mut()
                    .ok_or_else(|| anyhow::anyhow!("fused_ctx required for Q-value estimation"))?;
                let stats = fused.compute_q_stats(states_data, sample_size)
                    .map_err(|e| anyhow::anyhow!("cuBLAS Q-stats: {e}"))?;
                let num_actions = fused.total_actions();

                // Q-value divergence check (early stopping)
                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 (Q-value divergence): {}", e);
                    anyhow::anyhow!("Early stopping: {}", e)
                })?;

                self.cached_avg_q = stats.q_mean as f64;

                // Accumulate Q-value mean via Welford running mean (CPU-side f32)
                if let Some(ref mut guard) = self.training_guard {
                    guard.accumulate_q_value(stats.q_mean);
                }
            }
        }
        let avg_q_value = self.cached_avg_q;

        Ok((loss_clipped, avg_q_value, final_grad_norm_f64))
    }

    /// Lazy-init fused CUDA training context (Standard DQN only).
    /// Recreate if batch_size changed (OOM recovery).
    pub(crate) async fn ensure_fused_ctx(&mut self) -> anyhow::Result<()> {
        let needs_init = match &self.fused_ctx {
            None => self.device.is_cuda(),
            Some(ctx) => ctx.batch_size() != self.current_batch_size,
        };
        if !needs_init {
            return Ok(());
        }
        let stream = match self.cuda_stream {
            Some(ref s) => std::sync::Arc::clone(s),
            None => return Ok(()),
        };
        if self.fused_ctx.is_some() {
            info!("Fused CUDA context: batch_size changed, recreating");
            self.fused_ctx = None;
        }
        let agent = self.agent.read().await;
        match super::super::fused_training::FusedTrainingCtx::new(
            &self.device, &*agent, &self.hyperparams, self.current_batch_size, stream,
        ) {
            Ok(ctx) => {
                info!("Fused CUDA training initialized (batch_size={})", self.current_batch_size);
                self.fused_ctx = Some(ctx);
            }
            Err(e) => {
                return Err(anyhow::anyhow!(
                    "Fused CUDA training init FAILED (no fallback — non-fused path \
                     has action-space mismatch with branching DQN): {e}"
                ));
            }
        }
        Ok(())
    }
}
