//! Beam search decoding for autoregressive language models.
//!
//! Implements standard beam search (Sutskever et al. 2014) and diverse beam search
//! (Vijayakumar et al. 2018) with Wu et al. 2016 length normalization.
//!
//! # Example
//!
//! ```rust
//! use rtx_inference::beam_search::{BeamSearchConfig, BeamSearchDecoder};
//!
//! let config = BeamSearchConfig {
//!     beam_width: 4,
//!     max_length: 32,
//!     eos_token_id: 2,
//!     length_penalty_alpha: 0.6,
//!     min_length: 1,
//!     no_repeat_ngram_size: 0,
//! };
//! let decoder = BeamSearchDecoder::new(config);
//!
//! let vocab_size = 32_000;
//! let results = decoder.decode(&[1], vocab_size, 2, |tokens| {
//!     let mut lp = vec![-10.0_f32; vocab_size];
//!     lp[5] = -0.1;
//!     lp[2] = -0.5; // EOS
//!     lp
//! });
//! assert!(!results.is_empty());
//! ```

use std::collections::HashSet;

// ---------------------------------------------------------------------------
// Beam
// ---------------------------------------------------------------------------

/// A single beam hypothesis tracking generated tokens and cumulative log-probability.
#[derive(Debug, Clone)]
pub struct Beam {
    /// Token IDs generated so far (including any initial prompt tokens).
    pub tokens: Vec<u32>,
    /// Cumulative sum of log-probabilities for all generated tokens.
    pub log_prob: f32,
    /// Whether this beam has emitted the EOS token and is considered finished.
    pub finished: bool,
}

impl Beam {
    /// Create a new, empty beam with the given initial log-probability.
    ///
    /// The token list starts empty; use `tokens` field to pre-populate prompt tokens.
    #[must_use]
    pub fn new(log_prob: f32) -> Self {
        Self {
            tokens: Vec::new(),
            log_prob,
            finished: false,
        }
    }

    /// Length-normalized score using the Wu et al. 2016 formula.
    ///
    /// `score = log_prob / length_penalty(len)`
    /// `length_penalty(l) = ((5 + l) / 6) ^ alpha`
    ///
    /// With `alpha = 0` the raw log-probability is returned unchanged.
    /// With `alpha = 1` full length normalization is applied.
    ///
    /// An empty token sequence returns `log_prob` directly to avoid division issues.
    #[must_use]
    pub fn normalized_score(&self, alpha: f32) -> f32 {
        let len = self.tokens.len();
        if len == 0 || alpha == 0.0 {
            return self.log_prob;
        }
        let penalty = ((5.0 + len as f32) / 6.0_f32).powf(alpha);
        self.log_prob / penalty
    }
}

// ---------------------------------------------------------------------------
// BeamSearchConfig
// ---------------------------------------------------------------------------

/// Configuration parameters for the beam search decoder.
#[derive(Debug, Clone)]
pub struct BeamSearchConfig {
    /// Number of beams to maintain at each step.
    pub beam_width: usize,
    /// Maximum number of tokens to generate per beam.
    pub max_length: usize,
    /// Token ID that marks end-of-sequence.
    pub eos_token_id: u32,
    /// Length penalty exponent (Wu et al. 2016). 0 = no penalty, 1 = full normalization.
    pub length_penalty_alpha: f32,
    /// Suppress EOS emission before this many generated tokens.
    pub min_length: usize,
    /// If > 0, block any token that would complete an already-seen n-gram of this size.
    pub no_repeat_ngram_size: usize,
}

impl Default for BeamSearchConfig {
    fn default() -> Self {
        Self {
            beam_width: 4,
            max_length: 128,
            eos_token_id: 2,
            length_penalty_alpha: 0.6,
            min_length: 1,
            no_repeat_ngram_size: 0,
        }
    }
}

// ---------------------------------------------------------------------------
// BeamSearchDecoder
// ---------------------------------------------------------------------------

/// Standard beam search decoder.
pub struct BeamSearchDecoder {
    config: BeamSearchConfig,
}

