//! Criterion bench for Q5_K + Q6_K matmul kernels (wave 28).

// 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_q5_k::{
    matmul_q5_k_amx_tiled, matmul_q5_k_scalar, Q5_K_BLOCK_BYTES, Q5_K_SUPER,
};
use jouleclaw_loader_gguf::matmul_q6_k::{
    matmul_q6_k_amx_tiled, matmul_q6_k_scalar, Q6_K_BLOCK_BYTES, Q6_K_SUPER,
};
#[cfg(target_arch = "aarch64")]
use jouleclaw_loader_gguf::matmul_q5_k::matmul_q5_k_neon;
#[cfg(target_arch = "aarch64")]
use jouleclaw_loader_gguf::matmul_q6_k::matmul_q6_k_neon;

fn synth_w(n: usize, k: usize, block_bytes: usize, super_size: usize) -> Vec<u8> {
    let blocks_per_row = k / super_size;
    let total = n * blocks_per_row * block_bytes;
    let mut w = vec![0u8; total];
    for i in 0..total {
        w[i] = ((i.wrapping_mul(48271)) >> 16) as u8;
    }
    // Set the f16 scale + min near 1.0 — last 2 (Q6_K) or
    // first 4 (Q5_K) bytes per block. Keep simple: just set
    // first f16 to 1.0 across the board, ignoring sub-structure.
    w
}

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

const CASES: &[(usize, usize, usize)] = &[
    (1, 256, 256),
    (1, 1024, 1024),
    (1, 2048, 2048),
    (1, 4096, 4096),
    (8, 1024, 1024),
    (32, 1024, 1024),
];

fn bench_q5_k(c: &mut Criterion) {
    let mut g = c.benchmark_group("matmul_q5_k_scalar");
    for &(m, n, k) in CASES {
        let w = synth_w(n, k, Q5_K_BLOCK_BYTES, Q5_K_SUPER);
        let a = synth_a(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_q5_k_scalar(
                black_box(&a), black_box(&w), black_box(&bias),
                m, n, k, black_box(&mut out),
            )),
        );
    }
    g.finish();
    let mut probe_out = vec![0f32; 1];
    if matmul_q5_k_amx_tiled(
        &vec![0f32; 256], &[0u8; Q5_K_BLOCK_BYTES],
        &[], 1, 1, 256, &mut probe_out,
    ).is_err() { return; }
    let mut g = c.benchmark_group("matmul_q5_k_amx_tiled");
    for &(m, n, k) in CASES {
        let w = synth_w(n, k, Q5_K_BLOCK_BYTES, Q5_K_SUPER);
        let a = synth_a(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_q5_k_amx_tiled(
                    black_box(&a), black_box(&w), black_box(&bias),
                    m, n, k, black_box(&mut out),
                );
            }),
        );
    }
    g.finish();
}

fn bench_q6_k(c: &mut Criterion) {
    let mut g = c.benchmark_group("matmul_q6_k_scalar");
    for &(m, n, k) in CASES {
        let w = synth_w(n, k, Q6_K_BLOCK_BYTES, Q6_K_SUPER);
        let a = synth_a(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_q6_k_scalar(
                black_box(&a), black_box(&w), black_box(&bias),
                m, n, k, black_box(&mut out),
            )),
        );
    }
    g.finish();
    let mut probe_out = vec![0f32; 1];
    if matmul_q6_k_amx_tiled(
        &vec![0f32; 256], &vec![0u8; Q6_K_BLOCK_BYTES],
        &[], 1, 1, 256, &mut probe_out,
    ).is_err() { return; }
    let mut g = c.benchmark_group("matmul_q6_k_amx_tiled");
    for &(m, n, k) in CASES {
        let w = synth_w(n, k, Q6_K_BLOCK_BYTES, Q6_K_SUPER);
        let a = synth_a(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_q6_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_q5_k_neon(c: &mut Criterion) {
    let mut g = c.benchmark_group("matmul_q5_k_neon");
    for &(m, n, k) in CASES {
        let w = synth_w(n, k, Q5_K_BLOCK_BYTES, Q5_K_SUPER);
        let a = synth_a(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_q5_k_neon(
                black_box(&a), black_box(&w), black_box(&bias),
                m, n, k, black_box(&mut out),
            )),
        );
    }
    g.finish();
}

#[cfg(target_arch = "aarch64")]
fn bench_q6_k_neon(c: &mut Criterion) {
    let mut g = c.benchmark_group("matmul_q6_k_neon");
    for &(m, n, k) in CASES {
        let w = synth_w(n, k, Q6_K_BLOCK_BYTES, Q6_K_SUPER);
        let a = synth_a(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_q6_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_q5_k_neon(_c: &mut Criterion) {}
#[cfg(not(target_arch = "aarch64"))]
fn bench_q6_k_neon(_c: &mut Criterion) {}

criterion_group!(benches, bench_q5_k, bench_q5_k_neon, bench_q6_k, bench_q6_k_neon);
criterion_main!(benches);