// M2b batched KV-cache decode — the token-identical gate.
//
// Batched decode rolls out G samples of one prompt in lockstep (one common decode
// position each step, uniform RoPE via rope_pos, KV cache carrying a G dimension).
// Under GREEDY decoding all G rows are deterministic and must each equal the
// single-sequence greedy decode (generate_greedy_cached, itself gated token-
// identical to the naive sampler). This pins that the G-way batching indexes each
// sequence's K/V correctly (no cross-row contamination) and reproduces M2a exactly.
#![cfg(not(no_cuda))]

use xtrain_cuda::device;
use xtrain_model::{generate_cached_batch, generate_greedy_cached, Config, TinyTransformer};
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) -> TinyTransformer {
    let mut seed = 1u64;
    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)
        }
    })
    .with_compute_dtype(DType::F32)
}

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

    // Real GQA (8 query / 2 kv heads → group 4) so repeat_kv(nh, batch=G) is exercised.
    let cfg = Config::from_arch(48, 8, 16, 4, 256).with_kv_heads(2);
    let model = build(cfg, device);
    let prompt: Vec<i32> = vec![3, 9, 1, 14, 5];
    let max_new = 24usize;
    let g = 5usize;

    let single = generate_greedy_cached(&model, device, &prompt, max_new);
    let mut rng = 0u64;
    let batched = generate_cached_batch(&model, device, &prompt, g, max_new, 0.0, &mut rng);

    assert_eq!(batched.len(), g, "expected {g} sample rows");
    for (row, seq) in batched.iter().enumerate() {
        assert_eq!(
            seq.len(),
            single.len(),
            "row {row} length {} vs single {}",
            seq.len(),
            single.len()
        );
        if seq != &single {
            let first = seq.iter().zip(&single).position(|(a, b)| a != b).unwrap();
            panic!(
                "batched row {row} diverges from single-seq at index {first}: {:?} vs {:?}",
                seq[first], single[first]
            );
        }
    }
    println!(
        "batched decode OK: all {g} greedy rows token-identical to single-seq over {max_new} tokens"
    );
}