impl BeamSearchDecoder {
    /// Create a decoder with a custom configuration.
    #[must_use]
    pub fn new(config: BeamSearchConfig) -> Self {
        Self { config }
    }

    /// Create a decoder with default configuration.
    #[must_use]
    pub fn default_decoder() -> Self {
        Self::new(BeamSearchConfig::default())
    }

    /// Run beam search given a scoring function.
    ///
    /// # Arguments
    ///
    /// * `initial_tokens` — Prompt tokens prepended to every beam (not generated; used as
    ///   context only). Their length offsets `min_length` enforcement.
    /// * `vocab_size` — Total vocabulary size; used as the upper bound for candidate tokens.
    /// * `num_return` — Number of best beams to return.
    /// * `score_fn` — Callable that receives the full current token sequence (prompt +
    ///   generated) and returns a `Vec<f32>` of length `vocab_size` containing log-probs.
    ///
    /// # Returns
    ///
    /// Up to `num_return` beams sorted by descending normalized score. If no beam has
    /// emitted EOS the best still-active beams are returned instead.
    pub fn decode<F>(
        &self,
        initial_tokens: &[u32],
        vocab_size: usize,
        num_return: usize,
        score_fn: F,
    ) -> Vec<Beam>
    where
        F: Fn(&[u32]) -> Vec<f32>,
    {
        let bw = self.config.beam_width;
        let alpha = self.config.length_penalty_alpha;

        // Initialise with a single beam carrying all prompt tokens.
        let mut active: Vec<Beam> = vec![Beam {
            tokens: initial_tokens.to_vec(),
            log_prob: 0.0,
            finished: false,
        }];
        let mut finished: Vec<Beam> = Vec::new();

        for _step in 0..self.config.max_length {
            if active.is_empty() {
                break;
            }

            let mut candidates: Vec<Beam> = Vec::new();

            for beam in &active {
                let log_probs = score_fn(&beam.tokens);

                // Determine which token ids are blocked by the n-gram filter.
                let blocked = if self.config.no_repeat_ngram_size > 0 {
                    Self::blocked_tokens_for_ngram(&beam.tokens, self.config.no_repeat_ngram_size)
                } else {
                    HashSet::new()
                };

                // EOS is suppressed until the beam has generated at least `min_length` tokens
                // beyond the prompt.
                let generated_so_far = beam.tokens.len().saturating_sub(initial_tokens.len());
                let suppress_eos = generated_so_far < self.config.min_length;

                // Request extra tokens so that after filtering (EOS suppression, n-gram
                // blocking) we still have at least `bw` live candidates.  The +2 absorbs
                // one EOS suppression slot and one n-gram block slot simultaneously.
                let fetch_k = bw * 2 + 2;
                let top = Self::top_k_indices(&log_probs, fetch_k);
                let mut added = 0usize;
                for (token_id, lp) in top {
                    if added >= bw {
                        break;
                    }
                    let token_id = token_id as u32;
                    if blocked.contains(&token_id) {
                        continue;
                    }
                    if suppress_eos && token_id == self.config.eos_token_id {
                        continue;
                    }

                    let mut new_tokens = beam.tokens.clone();
                    new_tokens.push(token_id);
                    let new_lp = beam.log_prob + lp;
                    let is_finished = token_id == self.config.eos_token_id;
                    candidates.push(Beam {
                        tokens: new_tokens,
                        log_prob: new_lp,
                        finished: is_finished,
                    });
                    added += 1;
                }
            }

            // Sort all candidates by normalized score descending.
            candidates.sort_by(|a, b| {
                b.normalized_score(alpha)
                    .partial_cmp(&a.normalized_score(alpha))
                    .unwrap_or(std::cmp::Ordering::Equal)
            });

            // Partition into finished and active, keeping at most `bw` active beams.
            active.clear();
            for beam in candidates.into_iter().take(bw * 2) {
                if beam.finished {
                    finished.push(beam);
                } else if active.len() < bw {
                    active.push(beam);
                }
            }

            // Early stopping: if we already have enough finished beams and no active beam
            // can possibly outscore the worst finished beam, we are done.
            if finished.len() >= num_return && !active.is_empty() {
                let worst_finished = finished
                    .iter()
                    .map(|b| b.normalized_score(alpha))
                    .fold(f32::INFINITY, f32::min);
                let best_active = active
                    .iter()
                    .map(|b| b.normalized_score(alpha))
                    .fold(f32::NEG_INFINITY, f32::max);
                if best_active < worst_finished {
                    break;
                }
            }
        }

        // Any remaining active beams are folded into finished so the caller always gets
        // results even when no EOS was produced within `max_length`.
        finished.extend(active);
        finished.sort_by(|a, b| {
            b.normalized_score(alpha)
                .partial_cmp(&a.normalized_score(alpha))
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        finished.into_iter().take(num_return).collect()
    }

    /// Return the top-`k` `(index, log_prob)` pairs from `log_probs`, sorted descending.
    ///
    /// If `k` is greater than or equal to the length of `log_probs`, all entries are
    /// returned (still sorted descending).
    #[must_use]
    pub fn top_k_indices(log_probs: &[f32], k: usize) -> Vec<(usize, f32)> {
        let k = k.min(log_probs.len());
        // Partial sort: O(n log k) via a min-heap maintained at size k.
        // For typical vocab sizes (32k–256k) and small k (beam width ≤ 16) this is
        // substantially faster than a full sort.
        use std::cmp::Reverse;
        use std::collections::BinaryHeap;

        // BinaryHeap is a max-heap; we wrap in Reverse to get a min-heap so we can
        // efficiently evict the smallest element as we scan.
        // The heap stores (ordered_float_bits, index) where bits are from f32::to_bits
        // reinterpreted for ordering. We use a small wrapper instead.
        #[derive(PartialEq)]
        struct OrdF32(f32);

        impl Eq for OrdF32 {}

        impl PartialOrd for OrdF32 {
            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
                Some(self.cmp(other))
            }
        }

        impl Ord for OrdF32 {
            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
                self.0
                    .partial_cmp(&other.0)
                    .unwrap_or(std::cmp::Ordering::Equal)
            }
        }

        // min-heap of size k: smallest value is at the top and will be evicted first.
        let mut heap: BinaryHeap<Reverse<(OrdF32, usize)>> = BinaryHeap::with_capacity(k + 1);

        for (idx, &lp) in log_probs.iter().enumerate() {
            heap.push(Reverse((OrdF32(lp), idx)));
            if heap.len() > k {
                heap.pop();
            }
        }

        // Drain heap into a Vec and sort descending.
        let mut result: Vec<(usize, f32)> = heap
            .into_iter()
            .map(|Reverse((OrdF32(lp), idx))| (idx, lp))
            .collect();
        result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        result
    }

