//! Wave 5b.1 iter 1 — chunk-parallel inter-chunk state-recurrence kernel.
//!
//! Spec source:
//! - arXiv 2412.06464 §4 (Yang–Hatamizadeh; chunkwise parallelization)
//! - FLA reference: `chunk_gated_delta_rule_fwd_kernel_h_blockdim64` at
//!   /opt/vllm/vllm/model_executor/layers/fla/ops/chunk_delta_h.py:43-298
//!
//! The reference fixture is generated by the deterministic PyTorch port at
//! `tests/fixtures/gated_delta_net_chunk_reference.py`. To regenerate:
//!
//! ```sh
//! python3 tests/fixtures/gated_delta_net_chunk_reference.py
//! ```
//!
//! Tolerance bar: `bf16 atol = 5e-3, rtol = 2e-2` for the bf16 outputs (`h`
//! and `v_new`); same bar for the f32 final-state output. This matches the
//! `flash_attn_prefill_d512` bf16 tolerance bar (see
//! `flash_attn_prefill_d512.metal:21`).

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

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

use mlx_native::ops::gated_delta_net_chunk::{
    build_gated_delta_net_chunk_params, dispatch_gated_delta_net_chunk_inter_state,
    GatedDeltaNetChunkParams,
};
use mlx_native::{DType, KernelRegistry, MlxBuffer, MlxDevice};

// ---------------------------------------------------------------------------
// Fixture metadata — must match `gated_delta_net_chunk_reference.py`.
// ---------------------------------------------------------------------------
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_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) {
        // bf16 = upper 16 bits of f32, padded with 16 low zeros.
        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
}

