//! GPU-native AdamW optimizer.
//!
//! Runs the AdamW update rule entirely on GPU via a single CUDA kernel launch
//! per parameter tensor. No Candle tensors, no GradStore, no CPU roundtrips.
//!
//! Reuses the same AdamW formula as the DQN fused trainer's `dqn_adam_update_kernel`,
//! but operates on individual named parameters from a [`GpuVarStore`] rather than
//! a flat parameter buffer.

use std::collections::BTreeMap;
use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;

use crate::MLError;
use super::var_store::GpuVarStore;
use super::gpu_tensor::GpuTensor;

/// AdamW optimizer hyperparameters.
#[derive(Debug, Clone)]
pub struct AdamWConfig {
    /// Learning rate (default: 3e-4).
    pub lr: f32,
    /// First moment decay (default: 0.9).
    pub beta1: f32,
    /// Second moment decay (default: 0.999).
    pub beta2: f32,
    /// Numerical stability epsilon (default: 1e-8).
    pub epsilon: f32,
    /// Decoupled weight decay (default: 1e-5).
    pub weight_decay: f32,
    /// Maximum gradient L2 norm for clipping (default: 10.0).
    /// Set to `f32::INFINITY` to disable clipping.
    pub max_grad_norm: f32,
    /// Multiplicative LR decay applied each step: `lr *= (1 - lr_decay)`.
    /// `None` = no decay (default). Replaces `candle_optimisers::Decay`.
    pub lr_decay: Option<f32>,
}

impl Default for AdamWConfig {
    fn default() -> Self {
        Self {
            lr: 3e-4,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            weight_decay: 1e-5,
            max_grad_norm: 10.0,
            lr_decay: None,
        }
    }
}

/// Per-parameter optimizer state (first and second moment estimates).
struct ParamState {
    /// First moment (moving average of gradients).
    m: CudaSlice<f32>,
    /// Second moment (moving average of squared gradients).
    v: CudaSlice<f32>,
}

/// GPU-native AdamW optimizer.
///
/// Maintains per-parameter first/second moment buffers on GPU.
/// Each `step()` call launches one CUDA kernel per parameter (typically 10-40
/// parameters for a DQN/PPO network = 10-40 kernel launches, each <1us).
///
/// For networks with hundreds of parameters, consider the flat-buffer approach
/// used by [`GpuDqnTrainer`] instead (single kernel launch over all params).
pub struct GpuAdamW {
    config: AdamWConfig,
    /// Per-parameter optimizer state, keyed by parameter name.
    states: BTreeMap<String, ParamState>,
    /// Global step counter (for bias correction).
    step_count: i32,
    /// Compiled AdamW kernel.
    kernel: CudaFunction,
    /// Compiled gradient norm kernel (for clipping).
    grad_norm_kernel: CudaFunction,
    /// CUDA stream for all operations.
    stream: Arc<CudaStream>,
}

impl std::fmt::Debug for GpuAdamW {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "GpuAdamW {{ config: {:?}, step: {}, params: {} }}",
            self.config,
            self.step_count,
            self.states.len()
        )
    }
}

impl GpuAdamW {
    /// Create a new AdamW optimizer.
    ///
    /// Does NOT pre-allocate moment buffers — they are lazily created on the
    /// first `step()` call for each parameter.
    pub fn new(
        config: AdamWConfig,
        stream: Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let context = stream.context();
        let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(ADAMW_CUDA_SRC, &context)
            .map_err(|e| MLError::ModelError(format!("AdamW kernel compilation: {e}")))?;
        let module = context.load_module(ptx).map_err(|e| {
            MLError::ModelError(format!("AdamW module load: {e}"))
        })?;

        let kernel = module.load_function("adamw_update").map_err(|e| {
            MLError::ModelError(format!("adamw_update load: {e}"))
        })?;
        let grad_norm_kernel = module.load_function("grad_l2_norm").map_err(|e| {
            MLError::ModelError(format!("grad_l2_norm load: {e}"))
        })?;

        Ok(Self {
            config,
            states: BTreeMap::new(),
            step_count: 0,
            kernel,
            grad_norm_kernel,
            stream,
        })
    }

    /// Current step count.
    pub fn step_count(&self) -> i32 {
        self.step_count
    }

    /// Update the learning rate.
    pub fn set_learning_rate(&mut self, lr: f32) {
        self.config.lr = lr;
    }

    /// Current learning rate.
    pub fn learning_rate(&self) -> f32 {
        self.config.lr
    }

