//! Speculative decoding of thoughts on the SHARED paged KV pool of the scheduler. `Spec` owns the draft head and a few small device buffers; every call
//! receives the target `GpuModel` and the thinking branches `(pages, pos, tok)`. Per step (one cycle): draft chain from (x0 = tok, h0 = remembered hidden
//! state of the previous accepted row) -> ONE batched verify forward of B(k+1) rows over the branches' own pages (KV appended in place at pos..pos+k) ->
//! greedy acceptance -> accepted tokens returned to the scheduler.
//! Lossless: a token is emitted only if it equals the target's argmax at the verify row before it, and the verify forward uses row-invariant launch
//! parameters, so the emitted tokens equal one-row greedy decoding. Rollback = the scheduler advances `pos` only by the accepted tokens; KV written for
//! rejected rows lies beyond it and is overwritten by the next step (never attended: kv_hi = position).
//! A branch seen for the first time (or after a gap) has no hidden state: it is verified with garbage drafts (they fail) and its h0 is kept for the next step,
//! unless `prime` seeded it from the target row that produced its token.
//! Number answers (`answer_step`) re-read every accepted token from the target's full logits, as the plain answer step does, for its log-probability.
//! Sources: SGLang eagle `verify_tree_greedy` (chain case), vLLM `rejection_greedy_sample` + `llm_base_proposer` (see cuda/spec.cu).
use crate::draft::{Draft, DraftCtx};
use crate::gpu_model::{DevMeta, GpuModel};
use anyhow::{ensure, Result};
use dllm_kernels::spec::{EMIT_W, KMAX};
use dllm_kernels::DBuf;
use std::collections::HashMap;

/// adaptive draft depth: running estimate of the conditional acceptance of draft position j given positions < j were accepted, and a verify cost model;
/// k minimises the expected time per emitted token
#[derive(Debug, Clone)]
pub struct KPolicy {
    pub att: [f64; KMAX],
    pub succ: [f64; KMAX],
    /// prior pseudo-counts: p_j ~ prior_p with weight prior_n
    pub prior_p: f64,
    pub prior_n: f64,
    pub kmin: usize,
    pub kmax: usize,
    /// ms per verify forward as a function of rows: piecewise-linear table (rows, ms)
    pub verify_ms: Vec<(f64, f64)>,
    /// ms per draft step (batch of <= 16 rows)
    pub draft_ms: f64,
}
impl Default for KPolicy {
    /// measured on the RTX 5090 with CUDA graphs
    fn default() -> Self {
        KPolicy { att: [0.0; KMAX], succ: [0.0; KMAX], prior_p: 0.7, prior_n: 8.0, kmin: 1, kmax: 6, verify_ms: vec![(1.0, 3.2), (2.0, 3.6), (3.0, 3.8), (5.0, 4.0), (7.0, 4.2), (16.0, 6.5)], draft_ms: 0.28 }
    }
}
impl KPolicy {
    pub fn p(&self, j: usize) -> f64 {
        (self.succ[j] + self.prior_p * self.prior_n) / (self.att[j] + self.prior_n)
    }
    pub fn verify_cost(&self, rows: usize) -> f64 {
        let r = rows as f64;
        let t = &self.verify_ms;
        if r <= t[0].0 {
            return t[0].1 * r / t[0].0;
        }
        for w in t.windows(2) {
            if r <= w[1].0 {
                return w[0].1 + (w[1].1 - w[0].1) * (r - w[0].0) / (w[1].0 - w[0].0);
            }
        }
        let (a, b) = (t[t.len() - 2], t[t.len() - 1]);
        b.1 + (b.1 - a.1) * (r - b.0) / (b.0 - a.0)
    }
    /// expected emitted tokens per sequence per cycle with k drafts
    pub fn expected_tokens(&self, k: usize) -> f64 {
        let (mut e, mut prod) = (1.0, 1.0);
        for j in 0..k {
            prod *= self.p(j);
            e += prod;
        }
        e
    }
    /// k minimising (verify(B(k+1)) + k * draft) / E[tokens]; the verify batch stays within 16 rows (decode GEMV path, row-invariant numerics)
    pub fn choose(&self, b: usize) -> usize {
        let mut best = (self.kmin, f64::MAX);
        let kcap = (16 / b.max(1)).saturating_sub(1).max(1);
        for k in self.kmin..=self.kmax.min(kcap) {
            let cost = (self.verify_cost(b * (k + 1)) + if k > 0 { k as f64 * self.draft_ms } else { 0.0 }) / self.expected_tokens(k);
            if cost < best.1 {
                best = (k, cost);
            }
        }
        best.0
    }
    pub fn update(&mut self, k: usize, accepted: usize) {
        for j in 0..k {
            if j <= accepted {
                self.att[j] += 1.0;
            }
            if j < accepted {
                self.succ[j] += 1.0;
            }
        }
    }
}