    /// Compute the set of token IDs that would complete an already-seen n-gram.
    ///
    /// For a token sequence `tokens` and n-gram size `n`, we look at the last `n-1`
    /// tokens as a suffix key and find all tokens `t` such that `suffix ++ [t]` appears
    /// earlier in `tokens`. Those tokens are "blocked" because appending them would
    /// produce a repeated n-gram.
    ///
    /// Returns an empty set when there are fewer than `n-1` tokens (no suffix to match).
    #[must_use]
    pub fn blocked_tokens_for_ngram(tokens: &[u32], n: usize) -> HashSet<u32> {
        if n == 0 || tokens.len() < n {
            return HashSet::new();
        }

        let suffix_len = n - 1;
        let suffix = &tokens[tokens.len() - suffix_len..];

        let mut blocked = HashSet::new();

        // Slide a window of length `suffix_len` over all earlier positions.
        // For each position i where tokens[i..i+suffix_len] == suffix, the token at
        // position i+suffix_len would complete a repeated n-gram.
        let end = tokens.len() - suffix_len; // exclusive upper bound for window starts
        for i in 0..end {
            if &tokens[i..i + suffix_len] == suffix {
                blocked.insert(tokens[i + suffix_len]);
            }
        }

        blocked
    }
}

impl Default for BeamSearchDecoder {
    fn default() -> Self {
        Self::new(BeamSearchConfig::default())
    }
}