    /// Perform one optimization step.
    ///
    /// `gradients` maps parameter names to their gradient `GpuTensor`.
    /// Only parameters with matching names in the `store` are updated.
    /// Gradient clipping is applied via a scale factor passed to the CUDA kernel.
    pub fn step(
        &mut self,
        store: &mut GpuVarStore,
        gradients: &BTreeMap<String, GpuTensor>,
    ) -> Result<f32, MLError> {
        self.step_count += 1;
        let t = self.step_count;

        // Apply multiplicative LR decay if configured
        if let Some(decay) = self.config.lr_decay {
            self.config.lr *= 1.0 - decay;
        }

        // ── Global gradient norm computation + clipping scale ────────
        let (grad_norm, grad_scale) = if self.config.max_grad_norm.is_finite() {
            self.compute_grad_norm_and_scale(gradients)?
        } else {
            (0.0_f32, 1.0_f32)
        };

        // ── Per-parameter AdamW update ───────────────────────────────
        for (name, grad) in gradients {
            let param = match store.get_mut(name) {
                Some(p) => p,
                None => continue, // Skip gradients for non-existent params
            };

            let n = param.data.len();
            if grad.numel() != n {
                return Err(MLError::DimensionMismatch {
                    expected: n,
                    actual: grad.numel(),
                });
            }

            // Lazily allocate moment buffers
            if !self.states.contains_key(name) {
                let m = self.stream.alloc_zeros::<f32>(n).map_err(|e| {
                    MLError::ModelError(format!("AdamW alloc m '{name}': {e}"))
                })?;
                let v = self.stream.alloc_zeros::<f32>(n).map_err(|e| {
                    MLError::ModelError(format!("AdamW alloc v '{name}': {e}"))
                })?;
                self.states.insert(name.clone(), ParamState { m, v });
            }

            let state = self.states.get_mut(name).ok_or_else(|| {
                MLError::ModelError(format!("AdamW state missing for '{name}'"))
            })?;

            // Launch AdamW kernel with gradient scale for clipping
            let n_i32 = n as i32;
            let threads = 256_u32;
            let blocks = ((n as u32) + threads - 1) / threads;
            let cfg = LaunchConfig {
                grid_dim: (blocks, 1, 1),
                block_dim: (threads, 1, 1),
                shared_mem_bytes: 0,
            };

            unsafe {
                self.stream
                    .launch_builder(&self.kernel)
                    .arg(&param.data)
                    .arg(&grad.data)
                    .arg(&state.m)
                    .arg(&state.v)
                    .arg(&self.config.lr)
                    .arg(&self.config.beta1)
                    .arg(&self.config.beta2)
                    .arg(&self.config.epsilon)
                    .arg(&self.config.weight_decay)
                    .arg(&grad_scale)
                    .arg(&t)
                    .arg(&n_i32)
                    .launch(cfg)
                    .map_err(|e| MLError::ModelError(format!("adamw_update '{name}': {e}")))?;
            }
        }

        Ok(grad_norm)
    }

    /// Compute global gradient L2 norm and return `(norm, scale)`.
    ///
    /// If `norm > max_grad_norm`, `scale = max_grad_norm / norm` (clamped to 1.0).
    /// The scale is passed to the AdamW kernel which applies it to each gradient
    /// element before the moment update — this avoids mutating the gradient tensors.
    fn compute_grad_norm_and_scale(
        &self,
        gradients: &BTreeMap<String, GpuTensor>,
    ) -> Result<(f32, f32), MLError> {
        // Accumulate squared norms per parameter into a scalar on GPU
        let norm_sq_buf = self.stream.alloc_zeros::<f32>(1).map_err(|e| {
            MLError::ModelError(format!("AdamW norm_sq alloc: {e}"))
        })?;

        for grad in gradients.values() {
            let n = grad.numel();
            let n_i32 = n as i32;
            let threads = 256_u32;
            let blocks = ((n as u32) + threads - 1) / threads;
            let cfg = LaunchConfig {
                grid_dim: (blocks, 1, 1),
                block_dim: (threads, 1, 1),
                shared_mem_bytes: 0,
            };

            unsafe {
                self.stream
                    .launch_builder(&self.grad_norm_kernel)
                    .arg(&grad.data)
                    .arg(&norm_sq_buf)
                    .arg(&n_i32)
                    .launch(cfg)
                    .map_err(|e| MLError::ModelError(format!("grad_l2_norm: {e}")))?;
            }
        }

        // Read back the squared norm (single scalar, 4 bytes)
        let mut norm_sq_host = [0.0_f32];
        self.stream.memcpy_dtoh(&norm_sq_buf, &mut norm_sq_host).map_err(|e| { // gpu-exit: 1 scalar grad_norm for logging
            MLError::ModelError(format!("AdamW norm readback: {e}"))
        })?;

        let norm = norm_sq_host[0].sqrt();

        // Compute scale factor: 1.0 if within budget, else max_norm / norm
        let scale = if norm > self.config.max_grad_norm {
            self.config.max_grad_norm / (norm + 1e-6)
        } else {
            1.0_f32
        };

        Ok((norm, scale))
    }
}