/// one thinking branch: feed `tok` at position `pos` (its KV slot is page `pages[pos / 16]`); `uid` = stable branch id
#[derive(Debug, Clone, Default)]
pub struct SpecReq<'a> {
    pub uid: usize,
    pub pages: &'a [i32],
    pub pos: usize,
    pub tok: i32,
}

pub struct Spec {
    pub draft: Draft,
    pub policy: KPolicy,
    /// draft depth of number answers (digits are accepted at a different rate than thought text)
    pub answer_policy: KPolicy,
    splits: usize,
    maxp: usize,
    hb: usize,
    hpool_n: usize,
    len: DBuf,
    x0: DBuf,
    gen: DBuf,
    done: DBuf,
    drafts: DBuf,
    h0: DBuf,
    dpos: DBuf,
    btoff_c: DBuf,
    seq_bt: DBuf,
    tokens: DBuf,
    pos: DBuf,
    slot: DBuf,
    bt_off: DBuf,
    kv_hi: DBuf,
    kv_hi_sl: DBuf,
    g: DBuf,
    am_val: DBuf,
    emit: DBuf,
    stop: DBuf,
    ctr: DBuf,
    hpool: DBuf,
    hslot: DBuf,
    head_sc: dllm_kernels::head::HeadScratch,
    /// (k, b, KV splits) -> captured cycle
    graphs: HashMap<(usize, usize, usize), Box<dyn dllm_kernels::GraphExec>>,
    graphs_ok: bool,
    /// uid -> (hpool slot, expected next pos, expected next token, last used step)
    state: HashMap<usize, (usize, usize, i32, u64)>,
    free: Vec<usize>,
    step_no: u64,
}

fn i32s(v: &[i32]) -> &[u8] {
    unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, std::mem::size_of_val(v)) }
}

/// most branches in one verify batch (k = 0)
pub const MAX_B: usize = 16;

impl Spec {
    /// `max_len`: max positions of a branch (block table width); `hpool_n`: remembered hidden states (>= max concurrent branches)
    pub fn new(tgt: &GpuModel, draft: Draft, max_len: usize, hpool_n: usize) -> Result<Self> {
        let be = tgt.gpu.be.clone();
        let hb = tgt.cfg.hidden_size;
        let maxp = (max_len + 15) / 16 + 2;
        let b = MAX_B;
        let z = |n: usize| DBuf::zeros(&be, n);
        let rows_max = 16usize;
        Ok(Spec {
            draft,
            policy: KPolicy::default(),
            answer_policy: KPolicy::default(),
            splits: 4,
            maxp,
            hb,
            hpool_n,
            len: z(b * 4)?,
            x0: z(b * 4)?,
            gen: z(b * 4)?,
            done: z(b * 4)?,
            drafts: z(b * KMAX * 4)?,
            h0: z(b * hb * 2)?,
            dpos: z(b * 4)?,
            btoff_c: DBuf::from_slice(&be, &(0..b).map(|s| (s * maxp) as i32).collect::<Vec<_>>())?,
            seq_bt: z(b * maxp * 4)?,
            tokens: z(rows_max * 4)?,
            pos: z(rows_max * 4)?,
            slot: z(rows_max * 4)?,
            bt_off: z(rows_max * 4)?,
            kv_hi: z(rows_max * 4)?,
            kv_hi_sl: z(rows_max * 4)?,
            g: z(rows_max * 4)?,
            am_val: z(rows_max * 4)?,
            emit: z(b * EMIT_W * 4)?,
            stop: DBuf::from_slice(&be, &[-1i32; 4])?,
            ctr: z(16)?,
            hpool: z(hpool_n * hb * 2)?,
            hslot: z(b * 4)?,
            head_sc: dllm_kernels::head::HeadScratch::new(&be)?,
            graphs: HashMap::new(),
            graphs_ok: true,
            state: HashMap::new(),
            free: (0..hpool_n).rev().collect(),
            step_no: 0,
        })
    }

