//! EAGLE-1 autoregressive draft head (Li et al. 2024, arXiv:2401.15077).
//!
//! EAGLE drafts multiple future tokens by running a single small draft head
//! *autoregressively* over a fused feature stream: at each step the head
//! consumes a fusion of the current hidden state and the embedding of the
//! previously drafted token, and produces both the next hidden state and a
//! logits distribution over the vocabulary for the next token.  This is in
//! contrast to Medusa, which uses `K` independent heads each predicting a
//! fixed future offset.
//!
//! # Architecture
//!
//! Each [`EagleDraftHead`] layer computes:
//!
//! ```text
//! fused   = fuse(hidden [H], token_embedding [H])          // per FusionMethod
//! act     = SiLU(W1 @ fused + b1)                          // [H]
//! hidden' = W_h @ act + b_h                                // [H] next hidden state
//! logits  = W_v @ act + b_v                                // [V] vocab logits
//! ```
//!
//! [`EagleHeads`] wraps a stack of `num_draft_layers` such heads plus a small
//! token-embedding table, and drives the autoregressive draft loop described
//! in the EAGLE paper:
//!
//! 1. Fuse the current hidden state with the embedding of the last token.
//! 2. Run the draft head to get a new hidden state and vocab logits.
//! 3. Pick the top token (argmax) — or the best of a top-k list for tree
//!    construction — feed it back in as the "last token" for the next step.
//! 4. Repeat for `steps` iterations.
//!
//! # Example
//!
//! ```rust
//! use rtx_inference::eagle::{EagleHeads, EagleHeadsConfig};
//!
//! let config = EagleHeadsConfig::default();
//! let heads  = EagleHeads::new(config);
//! let hidden = vec![0.0_f32; 64];
//! let draft  = heads.draft(&hidden, &[1, 2, 3], 4);
//! assert_eq!(draft.len(), 4);
//! ```

use crate::InferenceResult;
use crate::speculative::{EagleConfig, EagleDraftModel, FusionMethod, Token};

// ── LCG pseudo-random number generator ────────────────────────────────────────

/// Minimal LCG PRNG used for reproducible weight initialisation (no external deps).
///
/// Mirrors the generator in [`crate::medusa`] so seeded init is deterministic
/// and reproducible across runs without pulling in `rand`.
struct Lcg(u64);

impl Lcg {
    fn new(seed: u64) -> Self {
        Self(seed ^ 0x1234_5678_9abc_def0)
    }

    /// Advance the state and return a value in `[0, 1)`.
    fn next_f32(&mut self) -> f32 {
        self.0 = self
            .0
            .wrapping_mul(6_364_136_223_846_793_005)
            .wrapping_add(1_442_695_040_888_963_407);
        let bits = 0x3f80_0000_u32 | ((self.0 >> 41) as u32 & 0x007f_ffff);
        f32::from_bits(bits) - 1.0
    }

    /// Return a value uniformly in `[-scale, scale]`.
    fn next_scaled(&mut self, scale: f32) -> f32 {
        (self.next_f32() * 2.0 - 1.0) * scale
    }
}

// ── Fusion helper ─────────────────────────────────────────────────────────────

/// Fuse a hidden state with a token embedding per [`FusionMethod`].
///
/// * [`FusionMethod::Add`] and [`FusionMethod::Attention`] return a vector of
///   length `hidden_dim` (elementwise combination).
/// * [`FusionMethod::Concat`] returns a vector of length `2 * hidden_dim`.
///
/// `Attention` uses a lightweight scalar-gate approximation: the gate is the
/// sigmoid of the dot product between `hidden` and `embedding`, scaled by
/// `1/sqrt(hidden_dim)` (a single-head, single-query attention score), and is
/// used to interpolate between `hidden` and `embedding`.  This keeps the
/// draft head architecture-agnostic to the fusion strategy without requiring
/// a full multi-head attention implementation for a draft-only path.
fn fuse_features(hidden: &[f32], embedding: &[f32], method: FusionMethod) -> Vec<f32> {
    debug_assert_eq!(
        hidden.len(),
        embedding.len(),
        "hidden/embedding dims must match"
    );
    match method {
        FusionMethod::Concat => {
            let mut out = Vec::with_capacity(hidden.len() + embedding.len());
            out.extend_from_slice(hidden);
            out.extend_from_slice(embedding);
            out
        }
        FusionMethod::Add => hidden
            .iter()
            .zip(embedding.iter())
            .map(|(h, e)| h + e)
            .collect(),
        FusionMethod::Attention => {
            let dim = hidden.len().max(1);
            let score: f32 = hidden
                .iter()
                .zip(embedding.iter())
                .map(|(h, e)| h * e)
                .sum::<f32>()
                / (dim as f32).sqrt();
            let gate = 1.0 / (1.0 + (-score).exp());
            hidden
                .iter()
                .zip(embedding.iter())
                .map(|(h, e)| gate * h + (1.0 - gate) * e)
                .collect()
        }
    }
}

