//! Criterion bench for Q4_K matmul — scalar fused (wave 19)
//! vs AMX-tiled (wave 26).
//!
//! Q4_K is the dominant format in distributed Llama-family
//! models (Q4_K_M). These numbers + the wave-24 Q8_0 numbers
//! tell the full quantized-matmul story on Apple Silicon.

// Benches and integration tests are their own crates, so the decision
// recorded in src/lib.rs does not reach them: these loops fill arrays by
// index, which is what the index is for.
#![allow(clippy::needless_range_loop)]

use std::hint::black_box;

use criterion::{
    criterion_group, criterion_main, BenchmarkId, Criterion, Throughput,
};
use jouleclaw_loader_gguf::matmul_q4_k::{
    matmul_q4_k_amx_tiled, matmul_q4_k_scalar, Q4_K_BLOCK_BYTES, Q4_K_SUPER,
};
#[cfg(target_arch = "aarch64")]
use jouleclaw_loader_gguf::matmul_q4_k::matmul_q4_k_neon;

/// Build a synthetic Q4_K weight buffer of shape [n, k].
/// Bit-pattern derived from index; deterministic. Sets
/// d / dmin to fp16-exact values near 1.0 to keep the
/// numerics clean for parity asserts elsewhere.
fn synth_q4_k_weights(n: usize, k: usize) -> Vec<u8> {
    assert_eq!(k % Q4_K_SUPER, 0);
    let blocks_per_row = k / Q4_K_SUPER;
    let row_bytes = blocks_per_row * Q4_K_BLOCK_BYTES;
    let mut w = vec![0u8; n * row_bytes];
    for i in 0..w.len() {
        w[i] = ((i.wrapping_mul(48271)) >> 16) as u8;
    }
    for row in 0..n {
        for blk in 0..blocks_per_row {
            let off = row * row_bytes + blk * Q4_K_BLOCK_BYTES;
            w[off..off + 2].copy_from_slice(&0x3C00u16.to_le_bytes());   // d = 1.0
            w[off + 2..off + 4].copy_from_slice(&0x3800u16.to_le_bytes()); // dmin = 0.5
        }
    }
    w
}

fn synth_activations(m: usize, k: usize) -> Vec<f32> {
    (0..(m * k)).map(|i| (i % 13) as f32 * 0.01).collect()
}

fn bench_matmul_q4_k_scalar(c: &mut Criterion) {
    let mut g = c.benchmark_group("matmul_q4_k_scalar");
    let cases = [
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    for &(m, n, k) in &cases {
        let w = synth_q4_k_weights(n, k);
        let a = synth_activations(m, k);
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g.throughput(Throughput::Elements(flops as u64));
        g.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    matmul_q4_k_scalar(
                        black_box(&a), black_box(&w), black_box(&bias),
                        m, n, k, black_box(&mut out),
                    );
                });
            },
        );
    }
    g.finish();
}

fn bench_matmul_q4_k_amx_tiled(c: &mut Criterion) {
    let mut probe_out = vec![0f32; 1];
    let probe_a = vec![0f32; 256];
    let probe_w = vec![0u8; Q4_K_BLOCK_BYTES];
    if matmul_q4_k_amx_tiled(&probe_a, &probe_w, &[], 1, 1, 256, &mut probe_out).is_err() {
        return;
    }
    let mut g = c.benchmark_group("matmul_q4_k_amx_tiled");
    let cases = [
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    for &(m, n, k) in &cases {
        let w = synth_q4_k_weights(n, k);
        let a = synth_activations(m, k);
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g.throughput(Throughput::Elements(flops as u64));
        g.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    let _ = matmul_q4_k_amx_tiled(
                        black_box(&a), black_box(&w), black_box(&bias),
                        m, n, k, black_box(&mut out),
                    );
                });
            },
        );
    }
    g.finish();
}

#[cfg(target_arch = "aarch64")]
fn bench_matmul_q4_k_neon(c: &mut Criterion) {
    let mut g = c.benchmark_group("matmul_q4_k_neon");
    let cases = [
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    for &(m, n, k) in &cases {
        let w = synth_q4_k_weights(n, k);
        let a = synth_activations(m, k);
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g.throughput(Throughput::Elements(flops as u64));
        g.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    matmul_q4_k_neon(
                        black_box(&a), black_box(&w), black_box(&bias),
                        m, n, k, black_box(&mut out),
                    );
                });
            },
        );
    }
    g.finish();
}

#[cfg(not(target_arch = "aarch64"))]
fn bench_matmul_q4_k_neon(_c: &mut Criterion) {}

criterion_group!(
    benches,
    bench_matmul_q4_k_scalar,
    bench_matmul_q4_k_neon,
    bench_matmul_q4_k_amx_tiled,
);
criterion_main!(benches);