    /// extra positions beyond `pos` that one step may write
    pub fn max_lookahead(&self) -> usize {
        self.policy.kmax
    }

    fn meta(&self) -> DevMeta {
        DevMeta { tokens: self.tokens.ptr, pos: self.pos.ptr, slot: self.slot.ptr, bt_off: self.bt_off.ptr, kv_hi: self.kv_hi.ptr, kv_hi_sl: self.kv_hi_sl.ptr, bt: self.seq_bt.ptr }
    }

    fn body(&self, tgt: &GpuModel, b: usize, k: usize, splits: usize) -> Result<()> {
        tgt.suspend_graphs(true);
        // row-invariant launch parameters only around our own launches (the decode graph buckets / scratch sizes assume the default policy)
        dllm_kernels::set_row_invariant(true);
        let r = self.body_inner(tgt, b, k, splits);
        dllm_kernels::set_row_invariant(false);
        tgt.suspend_graphs(false);
        r
    }

    fn body_inner(&self, tgt: &GpuModel, b: usize, k: usize, splits: usize) -> Result<()> {
        let gpu = &tgt.gpu;
        let hb = self.hb;
        let maxp = self.maxp;
        gpu.spec_h_gather(self.h0.ptr, self.hpool.ptr, self.hslot.ptr, b, hb)?;
        if k > 0 {
            gpu.spec_draft_pos(self.dpos.ptr, self.len.ptr, b)?;
            let ctx = DraftCtx { bt: self.seq_bt.ptr, bt_off: self.btoff_c.ptr, dpos: self.dpos.ptr, splits };
            self.draft.chain(tgt, &ctx, self.x0.ptr, self.h0.ptr, self.drafts.ptr, b, k)?;
        }
        let r = b * (k + 1);
        gpu.spec_build_verify(self.tokens.ptr, self.pos.ptr, self.slot.ptr, self.bt_off.ptr, self.kv_hi.ptr, self.kv_hi_sl.ptr, self.seq_bt.ptr, self.x0.ptr, self.len.ptr, self.drafts.ptr, b, k, maxp)?;
        tgt.forward_dev(r, true, splits, self.meta())?;
        let (v, hdim) = (tgt.cfg.vocab_size, tgt.cfg.hidden_size);
        let (w, sc, _) = tgt.embed_desc();
        let mut c0 = 0;
        while c0 < r {
            let m = (r - c0).min(8);
            dllm_kernels::head::lm_head_argmax(&gpu.be, &self.head_sc, self.g.ptr + (c0 * 4) as u64, self.am_val.ptr + (c0 * 4) as u64, tgt.hidden_ptr() + (c0 * hdim * 2) as u64, w, sc, m, v, hdim, -1, -1, tgt.softcap(), false)?;
            c0 += m;
        }
        gpu.spec_accept(self.emit.ptr, self.ctr.ptr, 1, self.len.ptr, self.x0.ptr, self.gen.ptr, self.done.ptr, self.h0.ptr, self.g.ptr, self.drafts.ptr, tgt.hidden_ptr(), self.stop.ptr, 1 << 30, b, k, hb)?;
        gpu.spec_h_scatter(self.hpool.ptr, self.h0.ptr, self.hslot.ptr, b, hb)?;
        Ok(())
    }

    fn pick_k(pol: &KPolicy, b: usize) -> usize {
        let cap = (16 / b).saturating_sub(1).min(KMAX - 1);
        pol.choose(b).min(cap)
    }

    /// one speculative cycle for `reqs` (any number; processed in chunks whose verify batch is <= 16 rows). Returns the accepted tokens per request (>= 1 each).
    pub fn step(&mut self, tgt: &mut GpuModel, reqs: &[SpecReq]) -> Result<Vec<Vec<i32>>> {
        let mut out: Vec<Vec<i32>> = vec![vec![]; reqs.len()];
        let mut i = 0;
        while i < reqs.len() {
            let b = (reqs.len() - i).min(MAX_B);
            let k = Self::pick_k(&self.policy, b);
            self.step_chunk(tgt, &reqs[i..i + b], k, self.splits, false, &mut out[i..i + b])?;
            i += b;
        }
        Ok(out)
    }

