//! GPU test with the real weights (needs a CUDA device with ~24 GB free):
//!   DLLM_MODEL_DIR=<span-02 dir> DLLM_DRAFT_DIR=<gemma-4-26B-A4B-it-assistant dir> cargo test --release -p dllm-model -- --ignored
//! Speculative decoding is lossless: the tokens of a speculative thought (MTP draft + verify, adaptive k) equal one-row greedy decoding token for token,
//! for one branch alone and for several branches verified in the same batch. Speculative number answers equal the plain answer step (one-row forward +
//! full logits) in tokens and log-probabilities, bit for bit.
use dllm_core::config::{AssistantConfig, ModelConfig};
use dllm_core::loader::WeightStore;
use dllm_core::tokenizer::{state_text, Tok};
use dllm_kernels::Gpu;
use dllm_model::draft::{Draft, DraftVocab};
use dllm_model::gpu_model::{GpuModel, GpuOpts, RowBatch, PAGE};
use dllm_model::spec::{greedy_logprob, Spec, SpecReq};
use std::sync::Arc;

const MAX_LEN: usize = 4096;
const PAGES: usize = 1024;

fn env(k: &str) -> String {
    std::env::var(k).unwrap_or_else(|_| panic!("set {k}"))
}

/// prompt rows 0..n of one sequence on `pages` (prefill: one forward)
fn prefill(m: &mut GpuModel, toks: &[i32], pages: &[i32]) -> i32 {
    let mut b = RowBatch { block_table: pages.to_vec(), ..Default::default() };
    for (p, &t) in toks.iter().enumerate() {
        b.tokens.push(t);
        b.pos.push(p as i32);
        b.slot.push(pages[p / PAGE] * PAGE as i32 + (p % PAGE) as i32);
        b.bt_off.push(0);
        b.kv_hi.push(p as i32);
        b.kv_hi_sliding.push(p as i32);
    }
    m.forward(&b).unwrap();
    m.argmax_rows(&[toks.len() - 1], -1, -1).unwrap()[0]
}

/// `n` greedy tokens after `first`, one row per forward (the reference)
fn greedy(m: &mut GpuModel, first: i32, start: usize, pages: &[i32], n: usize) -> Vec<i32> {
    let mut out = vec![first];
    while out.len() < n {
        let pos = start + out.len() - 1;
        m.forward(&RowBatch::decode(&[*out.last().unwrap()], &[(pages, pos)])).unwrap();
        out.push(m.greedy_first_rows(1, -1, -1).unwrap()[0]);
    }
    out
}

fn load() -> (GpuModel, Spec, Tok) {
    let dir = std::path::PathBuf::from(env("DLLM_MODEL_DIR"));
    let draft_dir = std::path::PathBuf::from(env("DLLM_DRAFT_DIR"));
    let cfg = ModelConfig::from_file(dir.join("config.json")).unwrap();
    let st = WeightStore::open(&dir).unwrap();
    let gpu = Gpu::new(dllm_kernels::gpu::CudaBackend::new(0).unwrap() as Arc<dyn dllm_kernels::Backend>);
    let m = GpuModel::load(gpu, &st, &cfg.text, "model.language_model.", GpuOpts { max_rows: 2048, num_pages: PAGES, max_pos: MAX_LEN + 64, sliding_pages: PAGES }).unwrap();
    // identity sliding map: every full page keeps its own sliding page
    m.set_sliding_map(0, &(1..=PAGES as i32).collect::<Vec<_>>()).unwrap();
    let acfg = AssistantConfig::from_file(draft_dir.join("config.json")).unwrap();
    let vocab = DraftVocab::from_json_file(dir.join("draft_vocab_16384.json")).unwrap();
    let draft = Draft::load(m.gpu.clone(), &WeightStore::open(&draft_dir).unwrap(), &acfg, &cfg.text, &vocab, MAX_LEN + 64).unwrap();
    let spec = Spec::new(&m, draft, MAX_LEN, 64).unwrap();
    (m, spec, Tok::from_file(dir.join("tokenizer.json")).unwrap())
}

