//! Wave 5b.1 iter 2 — recompute_w_u_fwd kernel test.
//!
//! Spec source:
//! - FLA reference: `recompute_w_u_fwd_kernel` at
//!   /opt/vllm/vllm/model_executor/layers/fla/ops/wy_fast.py:29-117
//!
//! The reference fixture is generated by:
//! ```sh
//! python3 tests/fixtures/gated_delta_net_recompute_wu_reference.py
//! ```
//!
//! Tolerance bar: `atol = 5e-3, rtol = 2e-2` for bf16 outputs.

#![allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]

use std::fs;
use std::path::PathBuf;

use mlx_native::ops::gated_delta_net_recompute_wu::{
    build_gated_delta_net_recompute_wu_params, dispatch_gated_delta_net_recompute_wu,
    GatedDeltaNetRecomputeWuParams,
};
use mlx_native::{DType, KernelRegistry, MlxBuffer, MlxDevice};

const B: u32 = 1;
const T: u32 = 128;
const HG: u32 = 2;
const H: u32 = 4;
const K: u32 = 128;
const V: u32 = 128;
const BT: u32 = 64;

fn fixture_dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures")
}

fn read_bytes(name: &str) -> Vec<u8> {
    let path = fixture_dir().join(name);
    fs::read(&path).unwrap_or_else(|e| {
        panic!(
            "failed to read fixture {} — did you run `python3 tests/fixtures/gated_delta_net_recompute_wu_reference.py`? ({})",
            path.display(),
            e
        )
    })
}

fn read_f32(name: &str) -> Vec<f32> {
    let bytes = read_bytes(name);
    assert!(bytes.len() % 4 == 0, "f32 byte length not multiple of 4");
    let mut out = Vec::with_capacity(bytes.len() / 4);
    for chunk in bytes.chunks_exact(4) {
        out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
    }
    out
}

fn bf16_bytes_to_f32(bytes: &[u8]) -> Vec<f32> {
    assert!(bytes.len() % 2 == 0, "bf16 byte length not multiple of 2");
    let mut out = Vec::with_capacity(bytes.len() / 2);
    for chunk in bytes.chunks_exact(2) {
        let bf16 = u16::from_le_bytes([chunk[0], chunk[1]]);
        let bits = (bf16 as u32) << 16;
        out.push(f32::from_bits(bits));
    }
    out
}

fn upload_bf16(device: &MlxDevice, bytes: &[u8]) -> MlxBuffer {
    let n_elems = bytes.len() / 2;
    let mut buf = device
        .alloc_buffer(bytes.len(), DType::BF16, vec![n_elems])
        .expect("alloc bf16");
    let dst = buf.as_mut_slice::<u8>().expect("mut bytes");
    dst.copy_from_slice(bytes);
    buf
}

fn upload_f32(device: &MlxDevice, data: &[f32]) -> MlxBuffer {
    let mut buf = device
        .alloc_buffer(data.len() * 4, DType::F32, vec![data.len()])
        .expect("alloc f32");
    buf.as_mut_slice::<f32>()
        .expect("mut")
        .copy_from_slice(data);
    buf
}

fn alloc_zeroed_bf16(device: &MlxDevice, n_elems: usize) -> MlxBuffer {
    let mut buf = device
        .alloc_buffer(n_elems * 2, DType::BF16, vec![n_elems])
        .expect("alloc bf16");
    let dst = buf.as_mut_slice::<u8>().expect("mut bytes");
    dst.fill(0u8);
    buf
}

