//! Wave 5b.1 iter 2 — chunk_scaled_dot_kkt kernel test.
//!
//! Spec source:
//! - FLA reference: `chunk_scaled_dot_kkt_fwd_kernel` at
//!   /opt/vllm/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py:36-99
//!
//! The reference fixture is generated by:
//! ```sh
//! python3 tests/fixtures/gated_delta_net_kkt_reference.py
//! ```
//!
//! Tolerance bar: `atol = 5e-3, rtol = 2e-2` (matches iter-1 chunk fixture).
//! The output is f32 but the inputs/intermediates carry bf16 quantization
//! noise from FLA line 86 (b_kb cast to bf16 before the dot).

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

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

use mlx_native::ops::gated_delta_net_kkt::{
    build_gated_delta_net_kkt_params, dispatch_gated_delta_net_kkt,
    GatedDeltaNetKktParams,
};
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 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_kkt_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 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
}

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

    // Inputs.
    let k_bytes = read_bytes("gated_delta_net_kkt_input_k.bin");
    let beta_vals = read_f32("gated_delta_net_kkt_input_beta.bin");
    let g_vals = read_f32("gated_delta_net_kkt_input_g.bin");

    let exp_k = (B * T * HG * K) as usize;
    let exp_beta = (B * T * H) as usize;
    let exp_g = (B * T * H) as usize;
    assert_eq!(k_bytes.len(), exp_k * 2, "k size mismatch");
    assert_eq!(beta_vals.len(), exp_beta, "beta size mismatch");
    assert_eq!(g_vals.len(), exp_g, "g size mismatch");

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

    // Output: A [B, T, H, BT] f32.
    let a_elems = (B * T * H * BT) as usize;
    let a_buf = upload_f32(&device, &vec![0.0f32; a_elems]);

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

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

    let got_a = a_buf.as_slice::<f32>().expect("read A").to_vec();
    let ref_a = read_f32("gated_delta_net_kkt_A_ref.bin");

    assert_eq!(got_a.len(), ref_a.len(), "A length mismatch");

    let atol: f32 = 5e-3;
    let rtol: f32 = 2e-2;
    let mut max_err = 0.0f32;
    for (i, (&g, &r)) in got_a.iter().zip(ref_a.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!(
                "A[{}]: gpu={} ref={} abs_err={} rel_err={}",
                i, g, r, abs_err, rel_err
            );
        }
    }
    eprintln!("kkt OK   max_err={:.3e}", max_err);
}

#[test]
fn test_kkt_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 beta_elems = (B * T * H) as usize;
    let g_elems = (B * T * H) as usize;
    let a_elems = (B * T * H * BT) as usize;

    // bf16 zeros: bytes 0 = bf16 0.0
    let mut k_buf = device
        .alloc_buffer(k_elems * 2, DType::BF16, vec![k_elems])
        .expect("alloc k");
    k_buf.as_mut_slice::<u8>().expect("k bytes").fill(0);

    let beta_buf = upload_f32(&device, &vec![0.0f32; beta_elems]);
    let g_buf = upload_f32(&device, &vec![0.0f32; g_elems]);
    let a_buf = upload_f32(&device, &vec![0.0f32; a_elems]);

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

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

    // With k=0 and beta=0, b_A = 0 (the dot is zero); after gate exp(0-0)=1
    // it stays 0; after mask it stays 0. So A must be all-zero.
    let got_a = a_buf.as_slice::<f32>().expect("read A");
    assert!(
        got_a.iter().all(|&v| v == 0.0),
        "A should be all-zero with zero inputs"
    );
}
