//! GPU-resident `AdamW` optimizer using CUDA kernels.
//!
//! Flat-param-group API over `ml_core::cuda_autograd`'s `AdamW` CUDA kernel.
//! All optimizer state (m, v, parameters) stays on GPU. The update step
//! runs as a single CUDA kernel launch per parameter group -- zero CPU
//! round-trips per optimization step.
//!
//! This module provides the `CudaAdam` struct with indexed parameter groups,
//! suitable for networks that manage weights as raw `CudaSlice<f32>` buffers
//! rather than named parameters in a `GpuVarStore`.

use std::sync::Arc;

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

use ml_core::MLError;

/// CUDA kernel source for the `AdamW` update -- identical to `ml-core`'s
/// `adamw_update` kernel (decoupled weight decay, bias-corrected moments).
/// Kept as a string constant so compilation is self-contained per optimizer
/// instance.
const ADAMW_KERNEL: &str = r#"
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,
                  int t,
                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float g = grad[i];
        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;
    }
}
"#;

/// Configuration for the CUDA `AdamW` optimizer.
#[derive(Debug, Clone)]
pub struct CudaAdamConfig {
    pub lr: f32,
    pub beta1: f32,
    pub beta2: f32,
    pub eps: f32,
    pub weight_decay: f32,
}

impl Default for CudaAdamConfig {
    fn default() -> Self {
        Self {
            lr: 3e-4,
            beta1: 0.9,
            beta2: 0.999,
            eps: 1e-8,
            weight_decay: 0.0,
        }
    }
}

/// A single parameter group managed by the CUDA `AdamW` optimizer.
pub struct AdamParamGroup {
    /// First moment (m) buffer, same size as parameter
    pub m: CudaSlice<f32>,
    /// Second moment (v) buffer, same size as parameter
    pub v: CudaSlice<f32>,
    /// Number of elements
    pub len: usize,
}

/// GPU-resident `AdamW` optimizer with flat indexed parameter groups.
///
/// Manages first and second moment buffers for each parameter group.
/// The `step_group()` method launches one CUDA kernel per parameter to update
/// weights in-place using the accumulated gradients.
pub struct CudaAdam {
    config: CudaAdamConfig,
    stream: Arc<CudaStream>,
    adam_func: cudarc::driver::CudaFunction,
    /// Parameter groups: (m, v) for each registered parameter buffer
    groups: Vec<AdamParamGroup>,
    /// Global step counter for bias correction
    step_count: u64,
}

impl std::fmt::Debug for CudaAdam {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaAdam")
            .field("config", &self.config)
            .field("num_groups", &self.groups.len())
            .field("step_count", &self.step_count)
            .finish()
    }
}

impl CudaAdam {
    /// Create a new `AdamW` optimizer.
    ///
    /// Call `add_param_group()` for each (weight, bias) pair before calling `step_group()`.
    pub fn new(stream: Arc<CudaStream>, config: CudaAdamConfig) -> Result<Self, MLError> {
        let context = stream.context();
        let ptx = ml_core::cuda_compile::compile_ptx_for_device(ADAMW_KERNEL, context)
            .map_err(|e| MLError::InitializationError {
                component: "CudaAdam".to_owned(),
                message: format!("Failed to compile adamw kernel: {e}"),
            })?;
        let module = context.load_module(ptx).map_err(|e| {
            MLError::InitializationError {
                component: "CudaAdam".to_owned(),
                message: format!("Failed to load adamw module: {e}"),
            }
        })?;
        let adam_func = module.load_function("adamw_update").map_err(|e| {
            MLError::InitializationError {
                component: "CudaAdam".to_owned(),
                message: format!("Failed to load adamw_update: {e}"),
            }
        })?;

        Ok(Self {
            config,
            stream,
            adam_func,
            groups: Vec::new(),
            step_count: 0,
        })
    }

    /// Register a parameter buffer with the optimizer.
    ///
    /// Allocates first (m) and second (v) moment buffers on GPU.
    pub fn add_param_group(&mut self, param_len: usize) -> Result<usize, MLError> {
        let m = self.stream.alloc_zeros::<f32>(param_len).map_err(|e| {
            MLError::InitializationError {
                component: "CudaAdam".to_owned(),
                message: format!("Failed to alloc m buffer: {e}"),
            }
        })?;
        let v = self.stream.alloc_zeros::<f32>(param_len).map_err(|e| {
            MLError::InitializationError {
                component: "CudaAdam".to_owned(),
                message: format!("Failed to alloc v buffer: {e}"),
            }
        })?;
        let idx = self.groups.len();
        self.groups.push(AdamParamGroup { m, v, len: param_len });
        Ok(idx)
    }

    /// Perform one `AdamW` optimization step for a specific parameter group.
    ///
    /// Updates `param` in-place using `grad`.
    pub fn step_group(
        &mut self,
        group_idx: usize,
        param: &mut CudaSlice<f32>,
        grad: &CudaSlice<f32>,
    ) -> Result<(), MLError> {
        let group = self.groups.get_mut(group_idx).ok_or_else(|| {
            MLError::ConfigError(format!("Adam group index {} out of range", group_idx))
        })?;

        let t = (self.step_count + 1) as i32;

        let n = group.len as i32;
        let block_size = 256_u32;
        let grid_size = (group.len as u32).div_ceil(block_size);

        // SAFETY: adamw_update kernel reads `n` gradient elements and updates `n`
        // parameter, momentum (m), and velocity (v) elements in-place. All
        // buffers are registered with matching sizes during `add_param_group`.
        unsafe {
            self.stream
                .launch_builder(&self.adam_func)
                .arg(param)
                .arg(grad)
                .arg(&mut group.m)
                .arg(&mut group.v)
                .arg(&self.config.lr)
                .arg(&self.config.beta1)
                .arg(&self.config.beta2)
                .arg(&self.config.eps)
                .arg(&self.config.weight_decay)
                .arg(&t)
                .arg(&n)
                .launch(LaunchConfig {
                    grid_dim: (grid_size, 1, 1),
                    block_dim: (block_size, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::TrainingError(format!("adamw_update launch failed: {e}")))?;
        }

        Ok(())
    }

    /// Increment the global step counter. Call once per full optimizer step.
    pub const fn increment_step(&mut self) {
        self.step_count += 1;
    }

    /// Get the current step count.
    pub const fn step_count(&self) -> u64 {
        self.step_count
    }

    /// Get the learning rate.
    pub const fn lr(&self) -> f32 {
        self.config.lr
    }

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