// ── CUDA source ──────────────────────────────────────────────────────────

const ADAMW_CUDA_SRC: &str = r#"
// ═══════════════════════════════════════════════════════════════════════
// AdamW optimizer kernel — one launch per parameter tensor.
// Implements bias-corrected AdamW with decoupled weight decay.
// ═══════════════════════════════════════════════════════════════════════

extern "C" __global__
void adamw_update(float* __restrict__ param,
                  const float* __restrict__ grad,
                  float* __restrict__ m,
                  float* __restrict__ v,
                  float lr,
                  float beta1,
                  float beta2,
                  float epsilon,
                  float weight_decay,
                  float grad_scale,
                  int t,
                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        // Apply gradient clipping scale (1.0 = no clip, <1.0 = clipped)
        float g = grad[i] * grad_scale;
        float p = param[i];

        // Decoupled weight decay: p = p * (1 - lr * wd)
        p *= (1.0f - lr * weight_decay);

        // Moment updates
        float mi = beta1 * m[i] + (1.0f - beta1) * g;
        float vi = beta2 * v[i] + (1.0f - beta2) * g * g;
        m[i] = mi;
        v[i] = vi;

        // Bias correction
        float bc1 = 1.0f - powf(beta1, (float)t);
        float bc2 = 1.0f - powf(beta2, (float)t);
        float m_hat = mi / bc1;
        float v_hat = vi / bc2;

        // Parameter update
        p -= lr * m_hat / (sqrtf(v_hat) + epsilon);
        param[i] = p;
    }
}

// Compute sum of squared gradient elements via atomicAdd.
// Called once per parameter tensor to accumulate into a global norm^2 scalar.
extern "C" __global__
void grad_l2_norm(const float* __restrict__ grad,
                  float* __restrict__ norm_sq,
                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float g = grad[i];
        atomicAdd(norm_sq, g * g);
    }
}
"#;

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::var_store::GpuVarStore;
    use super::super::gpu_tensor::GpuTensor;
    use super::super::init;

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
                .expect("CUDA required")
                .new_stream()
                .expect("fork stream")
    }

    #[test]
    fn test_adamw_basic_step() {
        let stream = make_stream();
        let mut store = GpuVarStore::new(stream.clone());
        store.linear("fc", 4, 4).unwrap();

        // Set weights to 1.0
        let ones = vec![1.0_f32; 16];
        stream.memcpy_htod(&ones, &mut store.get_mut("fc.weight").unwrap().data).unwrap();
        let bias_ones = vec![1.0_f32; 4];
        stream.memcpy_htod(&bias_ones, &mut store.get_mut("fc.bias").unwrap().data).unwrap();

        let mut optimizer = GpuAdamW::new(
            AdamWConfig { max_grad_norm: f32::INFINITY, ..Default::default() },
            stream.clone(),
        ).unwrap();

        // Create gradients (all 0.1)
        let mut grads = BTreeMap::new();
        grads.insert(
            "fc.weight".to_string(),
            GpuTensor::from_host(&vec![0.1_f32; 16], vec![4, 4], &stream).unwrap(),
        );
        grads.insert(
            "fc.bias".to_string(),
            GpuTensor::from_host(&vec![0.1_f32; 4], vec![4], &stream).unwrap(),
        );

        let _grad_norm = optimizer.step(&mut store, &grads).unwrap();

        // Weights should have decreased from 1.0
        let w_host = {
            let param = store.get("fc.weight").unwrap();
            let mut h = vec![0.0_f32; param.data.len()];
            stream.memcpy_dtoh(&param.data, &mut h).unwrap(); // test-only readback
            h
        };

        for &w in &w_host {
            assert!(w < 1.0, "Weight should decrease after AdamW step, got {w}");
            assert!(w > 0.99, "Weight should not decrease too much after 1 step, got {w}");
        }
    }

    #[test]
    fn test_adamw_step_counter() {
        let stream = make_stream();
        let mut store = GpuVarStore::new(stream.clone());
        let d = init::zeros(4, &stream).unwrap();
        store.register("x", d, vec![4]).unwrap();

        let mut optimizer = GpuAdamW::new(
            AdamWConfig::default(),
            stream.clone(),
        ).unwrap();

        assert_eq!(optimizer.step_count(), 0);

        let mut grads = BTreeMap::new();
        grads.insert(
            "x".to_string(),
            GpuTensor::from_host(&[0.0_f32; 4], vec![4], &stream).unwrap(),
        );

        optimizer.step(&mut store, &grads).unwrap();
        assert_eq!(optimizer.step_count(), 1);

        optimizer.step(&mut store, &grads).unwrap();
        assert_eq!(optimizer.step_count(), 2);
    }
}