#[test]
#[ignore]
fn speculative_thoughts_equal_greedy() {
    let (mut m, mut spec, tok) = load();
    let questions = [
        ("A train leaves at 9:40 and arrives at 13:05. The ticket says the trip takes 3 h 20 min.", "Is the ticket right?", ["Yes", "No, it is longer", "No, it is shorter"]),
        ("Lena has twice as many apples as Omar. Omar has 7 more than Pia. Pia has 5.", "How many apples does Lena have?", ["24", "17", "34"]),
        ("{\"order\": 4417, \"items\": [\"lamp\", \"cable\"], \"status\": \"returned\", \"refund\": null}", "What should support do next?", ["Issue the refund", "Ask for photos", "Close the ticket"]),
    ];
    let n = 160;
    let mut prompts = vec![];
    let mut want = vec![];
    for (i, (state, q, opts)) in questions.iter().enumerate() {
        let opts: Vec<String> = opts.iter().map(|s| s.to_string()).collect();
        let ids: Vec<i32> = tok.build(&state_text(&serde_json::Value::String(state.to_string())), q, &opts, 0).unwrap().unwrap().ids.iter().map(|&t| t as i32).collect();
        let pages: Vec<i32> = ((i * 64) as i32..(i * 64 + 64) as i32).collect();
        let first = prefill(&mut m, &ids, &pages);
        want.push(greedy(&mut m, first, ids.len(), &pages, n));
        prompts.push((ids, pages, first));
    }
    // one branch at a time, then all three branches in the same verify batches
    for batch in [vec![0usize], vec![1], vec![2], vec![0, 1, 2]] {
        let mut got: Vec<Vec<i32>> = vec![];
        for &i in &batch {
            let (ids, pages, _) = &prompts[i];
            assert_eq!(prefill(&mut m, ids, pages), prompts[i].2);
            got.push(vec![prompts[i].2]);
        }
        while got.iter().any(|g| g.len() < n) {
            let reqs: Vec<SpecReq> = batch.iter().zip(&got).map(|(&i, g)| SpecReq { uid: i + 100 * batch.len(), pages: &prompts[i].1, pos: prompts[i].0.len() + g.len() - 1, tok: *g.last().unwrap() }).collect();
            let out = spec.step(&mut m, &reqs).unwrap();
            for (g, o) in got.iter_mut().zip(out) {
                g.extend(o);
            }
        }
        for (&i, g) in batch.iter().zip(&got) {
            assert_eq!(&g[..n], &want[i][..], "prompt {i} (batch {batch:?}): speculative thought differs from greedy");
        }
        eprintln!("batch {batch:?}: {} tokens each, identical to greedy", n);
    }
}

/// the plain answer step: one-row forward, greedy token + log-probability of the full logits, until a stop token or 32 tokens
fn plain_answer(m: &mut GpuModel, first: (i32, f64), start: usize, pages: &[i32]) -> Vec<(i32, f64)> {
    let mut out = vec![first];
    while out.len() < 32 {
        let pos = start + out.len() - 1;
        m.forward(&RowBatch::decode(&[out.last().unwrap().0], &[(pages, pos)])).unwrap();
        let t = greedy_logprob(&m.logits_rows(&[0]).unwrap()[0]);
        if [1, 101, 106].contains(&t.0) {
            break;
        }
        out.push(t);
    }
    out
}

