//! Wave 5b.1 iter 4 — end-to-end chunk_gated_delta_rule_fwd orchestrator test.
//!
//! Spec source: FLA chunk.py:23-84 + chunk.py:106-108 (l2norm pre-step).
//!
//! The reference fixture is generated by:
//! ```sh
//! python3 tests/fixtures/chunk_gated_delta_rule_fwd_reference.py
//! ```
//!
//! The independent oracle (naive O(T^2) per-token recurrence) cross-validates
//! the reference at:
//! ```sh
//! python3 tests/fixtures/chunk_gated_delta_rule_fwd_oracle.py
//! ```
//!
//! Tolerance bar (looser than per-kernel fixtures because 6 stages of
//! bf16 round-off compound):
//!   o:           atol = 1e-2, rtol = 5e-2
//!   final_state: atol = 1e-2

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

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

use mlx_native::ops::chunk_gated_delta_rule::{
    dispatch_chunk_gated_delta_rule_fwd, ChunkGatedDeltaRuleParams,
};
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/chunk_gated_delta_rule_fwd_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
}

fn alloc_zeroed_f32(device: &MlxDevice, n_elems: usize) -> MlxBuffer {
    let mut buf = device
        .alloc_buffer(n_elems * 4, DType::F32, vec![n_elems])
        .expect("alloc f32");
    buf.as_mut_slice::<f32>().expect("mut").fill(0.0f32);
    buf
}

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

    // Inputs.
    let q_bytes = read_bytes("chunk_gated_delta_rule_fwd_input_q.bin");
    let k_bytes = read_bytes("chunk_gated_delta_rule_fwd_input_k.bin");
    let v_bytes = read_bytes("chunk_gated_delta_rule_fwd_input_v.bin");
    let g_log_decay = read_f32("chunk_gated_delta_rule_fwd_input_g_log_decay.bin");
    let beta = read_f32("chunk_gated_delta_rule_fwd_input_beta.bin");
    let h0 = read_f32("chunk_gated_delta_rule_fwd_input_h0.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_g = (B * T * H) as usize;
    let exp_beta = (B * T * H) as usize;
    let exp_h0 = (B * H * V * K) 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!(g_log_decay.len(), exp_g, "g_log_decay size mismatch");
    assert_eq!(beta.len(), exp_beta, "beta size mismatch");
    assert_eq!(h0.len(), exp_h0, "h0 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 g_log_decay_buf = upload_f32(&device, &g_log_decay);
    let beta_buf = upload_f32(&device, &beta);
    let h0_buf = upload_f32(&device, &h0);

    // Outputs.
    let o_elems = (B * T * H * V) as usize;
    let final_state_elems = (B * H * V * K) as usize;
    let o_buf = alloc_zeroed_bf16(&device, o_elems);
    let final_state_buf = alloc_zeroed_f32(&device, final_state_elems);

    let scale: f32 = (K as f32).powf(-0.5);
    let p = ChunkGatedDeltaRuleParams {
        b: B,
        t: T,
        hg: HG,
        h: H,
        k: K,
        v: V,
        bt: BT,
        scale,
        use_qk_l2norm: true,
    };

    let mut enc = device.command_encoder().expect("enc");
    dispatch_chunk_gated_delta_rule_fwd(
        &mut enc,
        &mut registry,
        &device,
        &q_buf,
        &k_buf,
        &v_buf,
        &g_log_decay_buf,
        &beta_buf,
        &h0_buf,
        &o_buf,
        &final_state_buf,
        p,
    )
    .expect("dispatch");
    enc.commit_and_wait().expect("commit");

    // Read outputs.
    let o_bytes = o_buf.as_slice::<u8>().expect("read o").to_vec();
    let got_o = bf16_bytes_to_f32(&o_bytes);
    let got_final_state: Vec<f32> = final_state_buf
        .as_slice::<f32>()
        .expect("read final_state")
        .to_vec();

    let ref_o_bytes = read_bytes("chunk_gated_delta_rule_fwd_o_ref.bin");
    let ref_o = bf16_bytes_to_f32(&ref_o_bytes);
    let ref_final_state = read_f32("chunk_gated_delta_rule_fwd_final_state_ref.bin");

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

    // End-to-end tolerance bar — 6 bf16 stages compound (looser than the
    // 5e-3 atol / 2e-2 rtol bar each per-kernel test uses individually).
    let o_atol: f32 = 1e-2;
    let o_rtol: f32 = 5e-2;
    let fs_atol: f32 = 1e-2;
    let mut max_o_err = 0.0f32;
    let mut max_fs_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_o_err = max_o_err.max(abs_err);
        if abs_err > o_atol && rel_err > o_rtol {
            panic!(
                "o[{}]: gpu={} ref={} abs_err={} rel_err={}",
                i, g, r, abs_err, rel_err
            );
        }
    }

    for (i, (&g, &r)) in got_final_state.iter().zip(ref_final_state.iter()).enumerate() {
        let abs_err = (g - r).abs();
        max_fs_err = max_fs_err.max(abs_err);
        if abs_err > fs_atol {
            panic!(
                "final_state[{}]: gpu={} ref={} abs_err={}",
                i, g, r, abs_err
            );
        }
    }

    eprintln!(
        "chunk_gated_delta_rule_fwd OK   max_o_err={:.3e}    max_final_state_err={:.3e}",
        max_o_err, max_fs_err
    );
}