// ---------------------------------------------------------------------------
// DiverseBeamSearchDecoder
// ---------------------------------------------------------------------------

/// Diverse beam search decoder (Vijayakumar et al. 2018).
///
/// Divides the total beam budget into `num_groups` groups of size
/// `config.beam_width / num_groups`. Each group runs standard beam search at each
/// step, but tokens already chosen by earlier groups in the same step are penalised
/// by `diversity_penalty` subtracted from their log-probability.
///
/// This encourages the returned hypotheses to differ from one another while still
/// being high-quality.
pub struct DiverseBeamSearchDecoder {
    config: BeamSearchConfig,
    /// Number of beam groups. Must divide `config.beam_width` evenly.
    pub num_groups: usize,
    /// Score penalty applied per token already chosen by a preceding group at the same step.
    pub diversity_penalty: f32,
}

impl DiverseBeamSearchDecoder {
    /// Create a diverse beam search decoder.
    ///
    /// # Panics
    ///
    /// Panics in debug mode if `num_groups` is 0 or exceeds `config.beam_width`.
    #[must_use]
    pub fn new(config: BeamSearchConfig, num_groups: usize, diversity_penalty: f32) -> Self {
        debug_assert!(num_groups > 0, "num_groups must be at least 1");
        debug_assert!(
            num_groups <= config.beam_width,
            "num_groups must not exceed beam_width"
        );
        Self {
            config,
            num_groups,
            diversity_penalty,
        }
    }