#[test]
#[ignore]
fn speculative_number_answers_equal_greedy() {
    let (mut m, mut spec, tok) = load();
    let (mk, stop) = (101, [1, 101, 106]);
    let build = |state: &str, q: &str| -> Vec<i32> { tok.build_number(&state_text(&serde_json::Value::String(state.into())), q, 0).unwrap().ids.iter().map(|&t| t as i32).collect() };
    // padded so that the answer rows cross 1024 positions (where the one-row decode forward changes its KV splits)
    let long = (100..1000).map(|n| format!("{} The factory makes 2,345 widgets per day.", "Filler. ".repeat(n))).find(|st| build(st, "How many widgets does it make in 365 days?").len() + 1 >= 1021).unwrap();
    let qs: Vec<(String, &str)> = vec![
        ("A warehouse has 17 shelves with 23 boxes each. 45 boxes are shipped out, then 3 deliveries of 12 boxes arrive.".into(), "How many boxes are in the warehouse now?"),
        ("".into(), "What is 123456 * 7?"),
        ("".into(), "What is 4321 * 5678?"),
        ("{\"population\": 1234567, \"growth_rate\": 0.02}".into(), "What is the population after one year of growth?"),
        ("".into(), "What is pi to 10 decimal places?"),
        (long, "How many widgets does it make in 365 days?"),
    ];
    let (mut n_multi, mut n_tok, mut n_cyc) = (0, 0, 0);
    for (i, (state, q)) in qs.iter().enumerate() {
        let prompt = build(state, q);
        let pages: Vec<i32> = ((i * 128) as i32..(i * 128 + 128) as i32).collect();
        for think in [false, true] {
            // context = prompt (+ a greedy thought of up to 64 tokens) + end-of-thought; its last row gives the first answer token
            let mut ctx = prompt.clone();
            if think {
                prefill(&mut m, &prompt, &pages);
                let first = m.argmax_rows(&[prompt.len() - 1], mk, 1).unwrap()[0];
                let th: Vec<i32> = greedy(&mut m, first, prompt.len(), &pages, 64).into_iter().take_while(|t| !stop.contains(t)).collect();
                ctx.extend(th);
            }
            ctx.push(mk);
            let row = ctx.len() - 1;
            prefill(&mut m, &ctx, &pages);
            let first = greedy_logprob(&m.logits_rows(&[row]).unwrap()[0]);
            assert!(!stop.contains(&first.0));
            let t0 = std::time::Instant::now();
            let want = plain_answer(&mut m, first, ctx.len(), &pages);
            let ms_plain = t0.elapsed().as_secs_f64() * 1e3;
            // speculative: same context, the draft seeded from the end-of-thought row
            prefill(&mut m, &ctx, &pages);
            spec.prime(&m, 7, row, ctx.len(), first.0).unwrap();
            let mut got = vec![first];
            let mut cycles = 0;
            let t0 = std::time::Instant::now();
            'dec: while got.len() < 32 {
                let r = SpecReq { uid: 7, pages: &pages, pos: ctx.len() + got.len() - 1, tok: got.last().unwrap().0 };
                cycles += 1;
                for t in spec.answer_step(&mut m, &[r]).unwrap().remove(0) {
                    if stop.contains(&t.0) || got.len() >= 32 {
                        break 'dec;
                    }
                    got.push(t);
                }
            }
            let ms_spec = t0.elapsed().as_secs_f64() * 1e3;
            let bits = |v: &[(i32, f64)]| v.iter().map(|x| (x.0, x.1.to_bits())).collect::<Vec<_>>();
            assert_eq!(bits(&got), bits(&want), "question {i} think {think}: speculative answer differs from the plain answer step");
            // one plain step per token after the first, the end-of-turn included
            n_multi += (cycles < got.len()) as usize;
            (n_tok, n_cyc) = (n_tok + got.len(), n_cyc + cycles);
            eprintln!("question {i} think {think}: ctx {} answer {:?} ({} tokens + end, {cycles} cycles) identical; {ms_plain:.1} ms plain, {ms_spec:.1} ms speculative", ctx.len(), tok.decode(&got.iter().map(|x| x.0 as u32).collect::<Vec<_>>()), got.len());
        }
    }
    eprintln!("{n_tok} answer tokens (after the first, end-of-turn included) in {n_cyc} speculative cycles: {:.2} tokens per cycle", n_tok as f64 / n_cyc as f64);
    assert!(n_multi > 0, "no answer cycle accepted a draft");
}