//! Hand-written AdamW optimizer (Phase T6).
//!
//! AdamW = Adam with **decoupled** weight decay (Loshchilov & Hutter, 2019): the
//! weight-decay term is applied directly to the parameter, NOT folded into the
//! gradient (so it does not interact with the adaptive `v` denominator). This
//! matches `torch.optim.AdamW`.
//!
//! Update for parameter `θ` at step `t` (1-indexed), with gradient `g`:
//! ```text
//!   m ← β1·m + (1−β1)·g
//!   v ← β2·v + (1−β2)·g²
//!   m̂ ← m / (1 − β1ᵗ)          (bias correction)
//!   v̂ ← v / (1 − β2ᵗ)
//!   θ ← θ − lr·( m̂ / (√v̂ + ε) + wd·θ )
//! ```
//! The `lr·wd·θ` term is the decoupled decay. Note PyTorch applies decay as
//! `θ ← θ·(1 − lr·wd)` then the Adam step; both are algebraically the same
//! first-order update — we fold decay into the single subtraction above, which
//! is what PyTorch's default (`maximize=False`, no `amsgrad`) computes.
//!
//! The math operates on flat host `f32` buffers ([`AdamW::step_host`]) so it is
//! unit-testable on a GPU-less host; [`AdamW::step`] is a thin wrapper that
//! round-trips each parameter's value/grad through the GPU tensor and is gated
//! behind `not(no_cuda)`.

/// Per-parameter optimizer state: the first (`m`) and second (`v`) moment
/// estimates, one f32 per element, kept flat (matching the parameter layout).
struct ParamState {
    m: Vec<f32>,
    v: Vec<f32>,
}

/// Decoupled-weight-decay Adam. One instance owns the moment state for a fixed
/// list of parameters, keyed by their index in the slice passed to `step`
/// (the model's stable `params()` order).
pub struct AdamW {
    pub lr: f32,
    beta1: f32,
    beta2: f32,
    eps: f32,
    weight_decay: f32,
    /// Global step count (shared across all params for bias correction).
    t: u64,
    /// Lazily sized to the parameter list on the first `step`.
    state: Vec<ParamState>,
}

impl AdamW {
    /// PyTorch-default hyperparameters except `lr`/`weight_decay`, which you set
    /// (β1=0.9, β2=0.999, ε=1e-8).
    pub fn new(lr: f32, weight_decay: f32) -> Self {
        Self::with_betas(lr, weight_decay, 0.9, 0.999, 1e-8)
    }

    pub fn with_betas(lr: f32, weight_decay: f32, beta1: f32, beta2: f32, eps: f32) -> Self {
        Self {
            lr,
            beta1,
            beta2,
            eps,
            weight_decay,
            t: 0,
            state: Vec::new(),
        }
    }

    /// Current global step (number of `step` calls so far).
    pub fn step_count(&self) -> u64 {
        self.t
    }

    /// Pure-host AdamW step over flat parameter/gradient buffers. `params[i]` is
    /// updated in place using `grads[i]`; both are the i-th parameter's elements
    /// in the model's stable order. Lazily allocates moment state on first call.
    ///
    /// This is the testable core — no GPU, no autograd. `lr` is passed per call
    /// so a schedule can vary it each step.
    pub fn step_host(&mut self, lr: f32, params: &mut [Vec<f32>], grads: &[Vec<f32>]) {
        assert_eq!(params.len(), grads.len(), "param/grad count mismatch");
        if self.state.is_empty() {
            self.state = params
                .iter()
                .map(|p| ParamState {
                    m: vec![0.0; p.len()],
                    v: vec![0.0; p.len()],
                })
                .collect();
        }
        assert_eq!(self.state.len(), params.len(), "param count changed");

        self.t += 1;
        let bc1 = 1.0 - self.beta1.powi(self.t as i32);
        let bc2 = 1.0 - self.beta2.powi(self.t as i32);

        for (i, (p, g)) in params.iter_mut().zip(grads).enumerate() {
            assert_eq!(p.len(), g.len(), "param/grad len mismatch at {i}");
            let st = &mut self.state[i];
            for j in 0..p.len() {
                let gj = g[j];
                st.m[j] = self.beta1 * st.m[j] + (1.0 - self.beta1) * gj;
                st.v[j] = self.beta2 * st.v[j] + (1.0 - self.beta2) * gj * gj;
                let mhat = st.m[j] / bc1;
                let vhat = st.v[j] / bc2;
                // Decoupled weight decay: decay term uses the *current* param,
                // matching PyTorch's `p ← p − lr·wd·p` applied alongside the step.
                p[j] -= lr * (mhat / (vhat.sqrt() + self.eps) + self.weight_decay * p[j]);
            }
        }
    }
}

