//! Wave 5b.1 iter 3 — chunk_fwd_o kernel test.
//!
//! Spec source:
//! - FLA reference: `chunk_fwd_kernel_o` at
//!   /opt/vllm/vllm/model_executor/layers/fla/ops/chunk_o.py:42-138
//!
//! The reference fixture is generated by:
//! ```sh
//! python3 tests/fixtures/gated_delta_net_chunk_o_reference.py
//! ```
//!
//! Tolerance bar: `atol = 5e-3, rtol = 2e-2` for bf16 outputs (matches
//! iter-1 + iter-2 fixtures).

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

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

use mlx_native::ops::gated_delta_net_chunk_o::{
    build_gated_delta_net_chunk_o_params, dispatch_gated_delta_net_chunk_o,
    GatedDeltaNetChunkOParams,
};
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;
const NT: u32 = T / BT;

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_chunk_o_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_chunk_fwd_o_matches_fla_reference() {
    let device = MlxDevice::new().expect("MlxDevice::new");
    let mut registry = KernelRegistry::new();

    // Inputs.
    let q_bytes = read_bytes("gated_delta_net_chunk_o_input_q.bin");
    let k_bytes = read_bytes("gated_delta_net_chunk_o_input_k.bin");
    let v_bytes = read_bytes("gated_delta_net_chunk_o_input_v.bin");
    let h_bytes = read_bytes("gated_delta_net_chunk_o_input_h.bin");
    let g_vals = read_f32("gated_delta_net_chunk_o_input_g.bin");

    let exp_q = (B * T * HG * K) as usize;
    let exp_k = (B * T * HG * K) as usize;
    let exp_v = (B * T * H * V) as usize;
    let exp_h = (B * NT * H * V * K) as usize;
    let exp_g = (B * T * H) as usize;
    assert_eq!(q_bytes.len(), exp_q * 2, "q size mismatch");
    assert_eq!(k_bytes.len(), exp_k * 2, "k size mismatch");
    assert_eq!(v_bytes.len(), exp_v * 2, "v size mismatch");
    assert_eq!(h_bytes.len(), exp_h * 2, "h size mismatch");
    assert_eq!(g_vals.len(), exp_g, "g size mismatch");

    let q_buf = upload_bf16(&device, &q_bytes);
    let k_buf = upload_bf16(&device, &k_bytes);
    let v_buf = upload_bf16(&device, &v_bytes);
    let h_buf = upload_bf16(&device, &h_bytes);
    let g_buf = upload_f32(&device, &g_vals);

    // Output.
    let o_elems = (B * T * H * V) as usize;
    let o_buf = alloc_zeroed_bf16(&device, o_elems);

    let scale: f32 = (K as f32).powf(-0.5);
    let p = GatedDeltaNetChunkOParams {
        b: B,
        t: T,
        hg: HG,
        h: H,
        k: K,
        v: V,
        bt: BT,
        scale,
    };
    let params_buf = build_gated_delta_net_chunk_o_params(&device, p).expect("params");

    let mut enc = device.command_encoder().expect("enc");
    dispatch_gated_delta_net_chunk_o(
        &mut enc,
        &mut registry,
        device.metal_device(),
        &q_buf,
        &k_buf,
        &v_buf,
        &h_buf,
        &g_buf,
        &o_buf,
        &params_buf,
        p,
    )
    .expect("dispatch");
    enc.commit_and_wait().expect("commit");

    let o_bytes = o_buf.as_slice::<u8>().expect("read o").to_vec();
    let got_o = bf16_bytes_to_f32(&o_bytes);

    let ref_o_bytes = read_bytes("gated_delta_net_chunk_o_o_ref.bin");
    let ref_o = bf16_bytes_to_f32(&ref_o_bytes);

    assert_eq!(got_o.len(), ref_o.len(), "o length mismatch");

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

    for (i, (&g, &r)) in got_o.iter().zip(ref_o.iter()).enumerate() {
        let abs_err = (g - r).abs();
        let rel_err = abs_err / r.abs().max(f32::EPSILON);
        max_err = max_err.max(abs_err);
        if abs_err > atol && rel_err > rtol {
            panic!(
                "o[{}]: gpu={} ref={} abs_err={} rel_err={}",
                i, g, r, abs_err, rel_err
            );
        }
    }

    eprintln!("chunk_fwd_o OK   max_o_err={:.3e}", max_err);
}

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

    let q_elems = (B * T * HG * K) as usize;
    let k_elems = (B * T * HG * K) as usize;
    let v_elems = (B * T * H * V) as usize;
    let h_elems = (B * NT * H * V * K) as usize;
    let g_elems = (B * T * H) as usize;
    let o_elems = (B * T * H * V) as usize;

    let q_buf = alloc_zeroed_bf16(&device, q_elems);
    let k_buf = alloc_zeroed_bf16(&device, k_elems);
    let v_buf = alloc_zeroed_bf16(&device, v_elems);
    let h_buf = alloc_zeroed_bf16(&device, h_elems);
    // g all-zero -> exp(g)=1 (no decay), but with q=k=v=h=0 the output is
    // still all zero.
    let g_buf = upload_f32(&device, &vec![0.0f32; g_elems]);
    let o_buf = alloc_zeroed_bf16(&device, o_elems);

    let scale: f32 = (K as f32).powf(-0.5);
    let p = GatedDeltaNetChunkOParams {
        b: B,
        t: T,
        hg: HG,
        h: H,
        k: K,
        v: V,
        bt: BT,
        scale,
    };
    let params_buf = build_gated_delta_net_chunk_o_params(&device, p).expect("params");

    let mut enc = device.command_encoder().expect("enc");
    dispatch_gated_delta_net_chunk_o(
        &mut enc,
        &mut registry,
        device.metal_device(),
        &q_buf,
        &k_buf,
        &v_buf,
        &h_buf,
        &g_buf,
        &o_buf,
        &params_buf,
        p,
    )
    .expect("dispatch");
    enc.commit_and_wait().expect("commit");

    let o_bytes = o_buf.as_slice::<u8>().expect("read o");
    assert!(
        o_bytes.iter().all(|&b| b == 0),
        "o should be all-zero with zero q/k/v/h"
    );
}