/// The input dimension a draft head layer expects for a given `hidden_dim`
/// and [`FusionMethod`].
fn fused_dim(hidden_dim: usize, method: FusionMethod) -> usize {
    match method {
        FusionMethod::Concat => hidden_dim * 2,
        FusionMethod::Add | FusionMethod::Attention => hidden_dim,
    }
}

// ── EagleDraftHead ────────────────────────────────────────────────────────────

/// One autoregressive EAGLE draft head layer.
///
/// Consumes fused features of length `fused_dim` and produces:
/// * a new hidden state of length `hidden_dim`
/// * vocab logits of length `vocab_size`
#[derive(Debug, Clone)]
pub struct EagleDraftHead {
    /// W1 weight matrix, row-major `[hidden_dim × fused_dim]`.
    pub w1: Vec<f32>,
    /// Bias for layer 1, `[hidden_dim]`.
    pub b1: Vec<f32>,
    /// Next-hidden-state projection, row-major `[hidden_dim × hidden_dim]`.
    pub w_hidden: Vec<f32>,
    /// Bias for the hidden-state projection, `[hidden_dim]`.
    pub b_hidden: Vec<f32>,
    /// Vocab logits projection, row-major `[vocab_size × hidden_dim]`.
    pub w_vocab: Vec<f32>,
    /// Bias for the vocab projection, `[vocab_size]`.
    pub b_vocab: Vec<f32>,
    /// Fused input dimension (depends on `fusion_method`).
    pub fused_dim: usize,
    /// Hidden dimension `H`.
    pub hidden_dim: usize,
    /// Vocabulary size `V`.
    pub vocab_size: usize,
}

/// Output of a single [`EagleDraftHead::forward`] call.
#[derive(Debug, Clone)]
pub struct EagleHeadOutput {
    /// The next hidden state, length `hidden_dim`.
    pub hidden: Vec<f32>,
    /// Vocab logits, length `vocab_size`.
    pub logits: Vec<f32>,
}

impl EagleDraftHead {
    /// Initialise with LCG-seeded random weights.
    ///
    /// Weights are drawn from `U[-scale, scale]` where `scale = sqrt(2 /
    /// fan_in)` (Xavier-like) for each projection independently.  Biases are
    /// zeroed.
    pub fn new_random(fused_dim: usize, hidden_dim: usize, vocab_size: usize, seed: u64) -> Self {
        let mut rng = Lcg::new(seed);

        let scale_w1 = (2.0_f32 / fused_dim.max(1) as f32).sqrt();
        let w1: Vec<f32> = (0..hidden_dim * fused_dim)
            .map(|_| rng.next_scaled(scale_w1))
            .collect();
        let b1 = vec![0.0_f32; hidden_dim];

        let scale_wh = (2.0_f32 / hidden_dim.max(1) as f32).sqrt();
        let w_hidden: Vec<f32> = (0..hidden_dim * hidden_dim)
            .map(|_| rng.next_scaled(scale_wh))
            .collect();
        let b_hidden = vec![0.0_f32; hidden_dim];

        let scale_wv = (2.0_f32 / hidden_dim.max(1) as f32).sqrt();
        let w_vocab: Vec<f32> = (0..vocab_size * hidden_dim)
            .map(|_| rng.next_scaled(scale_wv))
            .collect();
        let b_vocab = vec![0.0_f32; vocab_size];

        Self {
            w1,
            b1,
            w_hidden,
            b_hidden,
            w_vocab,
            b_vocab,
            fused_dim,
            hidden_dim,
            vocab_size,
        }
    }