// ---------------------------------------------------------------------------
// The TDD test — exercises the actual GPU kernel against the FLA reference.
// ---------------------------------------------------------------------------
#[test]
fn test_chunk_inter_state_matches_fla_reference_bf16() {
    let device = MlxDevice::new().expect("MlxDevice::new");
    let mut registry = KernelRegistry::new();

    // Inputs — loaded from the Python-generated fixture.
    let k_bytes = read_bytes("gated_delta_net_chunk_input_k.bin");
    let w_bytes = read_bytes("gated_delta_net_chunk_input_w.bin");
    let u_bytes = read_bytes("gated_delta_net_chunk_input_u.bin");
    let g_vals = read_f32("gated_delta_net_chunk_input_g.bin");
    let h0_vals = read_f32("gated_delta_net_chunk_input_h0.bin");

    let expected_k_elems = (B * T * HG * K) as usize;
    let expected_w_elems = (B * T * H * K) as usize;
    let expected_u_elems = (B * T * H * V) as usize;
    let expected_g_elems = (B * T * H) as usize;
    let expected_h0_elems = (B * H * V * K) as usize;
    assert_eq!(k_bytes.len(), expected_k_elems * 2, "k size mismatch");
    assert_eq!(w_bytes.len(), expected_w_elems * 2, "w size mismatch");
    assert_eq!(u_bytes.len(), expected_u_elems * 2, "u size mismatch");
    assert_eq!(g_vals.len(), expected_g_elems, "g size mismatch");
    assert_eq!(h0_vals.len(), expected_h0_elems, "h0 size mismatch");

    let k_buf = upload_bf16(&device, &k_bytes);
    let w_buf = upload_bf16(&device, &w_bytes);
    let u_buf = upload_bf16(&device, &u_bytes);
    let g_buf = upload_f32(&device, &g_vals);
    let h0_buf = upload_f32(&device, &h0_vals);

    // Output buffers.
    let h_out_elems = (B * NT * H * V * K) as usize;
    let v_new_elems = (B * T * H * V) as usize;
    let final_elems = (B * H * V * K) as usize;

    let h_out_buf = alloc_zeroed_bf16(&device, h_out_elems);
    let v_new_buf = alloc_zeroed_bf16(&device, v_new_elems);
    let final_buf = upload_f32(&device, &vec![0.0f32; final_elems]);

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

    let mut enc = device.command_encoder().expect("enc");
    dispatch_gated_delta_net_chunk_inter_state(
        &mut enc,
        &mut registry,
        device.metal_device(),
        &k_buf,
        &w_buf,
        &u_buf,
        &g_buf,
        &h0_buf,
        &h_out_buf,
        &v_new_buf,
        &final_buf,
        &params_buf,
        p,
    )
    .expect("dispatch");
    enc.commit_and_wait().expect("commit");

    // Read back outputs.
    let h_out_bytes = h_out_buf.as_slice::<u8>().expect("read h").to_vec();
    let v_new_bytes = v_new_buf.as_slice::<u8>().expect("read v_new").to_vec();
    let final_vec = final_buf.as_slice::<f32>().expect("read final").to_vec();

    let got_h = bf16_bytes_to_f32(&h_out_bytes);
    let got_v_new = bf16_bytes_to_f32(&v_new_bytes);

    // Reference.
    let ref_h_bytes = read_bytes("gated_delta_net_chunk_state_ref.bin");
    let ref_v_new_bytes = read_bytes("gated_delta_net_chunk_v_new_ref.bin");
    let ref_final = read_f32("gated_delta_net_chunk_final_ref.bin");
    let ref_h = bf16_bytes_to_f32(&ref_h_bytes);
    let ref_v_new = bf16_bytes_to_f32(&ref_v_new_bytes);

    assert_eq!(got_h.len(), ref_h.len(), "h length mismatch");
    assert_eq!(got_v_new.len(), ref_v_new.len(), "v_new length mismatch");
    assert_eq!(final_vec.len(), ref_final.len(), "final length mismatch");

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

    let mut max_h_err = 0.0f32;
    let mut max_v_err = 0.0f32;
    let mut max_f_err = 0.0f32;

    for (i, (&g, &r)) in got_h.iter().zip(ref_h.iter()).enumerate() {
        let abs_err = (g - r).abs();
        let rel_err = abs_err / r.abs().max(f32::EPSILON);
        max_h_err = max_h_err.max(abs_err);
        if abs_err > atol && rel_err > rtol {
            panic!(
                "h[{}]: gpu={} ref={} abs_err={} rel_err={}",
                i, g, r, abs_err, rel_err
            );
        }
    }
    for (i, (&g, &r)) in got_v_new.iter().zip(ref_v_new.iter()).enumerate() {
        let abs_err = (g - r).abs();
        let rel_err = abs_err / r.abs().max(f32::EPSILON);
        max_v_err = max_v_err.max(abs_err);
        if abs_err > atol && rel_err > rtol {
            panic!(
                "v_new[{}]: gpu={} ref={} abs_err={} rel_err={}",
                i, g, r, abs_err, rel_err
            );
        }
    }
    for (i, (&g, &r)) in final_vec.iter().zip(ref_final.iter()).enumerate() {
        let abs_err = (g - r).abs();
        let rel_err = abs_err / r.abs().max(f32::EPSILON);
        max_f_err = max_f_err.max(abs_err);
        if abs_err > atol && rel_err > rtol {
            panic!(
                "final[{}]: gpu={} ref={} abs_err={} rel_err={}",
                i, g, r, abs_err, rel_err
            );
        }
    }

    eprintln!(
        "chunk_inter_state OK   max_h_err={:.3e}   max_v_err={:.3e}   max_final_err={:.3e}",
        max_h_err, max_v_err, max_f_err
    );
}

