//! Medusa speculative decoding heads (Cai et al. 2024, arXiv:2401.10774).
//!
//! Medusa accelerates autoregressive decoding by attaching K small FFN "draft heads"
//! to the last hidden state of the base model.  Each head independently predicts
//! tokens at offsets +1, +2, …, +K from the current position.  At inference time the
//! heads collectively generate a *tree* of candidate continuations, which the base
//! model verifies in a single forward pass using tree attention.
//!
//! # Architecture
//!
//! Each [`MedusaHead`] is a two-layer FFN:
//!
//! ```text
//! hidden [H] → SiLU(W1 @ hidden + b1) [H] → W2 @ act + b2 [V]
//! ```
//!
//! The [`MedusaHeads`] ensemble holds K such heads and exposes:
//!
//! * [`MedusaHeads::generate_tree`] — cartesian-product candidate tree
//! * [`MedusaHeads::verify`]        — longest-accepted-prefix selection
//! * [`MedusaHeads::training_loss`] — per-head cross-entropy
//!
//! # Example
//!
//! ```rust
//! use rtx_inference::medusa::{MedusaHeads, MedusaConfig};
//!
//! let config = MedusaConfig::default();
//! let heads  = MedusaHeads::new(config);
//! let hidden = vec![0.0_f32; 64];
//! let tree   = heads.generate_tree(&hidden);
//! assert_eq!(tree.num_paths(), 3_usize.pow(4)); // top_k^num_heads
//! ```

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

/// Minimal LCG PRNG used for reproducible weight initialisation (no external deps).
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 {
        // Knuth's multiplier + addend, 64-bit
        self.0 = self
            .0
            .wrapping_mul(6_364_136_223_846_793_005)
            .wrapping_add(1_442_695_040_888_963_407);
        // Use the high 23 bits for the mantissa
        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
    }
}

// ── MedusaHead ────────────────────────────────────────────────────────────────

/// One Medusa draft head: a 2-layer FFN applied to the last hidden state.
///
/// Computes `hidden [H] → SiLU(W1 @ hidden + b1) [H] → W2 @ act + b2 [V]`.
#[derive(Debug, Clone)]
pub struct MedusaHead {
    /// W1 weight matrix, row-major `[hidden_dim × hidden_dim]`.
    pub w1: Vec<f32>,
    /// Bias for layer 1, `[hidden_dim]`.
    pub b1: Vec<f32>,
    /// W2 weight matrix, row-major `[vocab_size × hidden_dim]`.
    pub w2: Vec<f32>,
    /// Bias for layer 2, `[vocab_size]`.
    pub b2: Vec<f32>,
    /// Hidden dimension `H`.
    pub hidden_dim: usize,
    /// Vocabulary size `V`.
    pub vocab_size: usize,
}

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

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

        let scale_w2 = (2.0_f32 / hidden_dim as f32).sqrt();
        let w2: Vec<f32> = (0..vocab_size * hidden_dim)
            .map(|_| rng.next_scaled(scale_w2))
            .collect();
        let b2 = vec![0.0_f32; vocab_size];

        Self {
            w1,
            b1,
            w2,
            b2,
            hidden_dim,
            vocab_size,
        }
    }

    /// Run the head forward pass.
    ///
    /// `hidden` must have length `hidden_dim`.  Returns logits of length
    /// `vocab_size`.
    ///
    /// # Panics
    ///
    /// Panics in debug builds if `hidden.len() != self.hidden_dim`.
    pub fn forward(&self, hidden: &[f32]) -> Vec<f32> {
        debug_assert_eq!(
            hidden.len(),
            self.hidden_dim,
            "hidden vector length must equal hidden_dim"
        );
        let h = self.hidden_dim;

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

        // Layer 2: logits = W2 @ act + b2
        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.w2[i * h + j] * act[j]).sum::<f32>() + self.b2[i];
        }
        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, hidden: &[f32], k: usize) -> Vec<u32> {
        let logits = self.forward(hidden);
        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].to_vec()
    }
}

// ── MedusaConfig ─────────────────────────────────────────────────────────────

/// Configuration for the Medusa head ensemble.
#[derive(Debug, Clone)]
pub struct MedusaConfig {
    /// `K`: number of draft heads (each predicts a different future offset).
    pub num_heads: usize,
    /// Hidden dimension shared by all heads and the base model.
    pub hidden_dim: usize,
    /// Vocabulary size.
    pub vocab_size: usize,
    /// Number of candidate tokens kept per head per position.
    pub top_k_per_head: usize,
}