    /// Run the head forward pass over pre-fused input features.
    ///
    /// `fused` must have length `fused_dim`.  Returns the next hidden state
    /// (`hidden_dim`) and vocab logits (`vocab_size`).
    ///
    /// # Panics
    ///
    /// Panics in debug builds if `fused.len() != self.fused_dim`.
    pub fn forward(&self, fused: &[f32]) -> EagleHeadOutput {
        debug_assert_eq!(
            fused.len(),
            self.fused_dim,
            "fused feature length must equal fused_dim"
        );
        let h = self.hidden_dim;
        let f = self.fused_dim;

        // Layer 1: act = SiLU(W1 @ fused + b1)
        let mut act = vec![0.0_f32; h];
        for i in 0..h {
            let dot: f32 = (0..f).map(|j| self.w1[i * f + j] * fused[j]).sum::<f32>() + self.b1[i];
            act[i] = dot * (1.0 / (1.0 + (-dot).exp()));
        }

        // Next hidden state: hidden' = W_hidden @ act + b_hidden
        let mut hidden = vec![0.0_f32; h];
        for i in 0..h {
            hidden[i] = (0..h)
                .map(|j| self.w_hidden[i * h + j] * act[j])
                .sum::<f32>()
                + self.b_hidden[i];
        }

        // Vocab logits: logits = W_vocab @ act + b_vocab
        let v = self.vocab_size;
        let mut logits = vec![0.0_f32; v];
        for i in 0..v {
            logits[i] = (0..h)
                .map(|j| self.w_vocab[i * h + j] * act[j])
                .sum::<f32>()
                + self.b_vocab[i];
        }

        EagleHeadOutput { hidden, logits }
    }

