// M2a KV-cache decode engine — the token-identical correctness gate.
//
// The centerpiece M2 invariant: greedy decode through the KV-cache incremental
// engine (`xtrain_model::generate_greedy_cached`) must be TOKEN-IDENTICAL to the
// naive full-recompute greedy (`xtrain_train::sample::generate` at temperature 0),
// which re-runs the whole forward over the growing prefix each step. Same tokens ⇒
// the cache + decode-time attention + RoPE-at-position reproduce the full forward.
//
// Numerics note: a randomly-initialised model has near-uniform logits, so argmax
// can be fragile to ~1e-6 differences. This unit gate therefore runs in F32 (the
// tightest path, and the dtype the eval harness actually uses) on a small model.
// The headline gate on the trained v12 checkpoint (peaked logits → robust argmax)
// is run on the GPU box and recorded in docs/18.
#![cfg(not(no_cuda))]

use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, generate_greedy_cached};
use xtrain_tensor::{DType, Device};

fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
    let mut state = seed
        .wrapping_mul(2862933555777941757)
        .wrapping_add(3037000493);
    (0..n)
        .map(|_| {
            state = state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
        })
        .collect()
}

fn build(cfg: Config, device: Device, dtype: DType) -> TinyTransformer {
    let mut seed = 1u64;
    let m = TinyTransformer::new(cfg, device, |shape| {
        seed = seed.wrapping_add(1);
        let n: usize = shape.iter().product();
        if shape.len() == 1 {
            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
        } else {
            fill(n, seed, 0.08)
        }
    });
    m.with_compute_dtype(dtype)
}

// A real GQA config (8 query / 2 kv heads → group 4) to exercise repeat_kv in the
// decode path; head_dim 16, dim 128, 4 layers.
fn gqa_cfg() -> Config {
    Config::from_arch(48, 8, 16, 4, 256).with_kv_heads(2)
}

#[test]
fn kv_cache_decode_is_token_identical_to_naive_f32() {
    assert!(
        device::device_count().expect("device count") > 0,
        "no CUDA device"
    );
    device::set_device(0).unwrap();
    let device = Device::Cuda(0);

    let model = build(gqa_cfg(), device, DType::F32);
    let prompt: Vec<i32> = vec![1, 5, 9, 13, 2, 7];
    let max_new = 24usize;

    let mut rng = 7u64;
    let naive = xtrain_train::sample::generate(&model, device, &prompt, max_new, 0.0, &mut rng);
    let cached = generate_greedy_cached(&model, device, &prompt, max_new);

    assert_eq!(
        naive.len(),
        cached.len(),
        "length mismatch: naive {} vs cached {}",
        naive.len(),
        cached.len()
    );
    if naive != cached {
        // Report the first divergence for debugging.
        let first = naive
            .iter()
            .zip(&cached)
            .position(|(a, b)| a != b)
            .unwrap();
        panic!(
            "token divergence at index {first}: naive={:?} cached={:?}\nnaive  ={naive:?}\ncached ={cached:?}",
            naive[first], cached[first]
        );
    }
    println!(
        "KV-cache decode token-identical to naive over {} generated tokens (F32, GQA 8/2)",
        max_new
    );
}