//! GPU kernel tests (need a CUDA device): `cargo test --release -p dllm-kernels -- --ignored`
//! FP8 GEMV + fused lm_head argmax and the NVFP4 MoE decode chain on analytic inputs; decode / prefill attention against an f32 reference,
//! and bit-identical results through a permuted sliding-window pool behind a page map (pages outside the window map to a NaN-filled null page).
use dllm_kernels::moe::{Experts, MoeDecodeScratch};
use dllm_kernels::{attn, head, linear};
use dllm_kernels::{Backend, DBuf, Gpu};
use std::sync::Arc;

fn bf(x: f32) -> u16 { (x.to_bits() >> 16) as u16 }
fn f(x: u16) -> f32 { f32::from_bits((x as u32) << 16) }
struct Rng(u64);
impl Rng { fn next(&mut self) -> f32 { self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); (((self.0 >> 33) as u32) as f32 / 2147483648.0) - 1.0 } }

fn rel_rms(a: &[u16], b: &[u16]) -> f32 {
    let (mut n, mut d) = (0f64, 0f64);
    for (x, y) in a.iter().zip(b) { let (x, y) = (f(*x) as f64, f(*y) as f64); n += (x - y) * (x - y); d += y * y; }
    (n / d.max(1e-30)).sqrt() as f32
}

/// permuted copy of a [npg][nkv][16][hd] pool with a NaN null page 0: page p (>= first_live) lives at physical page npg - p, pages below `first_live` map to 0
fn mapped_pool(k: &[u16], v: &[u16], npg: usize, nkv: usize, hd: usize, first_live: usize) -> (Vec<u16>, Vec<u16>, Vec<i32>) {
    let pe = nkv * 16 * hd;
    let (mut mk, mut mv) = (vec![0x7fc0u16; (npg + 1) * pe], vec![0x7fc0u16; (npg + 1) * pe]);
    let mut map = vec![0i32; npg];
    for p in first_live..npg {
        let s = npg - p;
        map[p] = s as i32;
        mk[s * pe..(s + 1) * pe].copy_from_slice(&k[p * pe..(p + 1) * pe]);
        mv[s * pe..(s + 1) * pe].copy_from_slice(&v[p * pe..(p + 1) * pe]);
    }
    (mk, mv, map)
}

/// f32 reference: scaling-1.0 GQA attention of query rows at positions `qpos` over the pool [page][nkv][16][hd] (identity block table),
/// causal, sliding layers additionally `qpos - kvpos < window`
#[allow(clippy::too_many_arguments)]
fn reference(q: &[u16], k: &[u16], v: &[u16], qpos: &[i32], nh: usize, nkv: usize, hd: usize, window: usize) -> Vec<u16> {
    let at = |pool: &[u16], p: usize, h: usize, d: usize| f(pool[((p / 16 * nkv + h) * 16 + p % 16) * hd + d]);
    let mut out = vec![0u16; qpos.len() * nh * hd];
    for (t, &qp) in qpos.iter().enumerate() {
        let qp = qp as usize;
        let lo = if window > 0 { (qp + 1).saturating_sub(window) } else { 0 };
        for h in 0..nh {
            let kh = h / (nh / nkv);
            let s: Vec<f64> = (lo..=qp).map(|p| (0..hd).map(|d| f(q[(t * nh + h) * hd + d]) as f64 * at(k, p, kh, d) as f64).sum()).collect();
            let m = s.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
            let e: Vec<f64> = s.iter().map(|x| (x - m).exp()).collect();
            let z: f64 = e.iter().sum();
            for d in 0..hd {
                let o: f64 = (lo..=qp).zip(&e).map(|(p, w)| w / z * at(v, p, kh, d) as f64).sum();
                out[(t * nh + h) * hd + d] = bf(o as f32);
            }
        }
    }
    out
}