    /// Return the indices of the top-`k` tokens sorted descending by logit.
    ///
    /// `k` is clamped to `vocab_size` if it exceeds it.
    pub fn top_k(&self, logits: &[f32], k: usize) -> Vec<(u32, f32)> {
        let k = k.min(logits.len());
        let mut idx: Vec<u32> = (0..logits.len() as u32).collect();
        idx.sort_unstable_by(|&a, &b| {
            logits[b as usize]
                .partial_cmp(&logits[a as usize])
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        idx[..k].iter().map(|&i| (i, logits[i as usize])).collect()
    }
}

// ── EagleHeadsConfig ───────────────────────────────────────────────────────────

/// Configuration for the [`EagleHeads`] draft model.
#[derive(Debug, Clone)]
pub struct EagleHeadsConfig {
    /// Hidden dimension shared by the draft head(s) and the base model.
    pub hidden_dim: usize,
    /// Vocabulary size.
    pub vocab_size: usize,
    /// Number of stacked draft-head layers (EAGLE-1 typically uses 1).
    pub num_draft_layers: usize,
    /// How to fuse the hidden state with the previous token's embedding.
    pub fusion_method: FusionMethod,
    /// Number of candidate tokens kept per autoregressive step.
    pub top_k: usize,
}

impl Default for EagleHeadsConfig {
    fn default() -> Self {
        Self {
            hidden_dim: 64,
            vocab_size: 256,
            num_draft_layers: 1,
            fusion_method: FusionMethod::Concat,
            top_k: 3,
        }
    }
}

impl From<&EagleConfig> for EagleHeadsConfig {
    fn from(cfg: &EagleConfig) -> Self {
        Self {
            hidden_dim: cfg.hidden_dim,
            vocab_size: 32_000, // orchestration EagleConfig has no vocab_size; use a safe default
            num_draft_layers: cfg.num_draft_layers.max(1),
            fusion_method: cfg.fusion_method,
            top_k: cfg.top_k,
        }
    }
}

// ── EagleDraftStep / EagleLossResult ──────────────────────────────────────────

/// One step of an autoregressive draft, as returned by [`EagleHeads::draft`].
#[derive(Debug, Clone)]
pub struct EagleDraftStep {
    /// Top token chosen for this step (argmax of `candidates`).
    pub token: u32,
    /// Probability-proxy (raw logit) of the chosen token.
    pub score: f32,
    /// Top-`top_k` candidates at this step, `(token_id, logit)`, sorted
    /// descending by logit — usable to build a candidate tree.
    pub candidates: Vec<(u32, f32)>,
    /// Hidden state produced at this step (fed forward to the next step).
    pub hidden: Vec<f32>,
}

/// Cross-entropy loss result returned by [`EagleHeads::training_loss`].
#[derive(Debug, Clone)]
pub struct EagleLossResult {
    /// Per-step cross-entropy loss, indexed by autoregressive step.
    pub per_step_loss: Vec<f32>,
    /// Mean cross-entropy across all steps.
    pub total_loss: f32,
}

// ── EagleHeads ─────────────────────────────────────────────────────────────────

/// The concrete EAGLE-1 draft model: a stack of [`EagleDraftHead`] layers plus
/// a token-embedding table, driving the autoregressive draft loop described
/// in arXiv:2401.15077.
///
/// Unlike Medusa's `K` independent heads (each predicting a fixed future
/// offset from the *same* base hidden state), EAGLE re-uses (a stack of) one
/// small head autoregressively: the hidden state produced at step `t` is fed
/// back in, fused with the embedding of the token drafted at step `t`, to
/// produce step `t + 1`.  This lets the draft model track feature-level
/// uncertainty from step to step, which the EAGLE paper shows improves
/// acceptance length over Medusa's independent-head approach.
pub struct EagleHeads {
    /// Shared configuration.
    pub config: EagleHeadsConfig,
    /// Stacked draft-head layers (only the first is used autoregressively in
    /// the paper's EAGLE-1 design; deeper stacks are supported for
    /// experimentation with `num_draft_layers > 1`).
    pub layers: Vec<EagleDraftHead>,
    /// Token embedding table, row-major `[vocab_size × hidden_dim]`.
    pub embedding: Vec<f32>,
}

impl EagleHeads {
    /// Create a new draft model with random weights.
    pub fn new(config: EagleHeadsConfig) -> Self {
        let fdim = fused_dim(config.hidden_dim, config.fusion_method);
        let layers: Vec<EagleDraftHead> = (0..config.num_draft_layers.max(1))
            .map(|i| {
                EagleDraftHead::new_random(
                    fdim,
                    config.hidden_dim,
                    config.vocab_size,
                    (i as u64)
                        .wrapping_mul(0x9e37_79b9_7f4a_7c15)
                        .wrapping_add(7),
                )
            })
            .collect();

        let mut emb_rng = Lcg::new(0xe5a1_e000);
        let emb_scale = (1.0_f32 / config.hidden_dim.max(1) as f32).sqrt();
        let embedding: Vec<f32> = (0..config.vocab_size * config.hidden_dim)
            .map(|_| emb_rng.next_scaled(emb_scale))
            .collect();

        Self {
            config,
            layers,
            embedding,
        }
    }

    /// Look up the embedding row for `token`, clamped to a valid vocab index
    /// (out-of-range tokens map to index 0, so drafting never panics on an
    /// unexpected context token id).
    fn embed(&self, token: u32) -> &[f32] {
        let h = self.config.hidden_dim;
        let idx = (token as usize).min(self.config.vocab_size.saturating_sub(1));
        &self.embedding[idx * h..idx * h + h]
    }

    /// The draft head used for autoregressive stepping (the first layer).
    fn head(&self) -> &EagleDraftHead {
        &self.layers[0]
    }

    /// Run the EAGLE autoregressive draft loop.
    ///
    /// `hidden` is the base model's last hidden state (`[hidden_dim]`).
    /// `context_tokens` supplies the most recently generated token (its last
    /// element seeds the first fusion step); if empty, token `0` is used.
    /// `steps` is the number of autoregressive draft tokens to generate.
    ///
    /// Returns one `(token, score)` pair per step — the argmax token and its
    /// raw logit — where `score` approximates a probability (higher is more
    /// confident); use [`EagleHeads::draft_steps`] for the full per-step
    /// detail (candidates, hidden states) needed to build a candidate tree.
    pub fn draft(&self, hidden: &[f32], context_tokens: &[u32], steps: usize) -> Vec<(u32, f32)> {
        self.draft_steps(hidden, context_tokens, steps)
            .into_iter()
            .map(|s| (s.token, s.score))
            .collect()
    }

    /// Like [`EagleHeads::draft`], but returns full per-step detail
    /// (top-`k` candidates and the resulting hidden state at each step) so a
    /// caller can build a candidate tree (cf. Medusa's `generate_tree` /
    /// EAGLE-3's dynamic tree).
    pub fn draft_steps(
        &self,
        hidden: &[f32],
        context_tokens: &[u32],
        steps: usize,
    ) -> Vec<EagleDraftStep> {
        let head = self.head();
        let mut cur_hidden = hidden.to_vec();
        let mut last_token = context_tokens.last().copied().unwrap_or(0);

        let mut out = Vec::with_capacity(steps);
        for _ in 0..steps {
            let embedding = self.embed(last_token);
            let fused = fuse_features(&cur_hidden, embedding, self.config.fusion_method);
            let result = head.forward(&fused);

            let candidates = head.top_k(&result.logits, self.config.top_k.max(1));
            let (top_token, top_score) = candidates
                .first()
                .copied()
                .unwrap_or((0, f32::NEG_INFINITY));

            cur_hidden = result.hidden.clone();
            last_token = top_token;

            out.push(EagleDraftStep {
                token: top_token,
                score: top_score,
                candidates,
                hidden: result.hidden,
            });
        }
        out
    }

    /// Compute per-step cross-entropy training loss.
    ///
    /// `hidden` is the base model's last hidden state (`[hidden_dim]`).
    /// `context_tokens` seeds the first fusion step, matching [`Self::draft`].
    /// `targets[t]` is the ground-truth token the model should draft at
    /// autoregressive step `t` (teacher-forced: the *target* token, not the
    /// model's own prediction, is fed forward into the next step's fusion —
    /// this is the standard teacher-forcing training regime for EAGLE).
    ///
    /// Loss is computed with the numerically stable log-sum-exp trick.
    pub fn training_loss(
        &self,
        hidden: &[f32],
        context_tokens: &[u32],
        targets: &[u32],
    ) -> EagleLossResult {
        let head = self.head();
        let mut cur_hidden = hidden.to_vec();
        let mut last_token = context_tokens.last().copied().unwrap_or(0);

        let mut per_step_loss = Vec::with_capacity(targets.len());
        for &target in targets {
            let embedding = self.embed(last_token);
            let fused = fuse_features(&cur_hidden, embedding, self.config.fusion_method);
            let result = head.forward(&fused);

            let logits = &result.logits;
            let max_l = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let log_sum = max_l + logits.iter().map(|l| (l - max_l).exp()).sum::<f32>().ln();
            let ce = -(logits[target as usize] - log_sum);
            per_step_loss.push(ce);

            // Teacher forcing: advance state using the *target* token.
            cur_hidden = result.hidden;
            last_token = target;
        }

        let total_loss = if per_step_loss.is_empty() {
            0.0
        } else {
            per_step_loss.iter().sum::<f32>() / per_step_loss.len() as f32
        };

        EagleLossResult {
            per_step_loss,
            total_loss,
        }
    }
}

// ── EagleDraftModel trait impl ─────────────────────────────────────────────────

#[async_trait::async_trait]
impl EagleDraftModel for EagleHeads {
    async fn draft_from_hidden(
        &self,
        hidden_states: &[f32],
        context_tokens: &[u32],
        config: &EagleConfig,
    ) -> InferenceResult<Vec<Token>> {
        let steps = self.draft_steps(hidden_states, context_tokens, config.draft_steps.max(1));
        Ok(steps
            .into_iter()
            .map(|step| {
                // Softmax over the top-k candidates for a normalised probability.
                let max_l = step
                    .candidates
                    .iter()
                    .map(|&(_, l)| l)
                    .fold(f32::NEG_INFINITY, f32::max);
                let sum_exp: f32 = step
                    .candidates
                    .iter()
                    .map(|&(_, l)| (l - max_l).exp())
                    .sum();
                let prob = if sum_exp > 0.0 {
                    ((step.score - max_l).exp()) / sum_exp
                } else {
                    0.0
                };
                Token {
                    id: step.token,
                    text: format!("tok_{}", step.token),
                    logits: step.candidates.iter().map(|&(_, l)| l).collect(),
                    probability: prob,
                }
            })
            .collect())
    }

    fn hidden_dim(&self) -> usize {
        self.config.hidden_dim
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    // ── helpers ──

    fn default_model() -> EagleHeads {
        EagleHeads::new(EagleHeadsConfig::default())
    }

    fn zero_hidden(dim: usize) -> Vec<f32> {
        vec![0.0_f32; dim]
    }

    fn ones_hidden(dim: usize) -> Vec<f32> {
        vec![1.0_f32; dim]
    }

    // ── config ───────────────────────────────────────────────────────────────

    #[test]
    fn test_default_config() {
        let cfg = EagleHeadsConfig::default();
        assert_eq!(cfg.hidden_dim, 64);
        assert_eq!(cfg.vocab_size, 256);
        assert_eq!(cfg.num_draft_layers, 1);
        assert_eq!(cfg.fusion_method, FusionMethod::Concat);
        assert_eq!(cfg.top_k, 3);
    }

    // ── fuse_features ────────────────────────────────────────────────────────

    #[test]
    fn test_fuse_concat_length() {
        let h = vec![1.0_f32, 2.0, 3.0];
        let e = vec![4.0_f32, 5.0, 6.0];
        let fused = fuse_features(&h, &e, FusionMethod::Concat);
        assert_eq!(fused.len(), 6);
        assert_eq!(fused, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_fuse_add_length_and_values() {
        let h = vec![1.0_f32, 2.0, 3.0];
        let e = vec![4.0_f32, 5.0, 6.0];
        let fused = fuse_features(&h, &e, FusionMethod::Add);
        assert_eq!(fused.len(), 3);
        assert_eq!(fused, vec![5.0, 7.0, 9.0]);
    }

    #[test]
    fn test_fuse_attention_length_and_finite() {
        let h = vec![1.0_f32, 2.0, 3.0];
        let e = vec![0.5_f32, -1.0, 2.0];
        let fused = fuse_features(&h, &e, FusionMethod::Attention);
        assert_eq!(fused.len(), 3);
        assert!(fused.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn test_fused_dim_matches_method() {
        assert_eq!(fused_dim(64, FusionMethod::Concat), 128);
        assert_eq!(fused_dim(64, FusionMethod::Add), 64);
        assert_eq!(fused_dim(64, FusionMethod::Attention), 64);
    }

    // ── EagleDraftHead::forward ──────────────────────────────────────────────

    #[test]
    fn test_head_forward_shape() {
        let cfg = EagleHeadsConfig::default();
        let fdim = fused_dim(cfg.hidden_dim, cfg.fusion_method);
        let head = EagleDraftHead::new_random(fdim, cfg.hidden_dim, cfg.vocab_size, 1);
        let fused = zero_hidden(fdim);
        let out = head.forward(&fused);
        assert_eq!(out.hidden.len(), cfg.hidden_dim);
        assert_eq!(out.logits.len(), cfg.vocab_size);
    }

    #[test]
    fn test_head_forward_finite() {
        let cfg = EagleHeadsConfig::default();
        let fdim = fused_dim(cfg.hidden_dim, cfg.fusion_method);
        let head = EagleDraftHead::new_random(fdim, cfg.hidden_dim, cfg.vocab_size, 2);
        let fused = ones_hidden(fdim);
        let out = head.forward(&fused);
        assert!(out.hidden.iter().all(|v| v.is_finite()));
        assert!(out.logits.iter().all(|v| v.is_finite()));
    }

    // ── EagleDraftHead::top_k ────────────────────────────────────────────────

    #[test]
    fn test_head_top_k_length_and_sorted() {
        let cfg = EagleHeadsConfig::default();
        let fdim = fused_dim(cfg.hidden_dim, cfg.fusion_method);
        let head = EagleDraftHead::new_random(fdim, cfg.hidden_dim, cfg.vocab_size, 3);
        let logits = head.forward(&ones_hidden(fdim)).logits;
        let result = head.top_k(&logits, 5);
        assert_eq!(result.len(), 5);
        for w in result.windows(2) {
            assert!(w[0].1 >= w[1].1, "top-k must be sorted descending by logit");
        }
    }

    #[test]
    fn test_head_top_k_clamped_to_vocab() {
        let cfg = EagleHeadsConfig::default();
        let fdim = fused_dim(cfg.hidden_dim, cfg.fusion_method);
        let head = EagleDraftHead::new_random(fdim, cfg.hidden_dim, cfg.vocab_size, 4);
        let logits = head.forward(&zero_hidden(fdim)).logits;
        let result = head.top_k(&logits, cfg.vocab_size + 50);
        assert_eq!(result.len(), cfg.vocab_size);
    }

    // ── EagleHeads::new ──────────────────────────────────────────────────────

    #[test]
    fn test_new_builds_correct_layer_count() {
        let cfg = EagleHeadsConfig {
            num_draft_layers: 3,
            ..EagleHeadsConfig::default()
        };
        let model = EagleHeads::new(cfg.clone());
        assert_eq!(model.layers.len(), 3);
    }

    #[test]
    fn test_new_embedding_table_shape() {
        let cfg = EagleHeadsConfig::default();
        let model = EagleHeads::new(cfg.clone());
        assert_eq!(model.embedding.len(), cfg.vocab_size * cfg.hidden_dim);
    }

    #[test]
    fn test_seeded_init_is_deterministic() {
        let cfg = EagleHeadsConfig::default();
        let a = EagleHeads::new(cfg.clone());
        let b = EagleHeads::new(cfg);
        assert_eq!(a.layers[0].w1, b.layers[0].w1);
        assert_eq!(a.embedding, b.embedding);
    }

    // ── EagleHeads::draft / draft_steps ──────────────────────────────────────

    #[test]
    fn test_draft_length() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let draft = model.draft(&hidden, &[1, 2, 3], 4);
        assert_eq!(draft.len(), 4);
    }

    #[test]
    fn test_draft_zero_steps() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let draft = model.draft(&hidden, &[1, 2, 3], 0);
        assert!(draft.is_empty());
    }

    #[test]
    fn test_draft_empty_context_does_not_panic() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let draft = model.draft(&hidden, &[], 2);
        assert_eq!(draft.len(), 2);
    }

    #[test]
    fn test_draft_out_of_range_context_token_does_not_panic() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        // Token id far beyond vocab_size must be clamped, not panic.
        let draft = model.draft(&hidden, &[u32::MAX], 2);
        assert_eq!(draft.len(), 2);
    }

    #[test]
    fn test_draft_steps_candidates_and_hidden_shapes() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let steps = model.draft_steps(&hidden, &[5], 3);
        assert_eq!(steps.len(), 3);
        for step in &steps {
            assert_eq!(step.hidden.len(), model.config.hidden_dim);
            assert_eq!(step.candidates.len(), model.config.top_k);
            assert!(step.score.is_finite());
        }
    }

    #[test]
    fn test_draft_steps_top_candidate_matches_token() {
        let model = default_model();
        let hidden = ones_hidden(model.config.hidden_dim);
        let steps = model.draft_steps(&hidden, &[0], 2);
        for step in &steps {
            assert_eq!(step.candidates[0].0, step.token);
            assert!((step.candidates[0].1 - step.score).abs() < 1e-6);
        }
    }

    #[test]
    fn test_draft_all_fusion_methods() {
        for method in [
            FusionMethod::Concat,
            FusionMethod::Add,
            FusionMethod::Attention,
        ] {
            let cfg = EagleHeadsConfig {
                fusion_method: method,
                ..EagleHeadsConfig::default()
            };
            let model = EagleHeads::new(cfg.clone());
            let hidden = zero_hidden(cfg.hidden_dim);
            let draft = model.draft(&hidden, &[1, 2], 3);
            assert_eq!(
                draft.len(),
                3,
                "fusion method {method:?} produced wrong draft length"
            );
            assert!(draft.iter().all(|&(_, s)| s.is_finite()));
        }
    }

    // ── EagleHeads::training_loss ────────────────────────────────────────────

    #[test]
    fn test_training_loss_shape() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let targets = vec![0_u32, 1, 2];
        let result = model.training_loss(&hidden, &[1], &targets);
        assert_eq!(result.per_step_loss.len(), targets.len());
    }

    #[test]
    fn test_training_loss_non_negative_and_finite() {
        let model = default_model();
        let hidden = ones_hidden(model.config.hidden_dim);
        let targets = vec![0_u32; 4];
        let result = model.training_loss(&hidden, &[2], &targets);
        for (i, &loss) in result.per_step_loss.iter().enumerate() {
            assert!(
                loss >= 0.0,
                "step {i}: CE loss must be non-negative, got {loss}"
            );
            assert!(loss.is_finite());
        }
        assert!(result.total_loss.is_finite());
    }

    #[test]
    fn test_training_loss_total_is_mean() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let targets = vec![1_u32; 5];
        let result = model.training_loss(&hidden, &[0], &targets);
        let expected_mean =
            result.per_step_loss.iter().sum::<f32>() / result.per_step_loss.len() as f32;
        assert!((result.total_loss - expected_mean).abs() < 1e-5);
    }

