//! AdamW optimizer state + kernel binding.
//!
//! Uses a host-side step counter passed as an i32 kernel argument. The
//! Rust `step()` method increments the counter before each launch, so
//! the first call uses step=1 (matching standard Adam convention).
//! This eliminates the former device-resident counter and its separate
//! `adamw_increment_counter` kernel launch, saving 3407 kernel
//! launches per training step.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
use cudarc::driver::sys::CUstream;
use ml_core::device::MlDevice;
use crate::trainer::raw_launch::{RawArgs, raw_launch};

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

pub struct AdamW {
    _stream: Arc<CudaStream>,
    raw_stream: CUstream,
    _module: Arc<CudaModule>,
    func: CudaFunction,
    m: CudaSlice<f32>,
    v: CudaSlice<f32>,
    /// Host-side step counter. Starts at 0; `step()` increments to 1
    /// before the first kernel launch, matching standard Adam convention.
    step_count_host: 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 m = stream.alloc_zeros::<f32>(n_params).context("m alloc")?;
        let v = stream.alloc_zeros::<f32>(n_params).context("v alloc")?;
        let raw_stream = stream.cu_stream();
        Ok(Self {
            _stream: stream,
            raw_stream,
            _module: module,
            func,
            m,
            v,
            step_count_host: 0,
            lr,
            beta1: 0.9,
            beta2: 0.999,
            eps: 1e-8,
            wd: 1e-2,
        })
    }

    /// Apply one AdamW update. Increments the host-side step counter
    /// from N-1 to N, then launches `adamw_step` which receives N as
    /// a scalar argument for bias correction.
    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");

        // Advance step counter (was N-1, becomes N for this update).
        self.step_count_host += 1;

        let n_params_i = n as i32;
        let grid_x = (n as u32).div_ceil(256);
        {
            let mut args = RawArgs::new();
            args.push_ptr(theta.raw_ptr());
            args.push_ptr(grad.raw_ptr());
            args.push_ptr(self.m.raw_ptr());
            args.push_ptr(self.v.raw_ptr());
            args.push_i32(n_params_i);
            args.push_f32(self.lr);
            args.push_f32(self.beta1);
            args.push_f32(self.beta2);
            args.push_f32(self.eps);
            args.push_f32(self.wd);
            args.push_i32(self.step_count_host);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.func.cu_function(),
                    (grid_x, 1, 1), (256, 1, 1), 0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("adamw launch: {:?}", e))?;
            }
        }
        // No per-call sync — caller batches multiple AdamW launches on the
        // same stream and synchronises once at step boundary.
        Ok(())
    }

    /// Mutable access to Adam's first-moment buffer `m`. ALPHA fix
    /// 2026-05-21: used by the Phase 1→2 transition to project
    /// `heads_w_skip` Adam state onto the block-diagonal subspace —
    /// without this, pre-transition gradient momentum would re-introduce
    /// off-bucket signal across the transition (Bug 3 from the
    /// checkpoint deep-dive). Per `pearl_adam_normalizes_loss_weights`,
    /// Adam moments are the actual state to constrain when projecting.
    pub fn m_mut(&mut self) -> &mut CudaSlice<f32> {
        &mut self.m
    }

    /// Mutable access to Adam's second-moment buffer `v`. Companion to
    /// `m_mut`; the projection kernel zeroes both together at transition
    /// so the subsequent Adam step from zero moments yields a zero update
    /// (modulo the ε term, which the per-step belt-and-suspenders
    /// zero_off_bucket kernel cleans up).
    pub fn v_mut(&mut self) -> &mut CudaSlice<f32> {
        &mut self.v
    }

    /// Immutable access to Adam's first-moment buffer `m`. Used by the
    /// determinism Phase 1 checksum diagnostic
    /// (`docs/superpowers/specs/2026-06-02-determinism-foundation.md` §1.2)
    /// to take a sum-of-squares snapshot post-Adam-step without
    /// needing `&mut self`.
    pub fn m(&self) -> &CudaSlice<f32> {
        &self.m
    }

    /// Immutable access to Adam's second-moment buffer `v`. Companion to
    /// `m()`, same use case (determinism Phase 1 checksums).
    pub fn v(&self) -> &CudaSlice<f32> {
        &self.v
    }

    /// Return the current step counter value. No GPU sync needed —
    /// the counter is host-resident.
    pub fn step_count(&self) -> i32 {
        self.step_count_host
    }
}