fn run(nh: usize, nkv: usize, hd: usize, window: usize, ctx: usize, rows: usize) {
    let cb = dllm_kernels::gpu::CudaBackend::new(0).expect("cuda");
    let be: Arc<dyn Backend> = cb;
    let g = Gpu::new(be.clone());
    let mut rng = Rng(12345 + hd as u64);
    let npg = (ctx + 15) / 16 + 1;
    let kv: Vec<u16> = (0..npg * nkv * 16 * hd).map(|_| bf(rng.next())).collect();
    let vv: Vec<u16> = (0..npg * nkv * 16 * hd).map(|_| bf(rng.next())).collect();
    let q: Vec<u16> = (0..rows * nh * hd).map(|_| bf(rng.next())).collect();
    let kp = DBuf::from_slice(&be, &kv).unwrap(); let vp = DBuf::from_slice(&be, &vv).unwrap(); let qd = DBuf::from_slice(&be, &q).unwrap();
    let bt: Vec<i32> = (0..(ctx + 15) / 16).map(|i| i as i32).collect();
    let btd = DBuf::from_slice(&be, &bt).unwrap();
    let qpos: Vec<i32> = (0..rows).map(|j| (ctx - rows + j) as i32).collect();
    let qposd = DBuf::from_slice(&be, &qpos).unwrap();
    let off = DBuf::from_slice(&be, &vec![0i32; rows]).unwrap();
    let splits = attn::choose_splits(if window > 0 { ctx.min(window) } else { ctx }, rows, nkv, rows);
    let part = DBuf::zeros(&be, attn::decode_part_bytes(rows, nh, splits, hd)).unwrap();
    let cnt = DBuf::zeros(&be, 4096 * 4).unwrap();
    let o_dev = DBuf::zeros(&be, rows * nh * hd * 2).unwrap();
    attn::attn_decode(&g, o_dev.ptr, part.ptr, qd.ptr, kp.ptr, vp.ptr, btd.ptr, off.ptr, qposd.ptr, qposd.ptr, rows, nh, nkv, hd, window, splits, cnt.ptr, 0).unwrap();
    be.sync().unwrap();
    let a: Vec<u16> = o_dev.to_vec(rows * nh * hd).unwrap();
    let e = rel_rms(&a, &reference(&q, &kv, &vv, &qpos, nh, nkv, hd, window));
    println!("decode hd{hd} ctx{ctx} rows{rows} splits{splits}: rel_rms vs f32 reference = {e:.2e}");
    assert!(e < 1e-2);
    // separate combine kernel (no merge counters) == merge in the last CTA
    let o_sep = DBuf::zeros(&be, rows * nh * hd * 2).unwrap();
    attn::attn_decode(&g, o_sep.ptr, part.ptr, qd.ptr, kp.ptr, vp.ptr, btd.ptr, off.ptr, qposd.ptr, qposd.ptr, rows, nh, nkv, hd, window, splits, 0, 0).unwrap();
    be.sync().unwrap();
    assert_eq!(o_sep.to_vec::<u16>(rows * nh * hd).unwrap(), a, "split merge kernels differ");
    // sliding-window pool: the same KV in a permuted pool behind a page map (pages wholly outside every row's window map to the NaN-filled null page 0) => bit-identical
    let lo_row = qpos[0] as usize + 1 - window.min(qpos[0] as usize + 1);
    let (mk, mv, pmap) = mapped_pool(&kv, &vv, npg, nkv, hd, if window > 0 { lo_row / 16 } else { 0 });
    let (kpm, vpm, pmd) = (DBuf::from_slice(&be, &mk).unwrap(), DBuf::from_slice(&be, &mv).unwrap(), DBuf::from_slice(&be, &pmap).unwrap());
    let o_map = DBuf::zeros(&be, rows * nh * hd * 2).unwrap();
    attn::attn_decode(&g, o_map.ptr, part.ptr, qd.ptr, kpm.ptr, vpm.ptr, btd.ptr, off.ptr, qposd.ptr, qposd.ptr, rows, nh, nkv, hd, window, splits, cnt.ptr, pmd.ptr).unwrap();
    be.sync().unwrap();
    assert_eq!(o_map.to_vec::<u16>(rows * nh * hd).unwrap(), a, "decode with page map differs");
    println!("decode with sliding page map: bit-identical");
    // prefill: all ctx rows causal
    let tq: Vec<u16> = (0..ctx * nh * hd).map(|_| bf(rng.next())).collect();
    let tqd = DBuf::from_slice(&be, &tq).unwrap();
    let pos: Vec<i32> = (0..ctx as i32).collect();
    let posd = DBuf::from_slice(&be, &pos).unwrap();
    let offs = DBuf::from_slice(&be, &vec![0i32; ctx]).unwrap();
    let (t0, tn) = attn::make_tiles(&vec![0i32; ctx], attn::prefill_rows_per_tile(hd));
    let t0d = DBuf::from_slice(&be, &t0).unwrap(); let tnd = DBuf::from_slice(&be, &tn).unwrap();
    let p_dev = DBuf::zeros(&be, ctx * nh * hd * 2).unwrap();
    attn::attn_prefill(&g, p_dev.ptr, tqd.ptr, kp.ptr, vp.ptr, btd.ptr, offs.ptr, posd.ptr, posd.ptr, t0d.ptr, tnd.ptr, t0.len(), nh, nkv, hd, window, 0).unwrap();
    be.sync().unwrap();
    let a: Vec<u16> = p_dev.to_vec(ctx * nh * hd).unwrap();
    let e = rel_rms(&a, &reference(&tq, &kv, &vv, &pos, nh, nkv, hd, window));
    println!("prefill hd{hd} T{ctx}: rel_rms vs f32 reference = {e:.2e}");
    assert!(e < 1e-2);
    let (mk, mv, pmap) = mapped_pool(&kv, &vv, npg, nkv, hd, 0);
    let (kpm, vpm, pmd) = (DBuf::from_slice(&be, &mk).unwrap(), DBuf::from_slice(&be, &mv).unwrap(), DBuf::from_slice(&be, &pmap).unwrap());
    let p_map = DBuf::zeros(&be, ctx * nh * hd * 2).unwrap();
    attn::attn_prefill(&g, p_map.ptr, tqd.ptr, kpm.ptr, vpm.ptr, btd.ptr, offs.ptr, posd.ptr, posd.ptr, t0d.ptr, tnd.ptr, t0.len(), nh, nkv, hd, window, pmd.ptr).unwrap();
    be.sync().unwrap();
    assert_eq!(p_map.to_vec::<u16>(ctx * nh * hd).unwrap(), a, "prefill with page map differs");
    println!("prefill with sliding page map: bit-identical");
}