    #[test]
    fn test_training_loss_empty_targets() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let result = model.training_loss(&hidden, &[0], &[]);
        assert!(result.per_step_loss.is_empty());
        assert_eq!(result.total_loss, 0.0);
    }

    #[test]
    fn test_training_loss_near_zero_for_dominant_logit() {
        // Build a single-layer head whose vocab row for token 0 dominates,
        // so CE loss for target=0 should be close to 0.
        let hidden_dim = 4;
        let vocab_size = 8;
        let fdim = fused_dim(hidden_dim, FusionMethod::Add);
        let target: u32 = 0;

        let mut w_vocab = vec![0.0_f32; vocab_size * hidden_dim];
        for j in 0..hidden_dim {
            w_vocab[target as usize * hidden_dim + j] = 100.0;
        }
        let head = EagleDraftHead {
            w1: {
                // Identity-like mapping from fused (Add => hidden_dim) to hidden_dim.
                let mut m = vec![0.0_f32; hidden_dim * fdim];
                for i in 0..hidden_dim {
                    m[i * fdim + i] = 1.0;
                }
                m
            },
            b1: vec![0.0; hidden_dim],
            w_hidden: vec![0.0; hidden_dim * hidden_dim],
            b_hidden: vec![0.0; hidden_dim],
            w_vocab,
            b_vocab: vec![0.0; vocab_size],
            fused_dim: fdim,
            hidden_dim,
            vocab_size,
        };

        let fused = vec![1.0_f32; fdim];
        let out = head.forward(&fused);
        let logits = out.logits;
        let max_l = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let log_sum = max_l + logits.iter().map(|l| (l - max_l).exp()).sum::<f32>().ln();
        let ce = -(logits[target as usize] - log_sum);
        assert!(
            ce < 0.01,
            "CE loss should be near 0 when target logit dominates, got {ce}"
        );
    }

    // ── EagleDraftModel trait impl ───────────────────────────────────────────

    #[tokio::test]
    async fn test_trait_draft_from_hidden_smoke() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let eagle_cfg = EagleConfig {
            draft_steps: 3,
            hidden_dim: model.config.hidden_dim,
            num_draft_layers: 1,
            fusion_method: model.config.fusion_method,
            top_k: model.config.top_k,
        };

        let tokens = (&model as &dyn EagleDraftModel)
            .draft_from_hidden(&hidden, &[1, 2, 3], &eagle_cfg)
            .await
            .expect("draft_from_hidden must succeed");

        assert_eq!(tokens.len(), 3);
        for t in &tokens {
            assert!(t.probability.is_finite());
            assert!(!t.text.is_empty());
        }
    }

    #[tokio::test]
    async fn test_trait_hidden_dim_accessor() {
        let model = default_model();
        let dim = (&model as &dyn EagleDraftModel).hidden_dim();
        assert_eq!(dim, model.config.hidden_dim);
    }

    #[tokio::test]
    async fn test_trait_zero_draft_steps_yields_empty() {
        let model = default_model();
        let hidden = zero_hidden(model.config.hidden_dim);
        let eagle_cfg = EagleConfig {
            draft_steps: 0,
            ..EagleConfig::default()
        };
        // `.max(1)` inside draft_from_hidden ensures at least 1 step even if
        // config requests 0, matching the crate-wide convention of never
        // silently no-op'ing a draft call.
        let tokens = (&model as &dyn EagleDraftModel)
            .draft_from_hidden(&hidden, &[0], &eagle_cfg)
            .await
            .expect("draft_from_hidden must succeed");
        assert_eq!(tokens.len(), 1);
    }
}