// GPU AdamW parity (Phase T7): the device-side AdamW kernel (m/v on device, no
// host roundtrip) must produce the same update as the host reference
// `AdamW::step_host` given identical params + grads across several steps with a
// varying lr. This is the new correctness gate for the GPU optimizer; the host
// path itself is already pinned to PyTorch by xtrain-train's adamw_parity test.
//
// Gated #![cfg(not(no_cuda))] (runs on dash5; needs a GPU to link + launch).
#![cfg(not(no_cuda))]

use xtrain_autodiff::tape::Var;
use xtrain_cuda::device;
use xtrain_optim::{AdamW, GpuAdamW};
use xtrain_tensor::{Device, Tensor};

fn grad(step: usize, idx: usize, j: usize) -> f32 {
    let s = (step * 13 + idx * 7 + j * 3) as f32;
    (s * 0.123).sin() * 0.5
}

#[test]
fn gpu_adamw_matches_host() {
    assert!(device::device_count().unwrap() > 0, "no CUDA device");
    device::set_device(0).unwrap();
    let dev = Device::Cuda(0);

    let wd = 0.1f32;
    // Two params of different sizes (exercises per-param device state).
    let shapes: Vec<Vec<usize>> = vec![vec![2, 2], vec![3]];
    let init: Vec<Vec<f32>> = vec![vec![0.5, -1.0, 2.0, 0.0], vec![1.5, -0.25, 0.75]];

    // GPU side: leaf Vars on device.
    let params: Vec<Var> = init
        .iter()
        .zip(&shapes)
        .map(|(d, s)| Var::leaf(Tensor::from_slice(d, s).to_device(dev)))
        .collect();
    let mut gpu_opt = GpuAdamW::new(wd);

    // Host reference.
    let mut host_params = init.clone();
    let mut host_opt = AdamW::new(0.0, wd);

    for step in 0..15 {
        let lr = 0.01 + 0.001 * step as f32; // varying lr
        let grads: Vec<Vec<f32>> = shapes
            .iter()
            .enumerate()
            .map(|(idx, s)| {
                let n: usize = s.iter().product();
                (0..n).map(|j| grad(step, idx, j)).collect()
            })
            .collect();

        // Push grads onto the GPU Vars, run the device step, then clear.
        for (p, (g, s)) in params.iter().zip(grads.iter().zip(&shapes)) {
            p.zero_grad();
            Var::push_grad(p, Tensor::from_slice(g, s).to_device(dev));
        }
        gpu_opt.step(lr, &params);
        for p in &params {
            p.zero_grad();
        }

        host_opt.step_host(lr, &mut host_params, &grads);
    }

    let mut max_err = 0.0f32;
    for (p, hp) in params.iter().zip(&host_params) {
        let got = p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec();
        for (a, b) in got.iter().zip(hp) {
            max_err = max_err.max((a - b).abs());
        }
    }
    println!("gpu vs host AdamW: max abs err = {max_err:.3e}");
    assert!(max_err < 1e-6, "GPU AdamW diverged from host: {max_err:e}");
}