    /// number answers: one speculative cycle per request, returning (token, log-probability) per accepted row from the target's FULL logits of that verify
    /// row (`greedy_logprob`, the plain answer step's readout), cut after a row whose full-logit token differs from the fused argmax (a near tie: later
    /// rows were fed the other token). The verify forward runs with the KV splits of the one-row decode forward at the same positions (a cycle stays
    /// within one 1024-position band), so its rows equal the plain answer step's bit for bit.
    pub fn answer_step(&mut self, tgt: &mut GpuModel, reqs: &[SpecReq]) -> Result<Vec<Vec<(i32, f64)>>> {
        let mut out = Vec::with_capacity(reqs.len());
        for chunk in reqs.chunks(MAX_B) {
            let b = chunk.len();
            let k = Self::pick_k(&self.answer_policy, b).min(chunk.iter().map(|r| 1023 - r.pos % 1024).min().unwrap());
            let splits = tgt.decode_splits(chunk.iter().map(|r| r.pos + k + 1).max().unwrap());
            let mut toks = vec![vec![]; b];
            self.step_chunk(tgt, chunk, k, splits, true, &mut toks)?;
            let rows: Vec<usize> = toks.iter().enumerate().flat_map(|(s, t)| (0..t.len()).map(move |j| s * (k + 1) + j)).collect();
            let lg = tgt.logits_rows(&rows)?;
            let scored: Vec<(i32, f64)> = std::thread::scope(|s| {
                let h: Vec<_> = lg.iter().map(|z| s.spawn(move || greedy_logprob(z))).collect();
                h.into_iter().map(|h| h.join().unwrap()).collect()
            });
            let mut off = 0;
            for t in toks {
                let mut acc = Vec::with_capacity(t.len());
                for (d, &(g, lp)) in t.iter().zip(&scored[off..]) {
                    acc.push((g, lp));
                    if g != *d {
                        break;
                    }
                }
                off += t.len();
                out.push(acc);
            }
        }
        Ok(out)
    }

    /// seed branch `uid`, which continues with `tok` at `pos`, with hidden row `row` of the target's last forward (the row that produced `tok`):
    /// its first step then drafts instead of verifying garbage drafts
    pub fn prime(&mut self, tgt: &GpuModel, uid: usize, row: usize, pos: usize, tok: i32) -> Result<()> {
        let slot = match self.state.get(&uid) {
            Some(&(sl, _, _, _)) => sl,
            None => self.alloc_slot(),
        };
        let hb2 = self.hb * 2;
        tgt.gpu.be.copy_d2d(self.hpool.ptr + (slot * hb2) as u64, tgt.hidden_ptr() + (row * hb2) as u64, hb2)?;
        self.state.insert(uid, (slot, pos, tok, self.step_no));
        Ok(())
    }