#[test]
#[ignore]
fn attention_matches_reference() {
    run(16, 8, 256, 1024, 1500, 5);
    run(16, 2, 512, 0, 700, 1);
}

#[test]
#[ignore]
fn linear_and_head() {
    let cb = dllm_kernels::gpu::CudaBackend::new(0).expect("cuda");
    let be: Arc<dyn Backend> = cb;
    // ---- fp8 gemv: x = 1.0 everywhere, w rows constant (1.0 / 2.0 / 0.5) -> y[t, n] = K * w_n
    let (t, k, n) = (3usize, 256usize, 64usize);
    let vals = [(0x38u8, 1.0f32), (0x40, 2.0), (0x30, 0.5)];
    let mut w = vec![0u8; n * k];
    for r in 0..n { for c in 0..k { w[r * k + c] = vals[r % 3].0; } }
    let wd = DBuf::from_bytes(&be, &w).unwrap();
    let sc = DBuf::from_slice(&be, &vec![1.0f32; n]).unwrap();
    let xq = DBuf::from_bytes(&be, &vec![0x38u8; t * k]).unwrap();
    let xs = DBuf::from_slice(&be, &vec![1.0f32; t]).unwrap();
    let y = DBuf::zeros(&be, t * n * 2).unwrap();
    linear::gemv_pqd(&be, y.ptr, n, xq.ptr, xs.ptr, t, k, &[linear::Seg { w: wd.ptr, s: sc.ptr, n }], false, None).unwrap();
    be.sync().unwrap();
    let raw: Vec<u16> = y.to_vec(t * n).unwrap();
    for tt in 0..t { for r in 0..n {
        let got = f32::from_bits((raw[tt * n + r] as u32) << 16);
        assert_eq!(got, k as f32 * vals[r % 3].1, "gemv t={tt} r={r}");
    } }
    // ---- lm_head argmax: h = 1.0 (bf16), E rows = 1.0 except row 777 = 2.0
    let (v, h) = (1024usize, 64usize);
    let mut e = vec![0x38u8; v * h];
    for c in 0..h { e[777 * h + c] = 0x40; }
    let ed = DBuf::from_bytes(&be, &e).unwrap();
    let rs = DBuf::from_slice(&be, &vec![1.0f32; v]).unwrap();
    let hd = DBuf::from_slice(&be, &vec![0x3f80u16; 2 * h]).unwrap();
    let idx = DBuf::zeros(&be, 8).unwrap();
    let val = DBuf::zeros(&be, 8).unwrap();
    let scr = head::HeadScratch::new(&be).unwrap();
    head::lm_head_argmax(&be, &scr, idx.ptr, val.ptr, hd.ptr, ed.ptr, rs.ptr, 2, v, h, -1, -1, 30.0, false).unwrap();
    be.sync().unwrap();
    let i: Vec<i32> = idx.to_vec(2).unwrap();
    let f: Vec<f32> = val.to_vec(2).unwrap();
    assert_eq!(i, vec![777, 777]);
    assert!((f[0] - 30.0 * (128.0f32 / 30.0).tanh()).abs() < 1e-3, "{f:?}");
    head::lm_head_argmax(&be, &scr, idx.ptr, val.ptr, hd.ptr, ed.ptr, rs.ptr, 2, v, h, 777, -1, 30.0, false).unwrap();
    be.sync().unwrap();
    let i: Vec<i32> = idx.to_vec(2).unwrap();
    assert_eq!(i, vec![0, 0]);
}