#[cfg(not(no_cuda))]
mod gpu {
    use super::AdamW;
    use xtrain_autodiff::tape::Var;
    use xtrain_tensor::{DType, Device, Tensor};

    impl AdamW {
        /// Apply one AdamW step to every parameter `Var`, using `lr` for this step
        /// (so an LR schedule can vary it). Pulls each param's value and `.grad()`
        /// to the host, runs [`AdamW::step_host`], and writes the updated value
        /// back with `set_value`. A param with no grad is fed a zero grad, so the
        /// Adam term vanishes and only decoupled weight decay applies (the model's
        /// params all receive grads each step, so this is just a safety default).
        ///
        /// Does NOT zero grads — the caller does that (matching the GD-step
        /// template in the T5 overfit test).
        ///
        /// This is the host-roundtrip reference path; training uses
        /// [`GpuAdamW`] (kernel, m/v on device). Both are checked against the
        /// torch parity in tests.
        pub fn step(&mut self, lr: f32, params: &[Var]) {
            let device = params[0].value().device();
            let shapes: Vec<Vec<usize>> =
                params.iter().map(|p| p.value().shape().to_vec()).collect();

            let mut host_params: Vec<Vec<f32>> = params
                .iter()
                .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
                .collect();
            let host_grads: Vec<Vec<f32>> = params
                .iter()
                .zip(&host_params)
                .map(|(p, hp)| match p.grad() {
                    Some(g) => g.to_device(Device::Cpu).as_slice::<f32>().to_vec(),
                    None => vec![0.0; hp.len()], // no grad → no update this step
                })
                .collect();

            self.step_host(lr, &mut host_params, &host_grads);

            for ((p, data), shape) in params.iter().zip(&host_params).zip(&shapes) {
                let t = Tensor::from_slice(data, shape).to_device(device);
                p.set_value(t);
            }
        }
    }

    /// GPU AdamW (Phase T7): the optimizer state (m/v moments) lives on the device
    /// as one tensor pair per parameter, and the update runs as a CUDA kernel that
    /// reads each param's `.grad()` and rewrites the param buffer in place — no
    /// per-step GPU↔host roundtrip of params/grads. Same math as
    /// [`AdamW::step_host`] (the parity reference).
    pub struct GpuAdamW {
        beta1: f32,
        beta2: f32,
        eps: f32,
        weight_decay: f32,
        t: u64,
        /// Per-parameter (m, v) device buffers, sized lazily on first step.
        state: Vec<(Tensor, Tensor)>,
    }

    impl GpuAdamW {
        /// PyTorch-default betas/eps; you set lr (per-step) + weight decay.
        pub fn new(weight_decay: f32) -> Self {
            Self {
                beta1: 0.9,
                beta2: 0.999,
                eps: 1e-8,
                weight_decay,
                t: 0,
                state: Vec::new(),
            }
        }

        pub fn step_count(&self) -> u64 {
            self.t
        }

        /// One in-place AdamW step over every parameter `Var` at learning rate
        /// `lr`. Updates the param value buffer and the device m/v state via the
        /// `adamw_step_f32` kernel. Params are mutated in place, so the leaf `Var`
        /// identities stay stable across steps (no `set_value`). Does NOT zero
        /// grads — the caller does. A param without a grad is skipped this step.
        pub fn step(&mut self, lr: f32, params: &[Var]) {
            let device = params[0].value().device();
            if self.state.is_empty() {
                self.state = params
                    .iter()
                    .map(|p| {
                        let shape = p.value().shape().to_vec();
                        (
                            Tensor::zeros(&shape, DType::F32, device),
                            Tensor::zeros(&shape, DType::F32, device),
                        )
                    })
                    .collect();
            }
            assert_eq!(self.state.len(), params.len(), "param count changed");

            self.t += 1;
            let bc1 = 1.0 - self.beta1.powi(self.t as i32);
            let bc2 = 1.0 - self.beta2.powi(self.t as i32);

            for (p, (m, v)) in params.iter().zip(&self.state) {
                let g = match p.grad() {
                    Some(g) => g,
                    None => continue,
                };
                let pv = p.value();
                let n = pv.numel() as i32;
                unsafe {
                    xtrain_cuda::ffi::launch_adamw_step_f32(
                        pv.data_ptr() as *mut f32,
                        g.data_ptr() as *const f32,
                        m.data_ptr() as *mut f32,
                        v.data_ptr() as *mut f32,
                        lr,
                        self.beta1,
                        self.beta2,
                        self.eps,
                        self.weight_decay,
                        bc1,
                        bc2,
                        n,
                        std::ptr::null_mut(),
                    );
                }
            }
            xtrain_cuda::device::synchronize().expect("adamw step sync failed");
        }
    }
}

#[cfg(not(no_cuda))]
pub use gpu::GpuAdamW;