    /// `answer`: the cycle decodes number answers (acceptance statistics of `answer_policy`)
    fn step_chunk(&mut self, tgt: &mut GpuModel, reqs: &[SpecReq], k: usize, splits: usize, answer: bool, out: &mut [Vec<i32>]) -> Result<()> {
        let b = reqs.len();
        ensure!(b >= 1 && b * (k + 1) <= 16);
        self.step_no += 1;
        let be = tgt.gpu.be.clone();
        let maxp = self.maxp;
        // host -> device state of this step
        let mut len = vec![0i32; b];
        let mut x0 = vec![0i32; b];
        let mut slots = vec![0i32; b];
        let mut tab = vec![0i32; b * maxp];
        let mut fresh = vec![false; b];
        for (s, r) in reqs.iter().enumerate() {
            ensure!(r.pages.len() <= maxp && (r.pos + k) / 16 < r.pages.len(), "branch pages ({}) do not cover pos {} + lookahead {}", r.pages.len(), r.pos, k);
            len[s] = r.pos as i32;
            x0[s] = r.tok;
            tab[s * maxp..s * maxp + r.pages.len()].copy_from_slice(r.pages);
            let valid = matches!(self.state.get(&r.uid), Some(&(_, np, nt, _)) if np == r.pos && nt == r.tok);
            fresh[s] = !valid;
            let slot = match self.state.get(&r.uid) {
                Some(&(sl, _, _, _)) => sl,
                None => self.alloc_slot(),
            };
            self.state.insert(r.uid, (slot, r.pos, r.tok, self.step_no));
            slots[s] = slot as i32;
        }
        be.upload(self.len.ptr, i32s(&len))?;
        be.upload(self.x0.ptr, i32s(&x0))?;
        be.upload(self.hslot.ptr, i32s(&slots))?;
        be.upload(self.seq_bt.ptr, i32s(&tab))?;
        be.upload(self.ctr.ptr, &[0u8; 4])?;
        be.upload(self.done.ptr, &vec![0u8; b * 4])?;
        let key = (k, b, splits);
        let mut launched = false;
        if self.graphs_ok {
            if !self.graphs.contains_key(&key) {
                match be.graph_begin() {
                    Ok(()) => {
                        let r = self.body(tgt, b, k, splits);
                        let g = be.graph_end();
                        r?;
                        self.graphs.insert(key, g?);
                    }
                    Err(_) => self.graphs_ok = false,
                }
            }
            if let Some(g) = self.graphs.get(&key) {
                g.launch()?;
                launched = true;
            }
        }
        if !launched {
            self.body(tgt, b, k, splits)?;
        }
        let emit: Vec<i32> = self.emit.to_vec(b * EMIT_W)?;
        for s in 0..b {
            let e = &emit[s * EMIT_W..(s + 1) * EMIT_W];
            let (n, a) = (e[0] as usize, e[1] as usize);
            out[s] = e[3..3 + n].to_vec();
            if !fresh[s] {
                if answer { &mut self.answer_policy } else { &mut self.policy }.update(k, a);
            }
            let uid = reqs[s].uid;
            if let Some(st) = self.state.get_mut(&uid) {
                st.1 = reqs[s].pos + a + 1;
                st.2 = *out[s].last().unwrap();
            }
        }
        // bound the remembered-state table (branches that vanished)
        if self.state.len() > self.hpool_n - 16 {
            let cutoff = self.step_no.saturating_sub(64);
            let stale: Vec<usize> = self.state.iter().filter(|(_, v)| v.3 < cutoff).map(|(k, _)| *k).collect();
            for u in stale {
                if let Some(v) = self.state.remove(&u) {
                    self.free.push(v.0);
                }
            }
        }
        Ok(())
    }

    fn alloc_slot(&mut self) -> usize {
        if let Some(s) = self.free.pop() {
            return s;
        }
        // evict the least recently used remembered state
        let (&u, _) = self.state.iter().min_by_key(|(_, v)| v.3).unwrap();
        self.state.remove(&u).unwrap().0
    }
}

/// greedy token of a full logits row (first maximum) and its log-probability. The exponentials run on 8 threads and are summed in index order:
/// the same f64 as the sequential sum (one core needs ~2 ms for a 262K-entry row)
pub fn greedy_logprob(z: &[f32]) -> (i32, f64) {
    let mut b = 0;
    for (i, &v) in z.iter().enumerate() {
        if v > z[b] {
            b = i;
        }
    }
    let m = z[b] as f64;
    let mut e = vec![0f64; z.len()];
    let n = z.len().div_ceil(8);
    std::thread::scope(|s| {
        for (o, x) in e.chunks_mut(n).zip(z.chunks(n)) {
            s.spawn(move || o.iter_mut().zip(x).for_each(|(o, &v)| *o = (v as f64 - m).exp()));
        }
    });
    let lse = m + e.iter().sum::<f64>().ln();
    (b as i32, z[b] as f64 - lse)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn greedy_logprob_equals_the_sequential_sum() {
        let z: Vec<f32> = (0..262_147u32).map(|i| ((i.wrapping_mul(2654435761) >> 8) as f32 / 65536.0).sin() * 30.0).collect();
        let m = z.iter().cloned().fold(f32::NEG_INFINITY, f32::max) as f64;
        let b = z.iter().position(|&v| v as f64 == m).unwrap();
        let lp = z[b] as f64 - (m + z.iter().map(|&v| (v as f64 - m).exp()).sum::<f64>().ln());
        assert_eq!(greedy_logprob(&z), (b as i32, lp));
    }
    #[test]
    fn k_policy_prefers_deeper_drafts_when_acceptance_is_high() {
        let mut p = KPolicy::default();
        let k0 = p.choose(1);
        for _ in 0..500 {
            p.update(6, 6);
        }
        assert!(p.choose(1) >= k0 && p.choose(1) == 6);
        let mut q = KPolicy::default();
        for _ in 0..500 {
            q.update(6, 0);
        }
        assert_eq!(q.choose(1), 1);
        // the verify batch never exceeds 16 rows
        for b in 1..=16 {
            assert!(b * (p.choose(b) + 1) <= 16 || p.choose(b) == 1);
        }
    }
}