impl Default for MedusaConfig {
    fn default() -> Self {
        Self {
            num_heads: 4,
            hidden_dim: 64,
            vocab_size: 256,
            top_k_per_head: 3,
        }
    }
}

// ── MedusaTree ────────────────────────────────────────────────────────────────

/// A tree of draft candidate paths generated by [`MedusaHeads::generate_tree`].
///
/// Each path is a `Vec<u32>` of length `num_heads`, representing predicted
/// tokens at offsets +1, +2, …, +K from the current position.
#[derive(Debug, Clone)]
pub struct MedusaTree {
    /// All candidate paths.  Length == `top_k_per_head ^ num_heads`.
    pub paths: Vec<Vec<u32>>,
    /// Raw logit vectors produced by each head, indexed `[head][vocab]`.
    pub head_logits: Vec<Vec<f32>>,
}

impl MedusaTree {
    /// Number of candidate paths in the tree.
    pub fn num_paths(&self) -> usize {
        self.paths.len()
    }

    /// Length of each path (equals `num_heads`).
    pub fn num_heads(&self) -> usize {
        self.paths.first().map_or(0, Vec::len)
    }
}

// ── MedusaVerifyResult ────────────────────────────────────────────────────────

/// Result returned by [`MedusaHeads::verify`].
#[derive(Debug, Clone)]
pub struct MedusaVerifyResult {
    /// Longest accepted token sequence (always includes at least `base_next_token`).
    pub accepted_tokens: Vec<u32>,
    /// Index into [`MedusaTree::paths`] of the accepted path, or `None` if the
    /// draft was entirely rejected and only the base token was kept.
    pub accepted_path_idx: Option<usize>,
    /// Number of *speculative* tokens accepted beyond the base model token
    /// (i.e. `accepted_tokens.len() - 1`).
    pub num_accepted: usize,
}

// ── MedusaLossResult ─────────────────────────────────────────────────────────

/// Cross-entropy losses returned by [`MedusaHeads::training_loss`].
#[derive(Debug, Clone)]
pub struct MedusaLossResult {
    /// Per-head cross-entropy loss, indexed `[head]`.
    pub per_head_loss: Vec<f32>,
    /// Mean cross-entropy across all heads.
    pub total_loss: f32,
}

// ── MedusaHeads ──────────────────────────────────────────────────────────────

/// The complete Medusa head ensemble.
///
/// Holds K [`MedusaHead`]s and orchestrates tree generation, candidate
/// verification, and training-loss computation.
pub struct MedusaHeads {
    /// Shared configuration.
    pub config: MedusaConfig,
    /// Individual draft heads.
    pub heads: Vec<MedusaHead>,
}

impl MedusaHeads {
    /// Create a new ensemble with random weights, seeding each head differently.
    pub fn new(config: MedusaConfig) -> Self {
        let heads: Vec<MedusaHead> = (0..config.num_heads)
            .map(|i| {
                MedusaHead::new_random(
                    config.hidden_dim,
                    config.vocab_size,
                    // Distinct seed per head so weights differ.
                    (i as u64)
                        .wrapping_mul(0xdead_beef_cafe_babe)
                        .wrapping_add(42),
                )
            })
            .collect();
        Self { config, heads }
    }

    /// Generate a draft candidate tree from the base model's last hidden state.
    ///
    /// Each head produces a top-`top_k_per_head` list; the tree is the
    /// cartesian product of these lists, giving `top_k_per_head ^ num_heads`
    /// paths, each of length `num_heads`.
    pub fn generate_tree(&self, hidden: &[f32]) -> MedusaTree {
        let k = self.config.top_k_per_head;

        // Collect logits and top-k token indices from every head.
        let head_logits: Vec<Vec<f32>> = self.heads.iter().map(|h| h.forward(hidden)).collect();

        let head_topk: Vec<Vec<u32>> = self
            .heads
            .iter()
            .zip(head_logits.iter())
            .map(|(_, logits)| {
                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.min(idx.len())].to_vec()
            })
            .collect();