    /// Run diverse beam search.
    ///
    /// Groups run sequentially at each step. Group `g` penalises every token that was
    /// already selected as the first new token by groups `0..g` at the current step.
    ///
    /// Returns up to `num_return` beams sorted by descending normalized score.
    pub fn decode<F>(
        &self,
        initial_tokens: &[u32],
        vocab_size: usize,
        num_return: usize,
        score_fn: F,
    ) -> Vec<Beam>
    where
        F: Fn(&[u32]) -> Vec<f32>,
    {
        let alpha = self.config.length_penalty_alpha;
        let num_groups = self.num_groups.max(1);
        // Each group maintains its own set of active beams.
        let beams_per_group = (self.config.beam_width / num_groups).max(1);

        let initial_beam = Beam {
            tokens: initial_tokens.to_vec(),
            log_prob: 0.0,
            finished: false,
        };

        // One Vec<Beam> per group.
        let mut group_active: Vec<Vec<Beam>> = (0..num_groups)
            .map(|_| vec![initial_beam.clone()])
            .collect();
        let mut finished: Vec<Beam> = Vec::new();

        for _step in 0..self.config.max_length {
            // Track which tokens have been committed by each group this step so that
            // subsequent groups can be penalised.
            let mut chosen_tokens: Vec<HashSet<u32>> = vec![HashSet::new(); num_groups];

            for g in 0..num_groups {
                let active = &group_active[g];
                if active.is_empty() {
                    continue;
                }

                // Tokens penalised are those committed by groups 0..g.
                let penalised: HashSet<u32> = chosen_tokens[..g]
                    .iter()
                    .flat_map(|s| s.iter().copied())
                    .collect();

                let mut candidates: Vec<Beam> = Vec::new();

                for beam in active {
                    let log_probs = score_fn(&beam.tokens);

                    let blocked = if self.config.no_repeat_ngram_size > 0 {
                        BeamSearchDecoder::blocked_tokens_for_ngram(
                            &beam.tokens,
                            self.config.no_repeat_ngram_size,
                        )
                    } else {
                        HashSet::new()
                    };

                    let generated_so_far = beam.tokens.len().saturating_sub(initial_tokens.len());
                    let suppress_eos = generated_so_far < self.config.min_length;

                    // Apply diversity penalty before selecting top-k.
                    let adjusted: Vec<f32> = log_probs
                        .iter()
                        .enumerate()
                        .map(|(i, &lp)| {
                            if penalised.contains(&(i as u32)) {
                                lp - self.diversity_penalty
                            } else {
                                lp
                            }
                        })
                        .collect();

                    let fetch_k = beams_per_group * 2 + 2;
                    let top = BeamSearchDecoder::top_k_indices(&adjusted, fetch_k);
                    let mut added = 0usize;
                    for (token_id, lp) in top {
                        if added >= beams_per_group {
                            break;
                        }
                        let token_id = token_id as u32;
                        if blocked.contains(&token_id) {
                            continue;
                        }
                        if suppress_eos && token_id == self.config.eos_token_id {
                            continue;
                        }

                        // Record what each group committed.
                        chosen_tokens[g].insert(token_id);
                        added += 1;

                        // Use the un-penalised log-prob to keep beam scores comparable.
                        let raw_lp = if token_id < vocab_size as u32 {
                            log_probs[token_id as usize]
                        } else {
                            lp
                        };

                        let mut new_tokens = beam.tokens.clone();
                        new_tokens.push(token_id);
                        let new_lp = beam.log_prob + raw_lp;
                        let is_finished = token_id == self.config.eos_token_id;
                        candidates.push(Beam {
                            tokens: new_tokens,
                            log_prob: new_lp,
                            finished: is_finished,
                        });
                    }
                }

                // Sort and prune.
                candidates.sort_by(|a, b| {
                    b.normalized_score(alpha)
                        .partial_cmp(&a.normalized_score(alpha))
                        .unwrap_or(std::cmp::Ordering::Equal)
                });

                let mut new_active: Vec<Beam> = Vec::new();
                for beam in candidates.into_iter().take(beams_per_group * 2) {
                    if beam.finished {
                        finished.push(beam);
                    } else if new_active.len() < beams_per_group {
                        new_active.push(beam);
                    }
                }
                group_active[g] = new_active;
            }

            // Check whether all groups are exhausted.
            if group_active.iter().all(|g| g.is_empty()) {
                break;
            }
        }

        // Collect remaining active beams as fallback results.
        for group in group_active {
            finished.extend(group);
        }

        finished.sort_by(|a, b| {
            b.normalized_score(alpha)
                .partial_cmp(&a.normalized_score(alpha))
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        finished.into_iter().take(num_return).collect()
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // ---- Helper ---------------------------------------------------------------

    fn make_score_fn(vocab_size: usize, eos_id: u32) -> impl Fn(&[u32]) -> Vec<f32> {
        move |_tokens: &[u32]| {
            let mut lp = vec![-10.0_f32; vocab_size];
            lp[5] = -0.1;
            lp[eos_id as usize] = -0.5;
            lp
        }
    }

    // ---- Beam::normalized_score -----------------------------------------------

    #[test]
    fn test_beam_normalized_score_alpha0() {
        // alpha=0 => score == log_prob regardless of length
        let mut beam = Beam::new(-3.0);
        beam.tokens = vec![1, 2, 3];
        let score = beam.normalized_score(0.0);
        assert!(
            (score - (-3.0)).abs() < 1e-6,
            "alpha=0 should return raw log_prob, got {score}"
        );
    }

    #[test]
    fn test_beam_normalized_score_alpha1() {
        // alpha=1, len=1 => penalty = ((5+1)/6)^1 = 1.0 => score == log_prob
        let mut beam = Beam::new(-6.0);
        beam.tokens = vec![42];
        let score = beam.normalized_score(1.0);
        let expected = -6.0_f32 / ((5.0 + 1.0) / 6.0_f32).powf(1.0);
        assert!(
            (score - expected).abs() < 1e-5,
            "Expected {expected}, got {score}"
        );
    }

    #[test]
    fn test_beam_normalized_score_longer_beam_wins() {
        // Two beams with the same total log_prob; the longer one gets a better
        // (less negative) normalised score when alpha > 0.
        let mut short_beam = Beam::new(-6.0);
        short_beam.tokens = vec![1];

        let mut long_beam = Beam::new(-6.0);
        long_beam.tokens = vec![1, 2, 3, 4, 5];

        let alpha = 0.6;
        let short_score = short_beam.normalized_score(alpha);
        let long_score = long_beam.normalized_score(alpha);
        assert!(
            long_score > short_score,
            "Longer beam should have higher normalised score: long={long_score}, short={short_score}"
        );
    }

    // ---- top_k_indices --------------------------------------------------------

    #[test]
    fn test_top_k_indices_returns_k() {
        let lp: Vec<f32> = (0..20).map(|i| -(i as f32)).collect();
        let result = BeamSearchDecoder::top_k_indices(&lp, 5);
        assert_eq!(result.len(), 5, "Should return exactly 5 elements");
    }

    #[test]
    fn test_top_k_indices_descending() {
        let lp = vec![-3.0, -1.0, -4.0, -1.5, -2.0];
        let result = BeamSearchDecoder::top_k_indices(&lp, 3);
        assert_eq!(result.len(), 3);
        // Verify descending order.
        for window in result.windows(2) {
            assert!(
                window[0].1 >= window[1].1,
                "Not sorted descending: {:?}",
                result
            );
        }
        // Top element should be index 1 (log_prob -1.0).
        assert_eq!(result[0].0, 1);
    }

    #[test]
    fn test_top_k_larger_than_vocab() {
        let lp = vec![-1.0, -2.0, -3.0];
        // k=10 > vocab_size=3 => return all 3
        let result = BeamSearchDecoder::top_k_indices(&lp, 10);
        assert_eq!(result.len(), 3);
    }

    // ---- blocked_tokens_for_ngram ---------------------------------------------

    #[test]
    fn test_blocked_ngrams_empty() {
        let blocked = BeamSearchDecoder::blocked_tokens_for_ngram(&[], 2);
        assert!(blocked.is_empty());
    }

    #[test]
    fn test_blocked_ngrams_bigram() {
        // Sequence: [A, B, C, A] — suffix is [A], look for earlier [A] => B is blocked.
        let tokens: Vec<u32> = vec![10, 20, 30, 10];
        let blocked = BeamSearchDecoder::blocked_tokens_for_ngram(&tokens, 2);
        // Earlier [10] at position 0 is followed by 20 => 20 is blocked.
        assert!(
            blocked.contains(&20),
            "Token 20 should be blocked; got {blocked:?}"
        );
    }

    #[test]
    fn test_blocked_ngrams_trigram() {
        // [A, B, C, A, B] — suffix is [A, B], look for earlier [A, B] at pos 0 => C(=30) blocked.
        let tokens: Vec<u32> = vec![10, 20, 30, 10, 20];
        let blocked = BeamSearchDecoder::blocked_tokens_for_ngram(&tokens, 3);
        assert!(
            blocked.contains(&30),
            "Token 30 should be blocked; got {blocked:?}"
        );
    }

    #[test]
    fn test_blocked_ngrams_no_repeat() {
        // No repeated n-gram => empty blocked set.
        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5];
        let blocked = BeamSearchDecoder::blocked_tokens_for_ngram(&tokens, 2);
        assert!(
            blocked.is_empty(),
            "No repeated bigrams, should be empty; got {blocked:?}"
        );
    }

    // ---- decode ---------------------------------------------------------------

    #[test]
    fn test_decode_single_step() {
        // With beam_width=1 and a deterministic scorer, we should get exactly one beam
        // after one step (prompt=[1], max_length=1).
        let vocab_size = 10;
        let eos_id = 2_u32;
        let config = BeamSearchConfig {
            beam_width: 1,
            max_length: 1,
            eos_token_id: eos_id,
            length_penalty_alpha: 0.0,
            min_length: 0,
            no_repeat_ngram_size: 0,
        };
        let decoder = BeamSearchDecoder::new(config);
        let score_fn = |_: &[u32]| {
            let mut lp = vec![-10.0_f32; vocab_size];
            lp[5] = -0.1;
            lp[eos_id as usize] = -0.5;
            lp
        };
        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
        assert_eq!(results.len(), 1);
        // With min_length=0, EOS is allowed. But token 5 has higher lp, so beam should end with 5.
        let last = *results[0].tokens.last().unwrap();
        assert_eq!(last, 5, "Greedy first token should be 5");
    }

    #[test]
    fn test_decode_returns_num_return() {
        let vocab_size = 32;
        let config = BeamSearchConfig {
            beam_width: 4,
            max_length: 10,
            eos_token_id: 2,
            length_penalty_alpha: 0.6,
            min_length: 1,
            no_repeat_ngram_size: 0,
        };
        let decoder = BeamSearchDecoder::new(config);
        let results = decoder.decode(&[1], vocab_size, 2, make_score_fn(vocab_size, 2));
        assert_eq!(results.len(), 2, "Should return exactly 2 beams");
    }

    #[test]
    fn test_decode_eos_stops_beam() {
        // Force EOS to be the highest-scoring token so the beam finishes immediately.
        let vocab_size = 10;
        let eos_id = 2_u32;
        let config = BeamSearchConfig {
            beam_width: 1,
            max_length: 50,
            eos_token_id: eos_id,
            length_penalty_alpha: 0.0,
            min_length: 0,
            no_repeat_ngram_size: 0,
        };
        let decoder = BeamSearchDecoder::new(config);
        let score_fn = |_: &[u32]| {
            let mut lp = vec![-10.0_f32; vocab_size];
            lp[eos_id as usize] = -0.1; // EOS is best
            lp
        };
        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
        assert_eq!(results.len(), 1);
        assert!(
            results[0].finished,
            "Beam should be marked finished after EOS"
        );
        assert_eq!(
            *results[0].tokens.last().unwrap(),
            eos_id,
            "Last token should be EOS"
        );
    }

    #[test]
    fn test_decode_min_length_suppresses_eos() {
        // With min_length=5 the beam must generate at least 5 tokens before EOS is allowed.
        let vocab_size = 10;
        let eos_id = 2_u32;
        let config = BeamSearchConfig {
            beam_width: 1,
            max_length: 20,
            eos_token_id: eos_id,
            length_penalty_alpha: 0.0,
            min_length: 5,
            no_repeat_ngram_size: 0,
        };
        let decoder = BeamSearchDecoder::new(config);
        // EOS is always the best token, but is suppressed until min_length is met.
        let score_fn = |_: &[u32]| {
            let mut lp = vec![-10.0_f32; vocab_size];
            lp[5] = -0.2;
            lp[eos_id as usize] = -0.1; // EOS beats others, but suppressed until len>=5
            lp
        };
        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
        assert!(!results.is_empty());
        // The beam must have at least 1 (prompt) + 5 (generated) = 6 tokens before finishing.
        let generated = results[0].tokens.len() - 1; // subtract prompt token
        assert!(
            generated >= 5,
            "Should have generated at least 5 tokens before EOS, got {generated}"
        );
    }

    #[test]
    fn test_decode_always_returns_result() {
        // Even if no EOS is emitted within max_length, active beams are returned.
        let vocab_size = 10;
        let config = BeamSearchConfig {
            beam_width: 2,
            max_length: 5,
            eos_token_id: 99, // EOS id that will never be produced by the scorer
            length_penalty_alpha: 0.0,
            min_length: 0,
            no_repeat_ngram_size: 0,
        };
        let decoder = BeamSearchDecoder::new(config);
        let score_fn = |_: &[u32]| {
            let mut lp = vec![-10.0_f32; vocab_size];
            lp[5] = -0.1;
            lp[3] = -0.2;
            lp
        };
        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
        assert!(
            !results.is_empty(),
            "Should always return at least one beam"
        );
    }

    #[test]
    fn test_decode_beam_width_1_greedy() {
        // beam_width=1 should behave identically to greedy decoding.
        let vocab_size = 10;
        let eos_id = 2_u32;
        let config = BeamSearchConfig {
            beam_width: 1,
            max_length: 8,
            eos_token_id: eos_id,
            length_penalty_alpha: 0.0,
            min_length: 3,
            no_repeat_ngram_size: 0,
        };
        let decoder = BeamSearchDecoder::new(config);
        // Scorer alternates between two tokens to create a deterministic sequence.
        let score_fn = |tokens: &[u32]| {
            let mut lp = vec![-10.0_f32; vocab_size];
            if tokens.len() % 2 == 0 {
                lp[7] = -0.1;
            } else {
                lp[3] = -0.1;
            }
            lp[eos_id as usize] = -0.5;
            lp
        };
        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
        assert_eq!(results.len(), 1);
        // Manually simulate greedy: prompt=[1](len=1, odd), so first pick is lp[3].
        // Sequence: [1, 3, 7, 3, ...] until min_length=3 allows EOS but EOS < 0.5 vs -0.1 so it
        // never fires. We just verify the length constraint held.
        let generated = results[0].tokens.len() - 1;
        assert!(generated >= 3 || results[0].finished);
    }

    #[test]
    fn test_decode_prefers_higher_logprob() {
        // With two beams and clear score separation, the higher-lp beam should rank first.
        let vocab_size = 10;
        let eos_id = 2_u32;
        let config = BeamSearchConfig {
            beam_width: 2,
            max_length: 2,
            eos_token_id: eos_id,
            length_penalty_alpha: 0.0,
            min_length: 0,
            no_repeat_ngram_size: 0,
        };
        let decoder = BeamSearchDecoder::new(config);
        // Always give token 5 a much better score.
        let score_fn = |_: &[u32]| {
            let mut lp = vec![-20.0_f32; vocab_size];
            lp[5] = -0.01; // very good
            lp[3] = -5.0; // much worse
            lp[eos_id as usize] = -0.5;
            lp
        };
        let results = decoder.decode(&[1], vocab_size, 2, score_fn);
        // The top beam should contain token 5.
        assert!(
            results[0].tokens.contains(&5),
            "Top beam should contain token 5"
        );
    }

    #[test]
    fn test_default_config() {
        let cfg = BeamSearchConfig::default();
        assert_eq!(cfg.beam_width, 4);
        assert_eq!(cfg.max_length, 128);
        assert_eq!(cfg.eos_token_id, 2);
        assert!((cfg.length_penalty_alpha - 0.6).abs() < 1e-6);
        assert_eq!(cfg.min_length, 1);
        assert_eq!(cfg.no_repeat_ngram_size, 0);
    }

    // ---- DiverseBeamSearchDecoder --------------------------------------------

    #[test]
    fn test_diverse_decode_returns_results() {
        let vocab_size = 20;
        let config = BeamSearchConfig {
            beam_width: 4,
            max_length: 8,
            eos_token_id: 2,
            length_penalty_alpha: 0.6,
            min_length: 1,
            no_repeat_ngram_size: 0,
        };
        let decoder = DiverseBeamSearchDecoder::new(config, 2, 1.0);
        let results = decoder.decode(&[1], vocab_size, 2, make_score_fn(vocab_size, 2));
        assert!(
            !results.is_empty(),
            "Diverse beam search should return at least one beam"
        );
    }

    #[test]
    fn test_diverse_decode_group1_equals_standard() {
        // With num_groups=1 and diversity_penalty=0, diverse beam search should be
        // equivalent to standard beam search.
        let vocab_size = 20;
        let config = BeamSearchConfig {
            beam_width: 4,
            max_length: 6,
            eos_token_id: 2,
            length_penalty_alpha: 0.0,
            min_length: 1,
            no_repeat_ngram_size: 0,
        };
        let std_decoder = BeamSearchDecoder::new(config.clone());
        let div_decoder = DiverseBeamSearchDecoder::new(config, 1, 0.0);

        let sf = make_score_fn(vocab_size, 2);
        let sf2 = make_score_fn(vocab_size, 2);

        let std_results = std_decoder.decode(&[1], vocab_size, 1, sf);
        let div_results = div_decoder.decode(&[1], vocab_size, 1, sf2);

        assert!(!std_results.is_empty());
        assert!(!div_results.is_empty());
        // Both should produce the same top beam sequence.
        assert_eq!(
            std_results[0].tokens, div_results[0].tokens,
            "1-group diverse should match standard beam search"
        );
    }
}