#[test]
fn test_recompute_wu_matches_fla_reference() {
    let device = MlxDevice::new().expect("MlxDevice::new");
    let mut registry = KernelRegistry::new();

    // Inputs.
    let k_bytes = read_bytes("gated_delta_net_recompute_wu_input_k.bin");
    let v_bytes = read_bytes("gated_delta_net_recompute_wu_input_v.bin");
    let beta_vals = read_f32("gated_delta_net_recompute_wu_input_beta.bin");
    let g_vals = read_f32("gated_delta_net_recompute_wu_input_g.bin");
    let a_vals = read_f32("gated_delta_net_recompute_wu_input_A.bin");

    let exp_k = (B * T * HG * K) as usize;
    let exp_v = (B * T * H * V) as usize;
    let exp_beta = (B * T * H) as usize;
    let exp_g = (B * T * H) as usize;
    let exp_a = (B * T * H * BT) as usize;
    assert_eq!(k_bytes.len(), exp_k * 2, "k size mismatch");
    assert_eq!(v_bytes.len(), exp_v * 2, "v size mismatch");
    assert_eq!(beta_vals.len(), exp_beta, "beta size mismatch");
    assert_eq!(g_vals.len(), exp_g, "g size mismatch");
    assert_eq!(a_vals.len(), exp_a, "A size mismatch");

    let k_buf = upload_bf16(&device, &k_bytes);
    let v_buf = upload_bf16(&device, &v_bytes);
    let beta_buf = upload_f32(&device, &beta_vals);
    let g_buf = upload_f32(&device, &g_vals);
    let a_buf = upload_f32(&device, &a_vals);

    // Outputs.
    let w_elems = (B * T * H * K) as usize;
    let u_elems = (B * T * H * V) as usize;
    let w_buf = alloc_zeroed_bf16(&device, w_elems);
    let u_buf = alloc_zeroed_bf16(&device, u_elems);

    let p = GatedDeltaNetRecomputeWuParams {
        b: B,
        t: T,
        hg: HG,
        h: H,
        k: K,
        v: V,
        bt: BT,
    };
    let params_buf = build_gated_delta_net_recompute_wu_params(&device, p).expect("params");

    let mut enc = device.command_encoder().expect("enc");
    dispatch_gated_delta_net_recompute_wu(
        &mut enc,
        &mut registry,
        device.metal_device(),
        &k_buf,
        &v_buf,
        &beta_buf,
        &g_buf,
        &a_buf,
        &w_buf,
        &u_buf,
        &params_buf,
        p,
    )
    .expect("dispatch");
    enc.commit_and_wait().expect("commit");

    let w_bytes = w_buf.as_slice::<u8>().expect("read w").to_vec();
    let u_bytes = u_buf.as_slice::<u8>().expect("read u").to_vec();
    let got_w = bf16_bytes_to_f32(&w_bytes);
    let got_u = bf16_bytes_to_f32(&u_bytes);

    let ref_w_bytes = read_bytes("gated_delta_net_recompute_wu_w_ref.bin");
    let ref_u_bytes = read_bytes("gated_delta_net_recompute_wu_u_ref.bin");
    let ref_w = bf16_bytes_to_f32(&ref_w_bytes);
    let ref_u = bf16_bytes_to_f32(&ref_u_bytes);

    assert_eq!(got_w.len(), ref_w.len(), "w length mismatch");
    assert_eq!(got_u.len(), ref_u.len(), "u length mismatch");

    let atol: f32 = 5e-3;
    let rtol: f32 = 2e-2;
    let mut max_w_err = 0.0f32;
    let mut max_u_err = 0.0f32;

    for (i, (&g, &r)) in got_w.iter().zip(ref_w.iter()).enumerate() {
        let abs_err = (g - r).abs();
        let rel_err = abs_err / r.abs().max(f32::EPSILON);
        max_w_err = max_w_err.max(abs_err);
        if abs_err > atol && rel_err > rtol {
            panic!(
                "w[{}]: gpu={} ref={} abs_err={} rel_err={}",
                i, g, r, abs_err, rel_err
            );
        }
    }
    for (i, (&g, &r)) in got_u.iter().zip(ref_u.iter()).enumerate() {
        let abs_err = (g - r).abs();
        let rel_err = abs_err / r.abs().max(f32::EPSILON);
        max_u_err = max_u_err.max(abs_err);
        if abs_err > atol && rel_err > rtol {
            panic!(
                "u[{}]: gpu={} ref={} abs_err={} rel_err={}",
                i, g, r, abs_err, rel_err
            );
        }
    }

    eprintln!(
        "recompute_w_u OK   max_w_err={:.3e}   max_u_err={:.3e}",
        max_w_err, max_u_err
    );
}

#[test]
fn test_recompute_wu_zero_inputs_yield_zero() {
    let device = MlxDevice::new().expect("MlxDevice::new");
    let mut registry = KernelRegistry::new();

    let k_elems = (B * T * HG * K) as usize;
    let v_elems = (B * T * H * V) as usize;
    let beta_elems = (B * T * H) as usize;
    let g_elems = (B * T * H) as usize;
    let a_elems = (B * T * H * BT) as usize;
    let w_elems = (B * T * H * K) as usize;
    let u_elems = (B * T * H * V) as usize;

    let k_buf = alloc_zeroed_bf16(&device, k_elems);
    let v_buf = alloc_zeroed_bf16(&device, v_elems);
    let beta_buf = upload_f32(&device, &vec![0.0f32; beta_elems]);
    let g_buf = upload_f32(&device, &vec![0.0f32; g_elems]);
    // A: identity per chunk would be nonzero, but zero A also produces zero
    // outputs (since dot(zero, anything) = zero). Use zero for simplicity.
    let a_buf = upload_f32(&device, &vec![0.0f32; a_elems]);

    let w_buf = alloc_zeroed_bf16(&device, w_elems);
    let u_buf = alloc_zeroed_bf16(&device, u_elems);

    let p = GatedDeltaNetRecomputeWuParams {
        b: B,
        t: T,
        hg: HG,
        h: H,
        k: K,
        v: V,
        bt: BT,
    };
    let params_buf = build_gated_delta_net_recompute_wu_params(&device, p).expect("params");

    let mut enc = device.command_encoder().expect("enc");
    dispatch_gated_delta_net_recompute_wu(
        &mut enc,
        &mut registry,
        device.metal_device(),
        &k_buf,
        &v_buf,
        &beta_buf,
        &g_buf,
        &a_buf,
        &w_buf,
        &u_buf,
        &params_buf,
        p,
    )
    .expect("dispatch");
    enc.commit_and_wait().expect("commit");

    let w_bytes = w_buf.as_slice::<u8>().expect("read w");
    let u_bytes = u_buf.as_slice::<u8>().expect("read u");
    assert!(
        w_bytes.iter().all(|&b| b == 0),
        "w should be all-zero with zero A"
    );
    assert!(
        u_bytes.iter().all(|&b| b == 0),
        "u should be all-zero with zero A"
    );
}