        // Cartesian product: start with a single empty path and extend.
        let mut paths: Vec<Vec<u32>> = vec![vec![]];
        for topk in &head_topk {
            paths = paths
                .iter()
                .flat_map(|path| {
                    topk.iter().map(move |&tok| {
                        let mut p = path.clone();
                        p.push(tok);
                        p
                    })
                })
                .collect();
        }

        MedusaTree { paths, head_logits }
    }

    /// Verify draft candidates against base model predictions.
    ///
    /// For each candidate path the verifier checks, token by token, whether the
    /// draft agrees with `path_oracle` (a closure that returns the base model's
    /// greedy token given a path prefix).  The first token in every path must
    /// match `base_next_token`; subsequent tokens are checked via `path_oracle`.
    ///
    /// When a mismatch is found the oracle's correction is appended and the
    /// comparison stops.  The path with the longest accepted sequence is returned.
    ///
    /// The result always contains at least `base_next_token`.
    pub fn verify(
        &self,
        tree: &MedusaTree,
        base_next_token: u32,
        path_oracle: &dyn Fn(&[u32]) -> u32,
    ) -> MedusaVerifyResult {
        let mut best: Vec<u32> = vec![];
        let mut best_path_idx: Option<usize> = None;

        for (idx, path) in tree.paths.iter().enumerate() {
            // The first element of every candidate path must be the base token.
            if path.is_empty() || path[0] != base_next_token {
                continue;
            }

            let mut accepted = vec![base_next_token];

            // Verify tokens path[1], path[2], …  against the oracle.
            for (pos, &draft_tok) in path[1..].iter().enumerate() {
                // Oracle: given path[0..=pos], what would the base model predict?
                let expected = path_oracle(&path[..pos + 1]);
                if draft_tok == expected {
                    accepted.push(draft_tok);
                } else {
                    // Accept the oracle's correction and stop.
                    accepted.push(expected);
                    break;
                }
            }

            if accepted.len() > best.len() {
                best = accepted;
                best_path_idx = Some(idx);
            }
        }

        // Fallback: if no path started with base_next_token, accept it alone.
        if best.is_empty() {
            best = vec![base_next_token];
        }

        let num_accepted = best.len().saturating_sub(1);
        MedusaVerifyResult {
            accepted_tokens: best,
            accepted_path_idx: best_path_idx,
            num_accepted,
        }
    }

    /// Compute per-head cross-entropy training loss.
    ///
    /// `hidden` is the base model's last hidden state `[hidden_dim]`.
    /// `targets[i]` is the ground-truth token that head `i` should predict
    /// (i.e. the token at offset `+i+1` in the training sequence).
    ///
    /// Loss is computed with the numerically stable log-sum-exp trick.
    ///
    /// # Panics
    ///
    /// Panics if `targets.len() != self.heads.len()`.
    pub fn training_loss(&self, hidden: &[f32], targets: &[u32]) -> MedusaLossResult {
        assert_eq!(
            targets.len(),
            self.heads.len(),
            "targets length must equal number of heads"
        );

        let per_head_loss: Vec<f32> = self
            .heads
            .iter()
            .zip(targets.iter())
            .map(|(head, &target)| {
                let logits = head.forward(hidden);
                // Numerically stable CE: loss = log_sum_exp(logits) - logits[target]
                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();
                -(logits[target as usize] - log_sum)
            })
            .collect();

        let total_loss = per_head_loss.iter().sum::<f32>() / per_head_loss.len() as f32;

        MedusaLossResult {
            per_head_loss,
            total_loss,
        }
    }
}

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

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

    // ── helpers ──

    fn default_heads() -> MedusaHeads {
        MedusaHeads::new(MedusaConfig::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 = MedusaConfig::default();
        assert_eq!(cfg.num_heads, 4);
        assert_eq!(cfg.hidden_dim, 64);
        assert_eq!(cfg.vocab_size, 256);
        assert_eq!(cfg.top_k_per_head, 3);
    }

    // ── MedusaHead::forward ───────────────────────────────────────────────────

    #[test]
    fn test_head_forward_shape() {
        let cfg = MedusaConfig::default();
        let head = MedusaHead::new_random(cfg.hidden_dim, cfg.vocab_size, 1);
        let hidden = zero_hidden(cfg.hidden_dim);
        let logits = head.forward(&hidden);
        assert_eq!(logits.len(), cfg.vocab_size);
    }

    #[test]
    fn test_head_forward_finite() {
        let cfg = MedusaConfig::default();
        let head = MedusaHead::new_random(cfg.hidden_dim, cfg.vocab_size, 2);
        let hidden = ones_hidden(cfg.hidden_dim);
        let logits = head.forward(&hidden);
        assert!(
            logits.iter().all(|v| v.is_finite()),
            "all logits must be finite"
        );
    }

    // ── MedusaHead::top_k ────────────────────────────────────────────────────

    #[test]
    fn test_head_top_k_length() {
        let cfg = MedusaConfig::default();
        let head = MedusaHead::new_random(cfg.hidden_dim, cfg.vocab_size, 3);
        let hidden = zero_hidden(cfg.hidden_dim);
        let k = 5;
        let result = head.top_k(&hidden, k);
        assert_eq!(result.len(), k);
    }

    #[test]
    fn test_head_top_k_sorted() {
        let cfg = MedusaConfig::default();
        let head = MedusaHead::new_random(cfg.hidden_dim, cfg.vocab_size, 4);
        let hidden = ones_hidden(cfg.hidden_dim);
        let result = head.top_k(&hidden, 8);
        let logits = head.forward(&hidden);
        // Each successive index should have a logit <= the previous.
        for w in result.windows(2) {
            assert!(
                logits[w[0] as usize] >= logits[w[1] as usize],
                "top-k must be sorted descending by logit"
            );
        }
    }

    #[test]
    fn test_head_top_k_unique() {
        let cfg = MedusaConfig::default();
        let head = MedusaHead::new_random(cfg.hidden_dim, cfg.vocab_size, 5);
        let hidden = ones_hidden(cfg.hidden_dim);
        let k = 10;
        let result = head.top_k(&hidden, k);
        let mut sorted = result.clone();
        sorted.sort_unstable();
        sorted.dedup();
        assert_eq!(sorted.len(), result.len(), "top-k indices must be unique");
    }

    #[test]
    fn test_head_top_k_clamped_to_vocab() {
        let cfg = MedusaConfig::default();
        let head = MedusaHead::new_random(cfg.hidden_dim, cfg.vocab_size, 6);
        let hidden = zero_hidden(cfg.hidden_dim);
        // Request more than vocab_size — should not panic, returns vocab_size items.
        let result = head.top_k(&hidden, cfg.vocab_size + 100);
        assert_eq!(result.len(), cfg.vocab_size);
    }

    // ── SiLU activation ──────────────────────────────────────────────────────

    #[test]
    fn test_head_silu_activation_at_zero() {
        // When the hidden state is zero AND biases are zero, W1@0 + b1 = 0,
        // so the pre-activation is 0 for every neuron.  SiLU(0) = 0 * 0.5 = 0.
        // Build a head with zero weights to isolate this property.
        let h = 4;
        let v = 8;
        let head = MedusaHead {
            w1: vec![0.0; h * h],
            b1: vec![0.0; h],
            w2: vec![1.0; v * h], // uniform W2 so logits == sum(act) = 0
            b2: vec![0.0; v],
            hidden_dim: h,
            vocab_size: v,
        };
        let logits = head.forward(&vec![0.0; h]);
        // All activations are SiLU(0) = 0, so all logits should be 0.
        for l in &logits {
            assert!(l.abs() < 1e-6, "SiLU(0) path: expected logit ≈ 0, got {l}");
        }
    }

    #[test]
    fn test_head_silu_activation_positive_input() {
        // With positive pre-activation x > 0, SiLU(x) = x * sigmoid(x) > 0.
        let h = 2;
        let v = 2;
        // Identity W1 with positive bias forces positive pre-activation.
        let w1 = vec![1.0_f32, 0.0, 0.0, 1.0]; // 2×2 identity
        let b1 = vec![1.0_f32, 1.0]; // shift pre-activation up
        let w2 = vec![1.0_f32; v * h];
        let b2 = vec![0.0_f32; v];
        let head = MedusaHead {
            w1,
            b1,
            w2,
            b2,
            hidden_dim: h,
            vocab_size: v,
        };
        let logits = head.forward(&vec![0.0; h]);
        // Pre-activation = 1.0; SiLU(1.0) = 1 / (1 + e^{-1}) ≈ 0.731 > 0.
        for l in &logits {
            assert!(
                *l > 0.0,
                "SiLU of positive input should produce positive output, got {l}"
            );
        }
    }

    // ── MedusaTree ────────────────────────────────────────────────────────────

    #[test]
    fn test_generate_tree_num_paths() {
        let cfg = MedusaConfig::default(); // top_k=3, num_heads=4
        let heads = MedusaHeads::new(cfg.clone());
        let hidden = zero_hidden(cfg.hidden_dim);
        let tree = heads.generate_tree(&hidden);
        let expected = cfg.top_k_per_head.pow(cfg.num_heads as u32);
        assert_eq!(tree.num_paths(), expected);
    }

    #[test]
    fn test_generate_tree_path_length() {
        let cfg = MedusaConfig::default();
        let heads = MedusaHeads::new(cfg.clone());
        let hidden = zero_hidden(cfg.hidden_dim);
        let tree = heads.generate_tree(&hidden);
        for path in &tree.paths {
            assert_eq!(path.len(), cfg.num_heads, "each path must span all heads");
        }
    }

    #[test]
    fn test_generate_tree_head_logits_count() {
        let cfg = MedusaConfig::default();
        let heads = MedusaHeads::new(cfg.clone());
        let hidden = zero_hidden(cfg.hidden_dim);
        let tree = heads.generate_tree(&hidden);
        assert_eq!(tree.head_logits.len(), cfg.num_heads);
    }

    #[test]
    fn test_tree_num_paths_accessor() {
        let tree = MedusaTree {
            paths: vec![vec![1, 2], vec![3, 4], vec![5, 6]],
            head_logits: vec![],
        };
        assert_eq!(tree.num_paths(), 3);
        assert_eq!(tree.num_heads(), 2);
    }

    // ── verify ────────────────────────────────────────────────────────────────

    /// Build a simple 1-head tree (top_k=3) and control path_oracle precisely.
    fn small_heads() -> (MedusaHeads, MedusaConfig) {
        let cfg = MedusaConfig {
            num_heads: 1,
            hidden_dim: 4,
            vocab_size: 8,
            top_k_per_head: 3,
        };
        (MedusaHeads::new(cfg.clone()), cfg)
    }

    #[test]
    fn test_verify_accepts_base_token() {
        let (heads, cfg) = small_heads();
        let hidden = zero_hidden(cfg.hidden_dim);
        let tree = heads.generate_tree(&hidden);
        let base_token = 99_u32; // not in any path (paths ⊂ [0,7])
        // oracle never called because no path starts with 99
        let result = heads.verify(&tree, base_token, &|_| unreachable!());
        assert_eq!(result.accepted_tokens, vec![base_token]);
        assert_eq!(result.num_accepted, 0);
    }

    #[test]
    fn test_verify_perfect_draft() {
        // 1-head tree: a path [tok] matches base_next_token == tok.
        // Since num_heads==1, after accepting tok the loop over path[1..] is empty,
        // so the whole length-1 path is accepted.
        let (heads, cfg) = small_heads();
        let hidden = zero_hidden(cfg.hidden_dim);
        let tree = heads.generate_tree(&hidden);

        // Find the first path's leading token and use it as base_next_token.
        let first_tok = tree.paths[0][0];
        let result = heads.verify(&tree, first_tok, &|_| first_tok);
        // accepted_tokens contains exactly [first_tok]; num_accepted == 0 (single head)
        assert!(result.accepted_tokens.contains(&first_tok));
        assert!(result.accepted_path_idx.is_some());
    }

    #[test]
    fn test_verify_perfect_draft_multi_head() {
        // 2-head tree, oracle always agrees with draft → full path accepted.
        let cfg = MedusaConfig {
            num_heads: 2,
            hidden_dim: 4,
            vocab_size: 8,
            top_k_per_head: 2,
        };
        let heads = MedusaHeads::new(cfg.clone());
        let hidden = zero_hidden(cfg.hidden_dim);
        let tree = heads.generate_tree(&hidden);

        // Pick a path and set oracle to always confirm draft.
        let path = tree.paths[0].clone();
        let base = path[0];
        let draft_second = path[1];
        let result = heads.verify(&tree, base, &move |_prefix| draft_second);

        // Should accept both tokens.
        assert_eq!(result.accepted_tokens.len(), 2);
        assert_eq!(result.num_accepted, 1);
    }

    #[test]
    fn test_verify_mismatch_accepts_correction() {
        // 2-head tree: oracle disagrees at position 1 → correction is appended.
        let cfg = MedusaConfig {
            num_heads: 2,
            hidden_dim: 4,
            vocab_size: 8,
            top_k_per_head: 2,
        };
        let heads = MedusaHeads::new(cfg.clone());
        let hidden = zero_hidden(cfg.hidden_dim);
        let tree = heads.generate_tree(&hidden);

        let path = tree.paths[0].clone();
        let base = path[0];
        let correction = 200_u32; // guaranteed to differ from any vocab token (0..7)
        let result = heads.verify(&tree, base, &move |_prefix| correction);

        // base accepted, then correction appended → length 2, num_accepted = 1
        assert_eq!(result.accepted_tokens.len(), 2);
        assert_eq!(*result.accepted_tokens.last().unwrap(), correction);
        assert_eq!(result.num_accepted, 1);
    }

    #[test]
    fn test_verify_num_accepted_equals_extra_tokens() {
        let (heads, cfg) = small_heads();
        let hidden = zero_hidden(cfg.hidden_dim);
        let tree = heads.generate_tree(&hidden);
        let base = tree.paths[0][0];
        let result = heads.verify(&tree, base, &|_| base);
        // num_accepted == accepted_tokens.len() - 1
        assert_eq!(
            result.num_accepted,
            result.accepted_tokens.len().saturating_sub(1)
        );
    }

    // ── training_loss ─────────────────────────────────────────────────────────

    #[test]
    fn test_training_loss_shape() {
        let cfg = MedusaConfig::default();
        let heads = MedusaHeads::new(cfg.clone());
        let hidden = zero_hidden(cfg.hidden_dim);
        let targets: Vec<u32> = (0..cfg.num_heads as u32).collect();
        let result = heads.training_loss(&hidden, &targets);
        assert_eq!(result.per_head_loss.len(), cfg.num_heads);
    }

    #[test]
    fn test_training_loss_positive() {
        let cfg = MedusaConfig::default();
        let heads = MedusaHeads::new(cfg.clone());
        let hidden = ones_hidden(cfg.hidden_dim);
        let targets = vec![0_u32; cfg.num_heads];
        let result = heads.training_loss(&hidden, &targets);
        for (i, &loss) in result.per_head_loss.iter().enumerate() {
            assert!(
                loss >= 0.0,
                "head {i}: CE loss must be non-negative, got {loss}"
            );
        }
    }

    #[test]
    fn test_training_loss_total_is_mean() {
        let cfg = MedusaConfig::default();
        let heads = MedusaHeads::new(cfg.clone());
        let hidden = zero_hidden(cfg.hidden_dim);
        let targets = vec![1_u32; cfg.num_heads];
        let result = heads.training_loss(&hidden, &targets);
        let expected_mean =
            result.per_head_loss.iter().sum::<f32>() / result.per_head_loss.len() as f32;
        assert!(
            (result.total_loss - expected_mean).abs() < 1e-5,
            "total_loss must equal mean of per_head_loss"
        );
    }

    #[test]
    fn test_training_loss_near_zero_for_dominant_logit() {
        // Construct a head whose W2 row for token 0 is large (+100)
        // while all other rows are small (0).  CE loss for target=0 should be ≈ 0.
        let h = 4;
        let v = 8;
        let target: u32 = 0;

        // Build: W1 = identity, b1 = 0 → act = SiLU(hidden)
        // W2 row 0 = [100, 100, 100, 100], other rows = 0 → logit[0] >> logit[k>0]
        let mut w2 = vec![0.0_f32; v * h];
        for j in 0..h {
            w2[target as usize * h + j] = 100.0;
        }
        let head = MedusaHead {
            w1: {
                let mut m = vec![0.0_f32; h * h];
                for i in 0..h {
                    m[i * h + i] = 1.0;
                } // identity
                m
            },
            b1: vec![0.0; h],
            w2,
            b2: vec![0.0; v],
            hidden_dim: h,
            vocab_size: v,
        };

        // hidden = ones so act = SiLU(1) per neuron, all positive.
        let hidden = vec![1.0_f32; h];
        let logits = head.forward(&hidden);
        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}"
        );
    }
}