// ---------------------------------------------------------------------------
// Sanity: zero inputs => zero outputs (the kernel is connected and wired).
// ---------------------------------------------------------------------------
#[test]
fn test_chunk_inter_state_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 w_elems = (B * T * H * K) as usize;
    let u_elems = (B * T * H * V) as usize;
    let g_elems = (B * T * H) as usize;
    let h0_elems = (B * H * V * K) as usize;
    let h_out_elems = (B * NT * H * V * K) as usize;
    let v_new_elems = (B * T * H * V) as usize;
    let final_elems = (B * H * V * K) as usize;

    let k_buf = alloc_zeroed_bf16(&device, k_elems);
    let w_buf = alloc_zeroed_bf16(&device, w_elems);
    let u_buf = alloc_zeroed_bf16(&device, u_elems);
    let g_buf = upload_f32(&device, &vec![0.0f32; g_elems]);
    let h0_buf = upload_f32(&device, &vec![0.0f32; h0_elems]);

    let h_out_buf = alloc_zeroed_bf16(&device, h_out_elems);
    let v_new_buf = alloc_zeroed_bf16(&device, v_new_elems);
    let final_buf = upload_f32(&device, &vec![0.0f32; final_elems]);

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

    let mut enc = device.command_encoder().expect("enc");
    dispatch_gated_delta_net_chunk_inter_state(
        &mut enc,
        &mut registry,
        device.metal_device(),
        &k_buf,
        &w_buf,
        &u_buf,
        &g_buf,
        &h0_buf,
        &h_out_buf,
        &v_new_buf,
        &final_buf,
        &params_buf,
        p,
    )
    .expect("dispatch");
    enc.commit_and_wait().expect("commit");

    let h_out_bytes = h_out_buf.as_slice::<u8>().expect("read h");
    let v_new_bytes = v_new_buf.as_slice::<u8>().expect("read v_new");
    let final_vec = final_buf.as_slice::<f32>().expect("read final");

    assert!(
        h_out_bytes.iter().all(|&b| b == 0),
        "h_out should be all-zero with zero inputs"
    );
    assert!(
        v_new_bytes.iter().all(|&b| b == 0),
        "v_new should be all-zero with zero inputs"
    );
    assert!(
        final_vec.iter().all(|&v| v == 0.0),
        "final_state should be all-zero with zero inputs"
    );
}

/// Wave 5b.2 iter-1.5 regression: K=192 must be rejected by validate().
///
/// Codex iter-1 audit caught that section 2f's simdgroup_matrix MMA loop
/// hard-codes 16 K-tiles (K=128/8). Generalizing to runtime `K/8u` defeats
/// Metal's MMA scheduler (measured 3.15× regression, 1.08 → 3.40 ms), so
/// MAX_K is narrowed to 128 to match the kernel's compile-time tile count.
/// This test asserts the contract narrowing — K=192 dispatch must error
/// rather than silently corrupting cols 128..191.
#[test]
fn test_chunk_inter_state_k192_rejected_by_validate() {
    use mlx_native::ops::gated_delta_net_chunk::MAX_K;
    assert_eq!(MAX_K, 128, "MAX_K narrowed to 128 in Wave 5b.2 iter 1.5");

    // Try to dispatch at K=192 and observe the validate-side rejection.
    let device = MlxDevice::new().expect("MlxDevice::new");
    let mut registry = KernelRegistry::new();

    const K_192: u32 = 192;
    let k_elems = (B * T * HG * K_192) as usize;
    let w_elems = (B * T * H * K_192) as usize;
    let u_elems = (B * T * H * V) as usize;
    let g_elems = (B * T * H) as usize;
    let h0_elems = (B * H * V * K_192) as usize;
    let h_out_elems = (B * NT * H * V * K_192) as usize;
    let v_new_elems = (B * T * H * V) as usize;
    let final_elems = (B * H * V * K_192) as usize;

    let k_buf = alloc_zeroed_bf16(&device, k_elems);
    let w_buf = alloc_zeroed_bf16(&device, w_elems);
    let u_buf = alloc_zeroed_bf16(&device, u_elems);
    let g_buf = upload_f32(&device, &vec![0.0f32; g_elems]);
    let h0_buf = upload_f32(&device, &vec![0.0f32; h0_elems]);
    let h_out_buf = alloc_zeroed_bf16(&device, h_out_elems);
    let v_new_buf = alloc_zeroed_bf16(&device, v_new_elems);
    let final_buf = upload_f32(&device, &vec![0.0f32; final_elems]);

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

    let mut enc = device.command_encoder().expect("enc");
    let result = dispatch_gated_delta_net_chunk_inter_state(
        &mut enc,
        &mut registry,
        device.metal_device(),
        &k_buf,
        &w_buf,
        &u_buf,
        &g_buf,
        &h0_buf,
        &h_out_buf,
        &v_new_buf,
        &final_buf,
        &params_buf,
        p,
    );
    let err = result.expect_err("K=192 dispatch must error");
    let msg = err.to_string();
    assert!(
        msg.contains("192") && (msg.contains("MAX_K") || msg.contains("threadgroup")),
        "expected K=192-rejection error citing MAX_K or threadgroup budget, got: {msg}"
    );
}
