//! AdamW invariants — element-wise correctness without a CPU mirror.

use approx::assert_relative_eq;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
use ml_alpha::pinned_mem::MappedF32Buffer;
use ml_alpha::trainer::optim::AdamW;
use ml_core::device::MlDevice;
use std::sync::Arc;

fn test_device() -> MlDevice {
    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
}

fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> CudaSlice<f32> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }.unwrap();
    staging.write_from_slice(host);
    let mut dst = stream.alloc_zeros::<f32>(n).unwrap();
    let nbytes = n * 4;
    unsafe {
        let (dst_ptr, _g) = dst.device_ptr_mut(stream);
        cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream()).unwrap();
    }
    dst
}

fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Vec<f32> {
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }.unwrap();
    let nbytes = n * 4;
    unsafe {
        let (src_ptr, _g) = src.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream()).unwrap();
    }
    stream.synchronize().unwrap();
    staging.read_all()
}

#[test]
fn zero_grad_moves_parameter_only_by_weight_decay() {
    let dev = test_device();
    let stream = dev.cuda_stream().unwrap().clone();
    let theta_init: Vec<f32> = vec![1.0, 2.0, -3.0, 0.5];
    let mut theta = upload(&stream, &theta_init);
    let grad = upload(&stream, &vec![0.0; 4]);
    let mut opt = AdamW::new(&dev, 4, 1e-3).unwrap();
    opt.wd = 0.1;
    opt.step(&mut theta, &grad).unwrap();
    let theta_after = download(&stream, &theta);
    // With g=0, m=v=0 too, so the Adam update is 0; only weight decay
    // applies: theta -= lr * wd * theta = theta * (1 - lr*wd)
    let factor = 1.0_f32 - 1e-3 * 0.1;
    for (a, b) in theta_after.iter().zip(&theta_init) {
        assert_relative_eq!(*a, *b * factor, epsilon = 1e-6);
    }
}

#[test]
fn positive_grad_decreases_parameter() {
    let dev = test_device();
    let stream = dev.cuda_stream().unwrap().clone();
    let theta_init: Vec<f32> = vec![5.0, 3.0, 1.0, 0.1];
    let mut theta = upload(&stream, &theta_init);
    let grad = upload(&stream, &vec![1.0, 1.0, 1.0, 1.0]);
    let mut opt = AdamW::new(&dev, 4, 1e-2).unwrap();
    opt.wd = 0.0; // isolate the Adam update
    opt.step(&mut theta, &grad).unwrap();
    let theta_after = download(&stream, &theta);
    for (a, b) in theta_after.iter().zip(&theta_init) {
        assert!(*a < *b, "positive grad must reduce parameter: {b} -> {a}");
    }
}

#[test]
fn adam_step_count_increments() {
    let dev = test_device();
    let stream = dev.cuda_stream().unwrap().clone();
    let mut theta = upload(&stream, &vec![1.0; 8]);
    let grad = upload(&stream, &vec![0.1; 8]);
    let mut opt = AdamW::new(&dev, 8, 1e-3).unwrap();
    assert_eq!(opt.step_count(), 0);
    opt.step(&mut theta, &grad).unwrap();
    assert_eq!(opt.step_count(), 1);
    opt.step(&mut theta, &grad).unwrap();
    assert_eq!(opt.step_count(), 2);
}

#[test]
fn repeated_grad_descent_drives_parameter_toward_zero() {
    // grad = theta (gradient of 0.5 * theta^2 wrt theta), repeated AdamW steps
    // must reduce |theta| toward 0 over many steps.
    let dev = test_device();
    let stream = dev.cuda_stream().unwrap().clone();
    let initial: Vec<f32> = vec![10.0; 16];
    let mut theta = upload(&stream, &initial);
    let mut opt = AdamW::new(&dev, 16, 1e-1).unwrap();
    opt.wd = 0.0;
    for _ in 0..200 {
        let current = download(&stream, &theta);
        let grad = upload(&stream, &current);
        opt.step(&mut theta, &grad).unwrap();
    }
    let final_theta = download(&stream, &theta);
    let norm: f32 = final_theta.iter().map(|v| v * v).sum::<f32>().sqrt();
    assert!(norm < 1.0, "‖θ‖ should drop well below initial 40 after 200 steps, got {norm}");
}
