//! AdamW optimizer state + kernel binding.
//!
//! Uses a device-resident step counter (`step_count_d`) instead of a
//! host scalar. Required for CUDA Graph capture: host scalars passed
//! via `.arg()` are baked into kernel args at capture time, so a
//! replayed graph would freeze the step counter and produce wrong
//! bias-correction values. The counter is incremented via a tiny
//! single-thread kernel launched after the parameter update — both
//! launches go through the same stream and end up inside any captured
//! region the caller brackets.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use ml_core::device::MlDevice;
use crate::pinned_mem::MappedI32Buffer;

const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/adamw_step.cubin"));

pub struct AdamW {
    stream: Arc<CudaStream>,
    _module: Arc<CudaModule>,
    func: CudaFunction,
    inc_func: CudaFunction,
    m: CudaSlice<f32>,
    v: CudaSlice<f32>,
    /// Device-resident step counter (i32). Initialized to 0; the
    /// increment kernel advances it after each parameter update, so the
    /// first call observes step=1 (matching standard Adam convention).
    step_count_d: CudaSlice<i32>,
    pub lr: f32,
    pub beta1: f32,
    pub beta2: f32,
    pub eps: f32,
    pub wd: f32,
}

impl AdamW {
    pub fn new(dev: &MlDevice, n_params: usize, lr: f32) -> Result<Self> {
        let stream = dev.cuda_stream().context("adamw stream")?.clone();
        let ctx = dev.cuda_context().context("adamw ctx")?;
        let module = ctx.load_cubin(CUBIN.to_vec()).context("load adamw cubin")?;
        let func = module.load_function("adamw_step").context("adamw fn")?;
        let inc_func = module
            .load_function("adamw_increment_counter")
            .context("adamw_increment_counter fn")?;
        let m = stream.alloc_zeros::<f32>(n_params).context("m alloc")?;
        let v = stream.alloc_zeros::<f32>(n_params).context("v alloc")?;
        // Step counter starts at 0; the increment kernel inside step()
        // advances it to 1 BEFORE the first read by the adamw_step
        // kernel — see step() for the ordering.
        let step_count_d = stream
            .alloc_zeros::<i32>(1)
            .context("step_count_d alloc")?;
        Ok(Self {
            stream,
            _module: module,
            func,
            inc_func,
            m,
            v,
            step_count_d,
            lr,
            beta1: 0.9,
            beta2: 0.999,
            eps: 1e-8,
            wd: 1e-2,
        })
    }

    /// Apply one AdamW update. Launches `adamw_increment_counter` first
    /// (advances device step counter from N-1 to N), then `adamw_step`
    /// which reads the counter and uses N for bias correction. Both
    /// launches are stream-ordered on the trainer's stream — they end
    /// up inside any CUDA Graph capture that brackets the call.
    pub fn step(&mut self, theta: &mut CudaSlice<f32>, grad: &CudaSlice<f32>) -> Result<()> {
        let n = theta.len();
        assert_eq!(grad.len(), n, "grad/theta size mismatch");
        assert_eq!(self.m.len(), n, "m size mismatch");
        assert_eq!(self.v.len(), n, "v size mismatch");

        // 1. Advance device step counter (was N-1, becomes N for this update).
        let inc_cfg = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (1, 1, 1),
            shared_mem_bytes: 0,
        };
        {
            let mut inc = self.stream.launch_builder(&self.inc_func);
            inc.arg(&mut self.step_count_d);
            unsafe { inc.launch(inc_cfg).context("adamw increment launch")?; }
        }

        // 2. Apply update — kernel reads step from device pointer.
        let n_params_i = n as i32;
        let cfg = LaunchConfig {
            grid_dim: ((n as u32).div_ceil(256), 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        let mut launch = self.stream.launch_builder(&self.func);
        launch
            .arg(theta).arg(grad).arg(&mut self.m).arg(&mut self.v)
            .arg(&n_params_i)
            .arg(&self.lr).arg(&self.beta1).arg(&self.beta2).arg(&self.eps).arg(&self.wd)
            .arg(&self.step_count_d);
        unsafe { launch.launch(cfg).context("adamw launch")?; }
        // No per-call sync — caller batches multiple AdamW launches on the
        // same stream and synchronises once at step boundary.
        Ok(())
    }

    /// Read the current device-resident step counter. Forces a stream
    /// sync and a mapped-pinned readback — NOT suitable for the hot
    /// training path. Intended for tests and one-off diagnostics; the
    /// counter is normally observed only indirectly through training
    /// loss trajectory.
    pub fn step_count(&self) -> i32 {
        let staging = unsafe { MappedI32Buffer::new(1) }
            .expect("MappedI32Buffer alloc");
        unsafe {
            let (src_ptr, _g) = self.step_count_d.device_ptr(&self.stream);
            cudarc::driver::result::memcpy_dtod_async(
                staging.dev_ptr, src_ptr, 4, self.stream.cu_stream(),
            ).expect("step_count dtod");
        }
        self.stream.synchronize().expect("step_count sync");
        unsafe { std::ptr::read_volatile(staging.host_ptr) }
    }
}
