//! Criterion bench for attention forward — unchunked (wave 21)
//! vs flash-chunked (wave 29).

use std::hint::black_box;

use criterion::{
    criterion_group, criterion_main, BenchmarkId, Criterion,
};
use jouleclaw_loader_gguf::attention::{
    attention_forward_f32, attention_forward_f32_flash,
};

fn synth(seq: usize, d: usize, off: usize) -> Vec<f32> {
    (0..(seq * d))
        .map(|i| (((i + off) % 17) as f32) * 0.05 - 0.4)
        .collect()
}

fn bench_attention(c: &mut Criterion) {
    let d_head = 128;
    let cases: &[(usize, usize)] = &[
        (1, 256),
        (1, 1024),
        (1, 4096),
        (1, 16384),
        (1, 65536),
        (32, 4096),
    ];
    let mut g_un = c.benchmark_group("attention_unchunked");
    for &(seq_q, seq_kv) in cases {
        let q = synth(seq_q, d_head, 0);
        let k = synth(seq_kv, d_head, 31);
        let v = synth(seq_kv, d_head, 7);
        let mut out = vec![0f32; seq_q * d_head];
        let label = format!("q={seq_q},kv={seq_kv},d={d_head}");
        g_un.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(seq_q, seq_kv),
            |b, _| {
                b.iter(|| {
                    attention_forward_f32(
                        black_box(&q), black_box(&k), black_box(&v),
                        seq_q, seq_kv, d_head, /*causal*/ false,
                        black_box(&mut out),
                    );
                });
            },
        );
    }
    g_un.finish();

    let mut g_fl = c.benchmark_group("attention_flash");
    for &(seq_q, seq_kv) in cases {
        let q = synth(seq_q, d_head, 0);
        let k = synth(seq_kv, d_head, 31);
        let v = synth(seq_kv, d_head, 7);
        let mut out = vec![0f32; seq_q * d_head];
        let label = format!("q={seq_q},kv={seq_kv},d={d_head}");
        g_fl.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(seq_q, seq_kv),
            |b, _| {
                b.iter(|| {
                    attention_forward_f32_flash(
                        black_box(&q), black_box(&k), black_box(&v),
                        seq_q, seq_kv, d_head, /*causal*/ false,
                        black_box(&mut out),
                    );
                });
            },
        );
    }
    g_fl.finish();
}

criterion_group!(benches, bench_attention);
criterion_main!(benches);