#[test]
#[ignore]
fn moe_decode_chain() {
    let be: Arc<dyn Backend> = dllm_kernels::gpu::CudaBackend::new(0).expect("cuda");
    let gpu = Gpu::new(be.clone());
    let (e, h, i, t) = (128usize, 2816usize, 704usize, 2usize);
    // zero fp4 weights with scale byte 0x38 (=1.0) -> all expert outputs are exactly 0
    let rt1 = 2 * i / 16 * (h / 64);
    let rt2 = h / 16 * (i / 64);
    let z = |n: usize| DBuf::zeros(&be, n).unwrap();
    let gu_d = z(e * rt1 * 512);
    let gu_s = DBuf::from_bytes(&be, &vec![0x38u8; e * rt1 * 64]).unwrap();
    let dn_d = z(e * rt2 * 512);
    let dn_s = DBuf::from_bytes(&be, &vec![0x38u8; e * rt2 * 64]).unwrap();
    let gs = DBuf::from_slice(&be, &vec![1.0f32; e]).unwrap();
    let act_g = DBuf::from_slice(&be, &[0.004f32, 0.004]).unwrap();
    let w = Experts { gu_data: gu_d.ptr, gu_scale: gu_s.ptr, gu_gscale: gs.ptr, dn_data: dn_d.ptr, dn_scale: dn_s.ptr, dn_gscale: gs.ptr, act_g: act_g.ptr };
    let bufs = [z(t * 128 * 4), z(t * 8 * 4), z(t * 8 * 4), z(4), z(t * 8 * i / 2), z(t * 8 * i / 16), z(t * 8 * h * 2)];
    let s = MoeDecodeScratch { logits: bufs[0].ptr, ids: bufs[1].ptr, wts: bufs[2].ptr, ticket: bufs[3].ptr, hq: bufs[4].ptr, hs: bufs[5].ptr, opair: bufs[6].ptr, grp: 0 };
    let ones_bf16 = |n: usize| DBuf::from_slice(&be, &vec![0x3f80u16; n]).unwrap();
    let (x, o, w_pa, w_pre, w_pre2, rscale) = (ones_bf16(t * h), ones_bf16(t * h), ones_bf16(h), ones_bf16(h), ones_bf16(h), ones_bf16(h));
    let (mid, rr) = (z(t * h * 2), z(t * h * 2));
    let (fq, fs, gq, gsc) = (z(t * h), z(t * 4), z(t * h / 2), z(t * h / 16));
    gpu.post_attn(mid.ptr, fq.ptr, fs.ptr, gq.ptr, gsc.ptr, rr.ptr, x.ptr, o.ptr, w_pa.ptr, w_pre.ptr, w_pre2.ptr, rscale.ptr, act_g.ptr, t, h, 1e-6).unwrap();
    let proj = ones_bf16(128 * h);
    let pes = ones_bf16(128);
    gpu.moe_router_decode(&s, rr.ptr, proj.ptr, pes.ptr, t, h).unwrap();
    gpu.moe_glu_decode(&w, &s, gq.ptr, gsc.ptr, t, h, i).unwrap();
    gpu.moe_down_decode(&w, &s, t, h, i).unwrap();
    let m = z(t * h * 2); // dense-MLP output = 0
    let out = z(t * h * 2);
    let ls = ones_bf16(1);
    gpu.layer_final(out.ptr, m.ptr, s.opair, s.wts, mid.ptr, w_pa.ptr, w_pa.ptr, w_pa.ptr, ls.ptr, t, h, 1e-6, None).unwrap();
    be.sync().unwrap();
    let raw: Vec<u16> = out.to_vec(t * h).unwrap();
    // mid = x + norm(o) = 1 + 1 = 2 ; branches are 0 -> out = mid * 1 = 2.0 (bf16 0x4000)
    assert!(raw.iter().all(|&b| b == 0x4000), "unexpected layer output {:?}", &raw[..4]);
}