megahal_gen/

lib.rs

//! MegaHAL reply generation: babble, seeding, bidirectional generation,
//! and surprise evaluation.
//!
//! This crate implements the core MegaHAL reply generation algorithm:
//!
//! 1. Seed a starting symbol from the keyword list.
//! 2. Run a forward phase from the seed to generate the rest of the sentence.
//! 3. Run a backward phase from the seed to generate the beginning.
//! 4. Evaluate candidates by surprise scoring.
//! 5. Select the highest-scoring candidate.
//!
//! It implements MegaHAL's keyword-biased random walk: keywords encountered during
//! the walk are greedily selected, while non-keywords fall back to probability-weighted
//! selection.

use std::collections::HashSet;
use std::time::{Duration, Instant};

use megahal_markov::{BidirectionalModel, ContextWindow};
use ngram_trie::Trie;
use rand::{Rng, RngExt};
use symbol_core::{ERROR_ID, FIN_ID, Symbol, SymbolId};
use symbol_dict::SymbolDict;

/// Controls how many candidate replies are generated before selecting the best.
#[derive(Debug, Clone)]
pub enum GenerationLimit {
    /// Stop after the given duration.
    Timeout(Duration),
    /// Stop after the given number of iterations.
    Iterations(usize),
    /// Stop when either limit is reached.
    Both {
        timeout: Duration,
        max_iterations: usize,
    },
}

impl Default for GenerationLimit {
    fn default() -> Self {
        // Match original MegaHAL: 1-second timeout.
        GenerationLimit::Timeout(Duration::from_secs(1))
    }
}

/// Generate the best reply for given input tokens and keywords.
///
/// Runs the candidate generation loop from MEGAHAL_SPEC.md Section 7.1:
/// 1. Generate a baseline reply with empty keywords.
/// 2. Repeatedly generate candidates with keywords, scoring each by surprise.
/// 3. Return the highest-scoring candidate that differs from the input.
///
/// `keywords` is the ordered keyword list from `extract_keywords`; input order
/// is preserved so that `seed` scans keywords in the same order as C's
/// `make_keywords` dictionary (`megahal.c:2273-2342`).  A `HashSet` is built
/// internally for the O(1) membership checks in `babble` and `evaluate_reply`.
///
/// C's loop is a `do/while` (`megahal.c:2228-2240`): it always generates and
/// evaluates at least one keyword-seeded candidate before checking the limit.
pub fn generate_reply<S, R>(
    model: &BidirectionalModel<S>,
    input_tokens: &[S],
    keywords: &[S],
    aux_set: &HashSet<S>,
    limit: &GenerationLimit,
    rng: &mut R,
) -> Vec<S>
where
    S: Symbol + AsRef<[u8]>,
    R: Rng,
{
    // Build a HashSet from the ordered keywords for O(1) membership checks.
    let keywords_set: HashSet<S> = keywords.iter().cloned().collect();

    let empty_keywords: &[S] = &[];
    let empty_keywords_set: HashSet<S> = HashSet::new();
    let empty_aux = HashSet::new();

    // Baseline reply (no keyword bias). Per C's `dissimilar()` check
    // (megahal.c:2215-2218), drop the baseline when it equals the input.
    let mut best = generate_one_reply(model, empty_keywords, &empty_keywords_set, &empty_aux, rng);
    if tokens_equal(&best, input_tokens) {
        best = Vec::new();
    }

    let mut max_surprise: f64 = -1.0;
    let start = Instant::now();
    let mut iterations: usize = 0;

    // C's loop is do/while: generate first, check limit after.
    loop {
        let candidate = generate_one_reply(model, keywords, &keywords_set, aux_set, rng);
        let surprise = evaluate_reply(model, &candidate, &keywords_set);

        if surprise > max_surprise && !tokens_equal(&candidate, input_tokens) {
            max_surprise = surprise;
            best = candidate;
        }

        iterations += 1;

        // Check limits after generating (matching C's do/while).
        match limit {
            GenerationLimit::Timeout(d) => {
                if start.elapsed() >= *d {
                    break;
                }
            }
            GenerationLimit::Iterations(n) => {
                if iterations >= *n {
                    break;
                }
            }
            GenerationLimit::Both {
                timeout,
                max_iterations,
            } => {
                if start.elapsed() >= *timeout || iterations >= *max_iterations {
                    break;
                }
            }
        }
    }

    best
}

/// Generate a single candidate reply (forward + backward phases).
///
/// MEGAHAL_SPEC.md Section 7.2.
fn generate_one_reply<S, R>(
    model: &BidirectionalModel<S>,
    keywords: &[S],
    keywords_set: &HashSet<S>,
    aux_set: &HashSet<S>,
    rng: &mut R,
) -> Vec<S>
where
    S: Symbol + AsRef<[u8]>,
    R: Rng,
{
    let mut reply: Vec<SymbolId> = Vec::new();
    let mut used_key = false;

    // Forward phase. Per C `reply()` (megahal.c:2420-2471), the backward phase
    // always runs even when forward produces nothing.
    let mut ctx = model.forward_context();

    let seed_id = seed(model, keywords, aux_set, rng);
    if seed_id != ERROR_ID && seed_id != FIN_ID {
        reply.push(seed_id);
        ctx.advance(&model.forward, seed_id);

        loop {
            let sym = babble(
                &model.forward,
                &ctx,
                &model.dictionary,
                keywords_set,
                aux_set,
                &reply,
                &mut used_key,
                rng,
            );
            if sym == ERROR_ID || sym == FIN_ID {
                break;
            }
            reply.push(sym);
            ctx.advance(&model.forward, sym);
        }
    }

    // Backward phase.
    let mut ctx = model.backward_context();

    // Re-establish backward context from the beginning of the reply.
    // Spec 7.2.3: walk from index min(reply_length-1, order) down to 0.
    // This matches the C code: for(i=MIN(size-1,order); i>=0; i--)
    if !reply.is_empty() {
        let start = (reply.len() - 1).min(model.order as usize);
        for i in (0..=start).rev() {
            ctx.advance(&model.backward, reply[i]);
        }
    }

    loop {
        let sym = babble(
            &model.backward,
            &ctx,
            &model.dictionary,
            keywords_set,
            aux_set,
            &reply,
            &mut used_key,
            rng,
        );
        if sym == ERROR_ID || sym == FIN_ID {
            break;
        }
        reply.insert(0, sym);
        ctx.advance(&model.backward, sym);
    }

    resolve_ids(model, &reply)
}

/// Select a seed symbol for forward generation.
///
/// MEGAHAL_SPEC.md Section 7.2.1.
///
/// `keywords` is the ordered slice from `extract_keywords`, preserving the
/// input order C's `make_keywords` builds (`megahal.c:2273-2342`).  The scan
/// starts at a random index and wraps around, matching C's `seed()`
/// (`megahal.c:2694-2706`).  The `.sort()` that was here before introduced a
/// different distribution because sorted order differs from input order.
fn seed<S, R>(
    model: &BidirectionalModel<S>,
    keywords: &[S],
    aux_set: &HashSet<S>,
    rng: &mut R,
) -> SymbolId
where
    S: Symbol + AsRef<[u8]>,
    R: Rng,
{
    let root = model.forward.root();
    let children = model.forward.children(root);

    // Keywords are scanned first: a keyword that is in the dictionary and not
    // auxiliary seeds the reply even when the forward root has no children,
    // matching C seed() (megahal.c:2697-2706).
    if !keywords.is_empty() {
        let start = rng.random_range(0..keywords.len());

        for offset in 0..keywords.len() {
            let idx = (start + offset) % keywords.len();
            let kw = &keywords[idx];

            if let Some(id) = model.dictionary.find(kw)
                && !aux_set.contains(kw)
            {
                return id;
            }
        }
    }

    // Default: a random child of the forward root, or ERROR if it has none.
    if children.is_empty() {
        return ERROR_ID;
    }
    let idx = rng.random_range(0..children.len());
    model.forward.node(children[idx]).symbol
}

/// Keyword-biased random symbol selection (the "babble" function).
///
/// MEGAHAL_SPEC.md Section 7.3.
#[allow(clippy::too_many_arguments)]
fn babble<S, R>(
    trie: &Trie,
    ctx: &ContextWindow,
    dict: &SymbolDict<S>,
    keywords: &HashSet<S>,
    aux_set: &HashSet<S>,
    reply: &[SymbolId],
    used_key: &mut bool,
    rng: &mut R,
) -> SymbolId
where
    S: Symbol + AsRef<[u8]>,
    R: Rng,
{
    // Find deepest available context.
    let node_ref = match ctx.deepest() {
        Some(r) => r,
        None => return ERROR_ID,
    };

    let node = trie.node(node_ref);
    let children = trie.children(node_ref);

    if children.is_empty() {
        return ERROR_ID;
    }

    let branch = children.len();
    // C `babble()` calls `rnd(node->usage)` which returns 0 when usage is 0
    // and the loop falls through; `rng.random_range(0..0)` panics, so guard
    // explicitly and treat usage==0 as sentence-terminating.
    if node.usage == 0 {
        return ERROR_ID;
    }
    let mut i = rng.random_range(0..branch);
    let mut count = rng.random_range(0..node.usage as i64);

    loop {
        let child_ref = children[i];
        let child = trie.node(child_ref);
        let sym = child.symbol;

        // Check if this symbol is a keyword we should greedily select.
        let word = dict.resolve(sym);
        let is_keyword = keywords.contains(word);
        let is_aux = aux_set.contains(word);
        let already_in_reply = reply.contains(&sym);

        if is_keyword && (*used_key || !is_aux) && !already_in_reply {
            *used_key = true;
            return sym;
        }

        // Otherwise, probability-weighted selection.
        count -= child.count as i64;
        if count < 0 {
            return sym;
        }

        i = (i + 1) % branch;
    }
}

/// Score a candidate reply by surprise (Shannon entropy of keywords in context).
///
/// MEGAHAL_SPEC.md Section 8.
fn evaluate_reply<S>(model: &BidirectionalModel<S>, candidate: &[S], keywords: &HashSet<S>) -> f64
where
    S: Symbol + AsRef<[u8]>,
{
    if candidate.is_empty() {
        return 0.0;
    }

    let mut entropy: f64 = 0.0;
    let mut num: usize = 0;

    // Forward evaluation.
    let mut ctx = model.forward_context();
    for token in candidate {
        let sym_id = match model.dictionary.find(token) {
            Some(id) => id,
            None => continue,
        };

        if keywords.contains(token) {
            let mut prob: f64 = 0.0;
            let mut ctx_count: usize = 0;

            for j in 0..model.order as usize {
                if let Some(parent_ref) = ctx.at_depth(j)
                    && let Some(child_ref) = model.forward.find_child(parent_ref, sym_id)
                {
                    let child = model.forward.node(child_ref);
                    let parent = model.forward.node(parent_ref);
                    if parent.usage > 0 {
                        prob += child.count as f64 / parent.usage as f64;
                        ctx_count += 1;
                    }
                }
            }

            if ctx_count > 0 {
                entropy -= (prob / ctx_count as f64).ln();
            }
            num += 1;
        }

        ctx.advance(&model.forward, sym_id);
    }

    // Backward evaluation.
    let mut ctx = model.backward_context();
    for token in candidate.iter().rev() {
        let sym_id = match model.dictionary.find(token) {
            Some(id) => id,
            None => continue,
        };

        if keywords.contains(token) {
            let mut prob: f64 = 0.0;
            let mut ctx_count: usize = 0;

            for j in 0..model.order as usize {
                if let Some(parent_ref) = ctx.at_depth(j)
                    && let Some(child_ref) = model.backward.find_child(parent_ref, sym_id)
                {
                    let child = model.backward.node(child_ref);
                    let parent = model.backward.node(parent_ref);
                    if parent.usage > 0 {
                        prob += child.count as f64 / parent.usage as f64;
                        ctx_count += 1;
                    }
                }
            }

            if ctx_count > 0 {
                entropy -= (prob / ctx_count as f64).ln();
            }
            num += 1;
        }

        ctx.advance(&model.backward, sym_id);
    }

    // Length penalty.
    if num >= 8 {
        entropy /= ((num - 1) as f64).sqrt();
    }
    if num >= 16 {
        entropy /= num as f64;
    }

    entropy
}

/// Capitalize a token sequence per MegaHAL sentence-case rules.
///
/// MEGAHAL_SPEC.md Section 9.1. Mirrors C `capitalize()` in megahal.c:
/// the sentence-start flag is set when a `!.?` is followed by whitespace,
/// not by the punctuation itself.
pub fn capitalize(tokens: &[String]) -> String {
    let raw: String = tokens.concat();
    let mut result = Vec::with_capacity(raw.len());
    let bytes = raw.as_bytes();
    let mut start = true;

    for (i, &b) in bytes.iter().enumerate() {
        if b.is_ascii_alphabetic() {
            if start {
                result.push(b.to_ascii_uppercase());
            } else {
                result.push(b.to_ascii_lowercase());
            }
            start = false;
        } else {
            result.push(b);
        }
        if i > 2 && b.is_ascii_whitespace() && matches!(bytes[i - 1], b'!' | b'.' | b'?') {
            start = true;
        }
    }

    String::from_utf8(result).unwrap_or(raw)
}

/// Check if two token sequences are equal (case-insensitive, for dissimilarity test).
fn tokens_equal<S: Symbol>(a: &[S], b: &[S]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    a.iter().zip(b.iter()).all(|(x, y)| x == y)
}

/// Resolve a sequence of SymbolIds back to Symbol values.
fn resolve_ids<S: Symbol>(model: &BidirectionalModel<S>, ids: &[SymbolId]) -> Vec<S> {
    ids.iter()
        .map(|&id| model.dictionary.resolve(id).clone())
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::SeedableRng;
    use rand::rngs::SmallRng;

    // --- Test infrastructure ---

    #[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
    struct TSym(String);

    impl Symbol for TSym {
        fn error() -> Self {
            TSym("<ERROR>".into())
        }
        fn fin() -> Self {
            TSym("<FIN>".into())
        }
    }

    impl AsRef<[u8]> for TSym {
        fn as_ref(&self) -> &[u8] {
            self.0.as_bytes()
        }
    }

    fn ts(s: &str) -> TSym {
        TSym(s.to_string())
    }

    fn trained_model(order: u8, sentences: &[&[&str]]) -> BidirectionalModel<TSym> {
        let mut model = BidirectionalModel::new(order);
        for sentence in sentences {
            let tokens: Vec<TSym> = sentence.iter().map(|&s| ts(s)).collect();
            model.learn(&tokens);
        }
        model
    }

    fn make_rng(s: u64) -> SmallRng {
        SmallRng::seed_from_u64(s)
    }

    // --- GenerationLimit tests ---

    #[test]
    fn generation_limit_default_is_timeout() {
        let limit = GenerationLimit::default();
        assert!(matches!(limit, GenerationLimit::Timeout(_)));
    }

    // --- capitalize tests ---

    #[test]
    fn capitalize_basic() {
        let tokens = vec![
            "hello".to_string(),
            " ".to_string(),
            "world".to_string(),
            ".".to_string(),
        ];
        assert_eq!(capitalize(&tokens), "Hello world.");
    }

    #[test]
    fn capitalize_after_period() {
        let tokens = vec![
            "hello".to_string(),
            ". ".to_string(),
            "world".to_string(),
            ".".to_string(),
        ];
        assert_eq!(capitalize(&tokens), "Hello. World.");
    }

    #[test]
    fn capitalize_empty() {
        let tokens: Vec<String> = vec![];
        assert_eq!(capitalize(&tokens), "");
    }

    #[test]
    fn capitalize_after_exclamation() {
        let tokens = vec![
            "wow".to_string(),
            "! ".to_string(),
            "amazing".to_string(),
            ".".to_string(),
        ];
        assert_eq!(capitalize(&tokens), "Wow! Amazing.");
    }

    #[test]
    fn capitalize_after_question() {
        let tokens = vec![
            "really".to_string(),
            "? ".to_string(),
            "yes".to_string(),
            ".".to_string(),
        ];
        assert_eq!(capitalize(&tokens), "Really? Yes.");
    }

    #[test]
    fn capitalize_no_space_after_period_does_not_capitalize() {
        let tokens = vec!["a.b.c".to_string()];
        assert_eq!(capitalize(&tokens), "A.b.c");
    }

    #[test]
    fn capitalize_glued_sentences_not_split() {
        let tokens = vec!["hello.world".to_string()];
        assert_eq!(capitalize(&tokens), "Hello.world");
    }

    #[test]
    fn capitalize_leading_ellipsis() {
        // First alpha after a leading run of dots still gets uppercased
        // (start flag was never cleared).
        let tokens = vec!["...hello".to_string()];
        assert_eq!(capitalize(&tokens), "...Hello");
    }

    // --- seed tests ---

    #[test]
    fn seed_selects_keyword() {
        let model = trained_model(
            2,
            &[
                &["THE", " ", "CAT", " ", "SAT"],
                &["THE", " ", "DOG", " ", "RAN"],
            ],
        );
        let kws = vec![ts("CAT")];
        let aux = HashSet::new();
        let mut rng = make_rng(42);
        let id = seed(&model, &kws, &aux, &mut rng);
        let cat_id = model.dictionary.find(&ts("CAT")).unwrap();
        assert_eq!(id, cat_id);
    }

    #[test]
    fn seed_skips_aux_keyword() {
        let model = trained_model(2, &[&["THE", " ", "MY", " ", "CAT"]]);
        let kws = vec![ts("MY")];
        let mut aux = HashSet::new();
        aux.insert(ts("MY"));
        let mut rng = make_rng(42);
        let id = seed(&model, &kws, &aux, &mut rng);
        // MY is aux-only → seed falls back to random child of forward root.
        let my_id = model.dictionary.find(&ts("MY")).unwrap();
        assert_ne!(id, my_id);
    }

    #[test]
    fn seed_with_empty_keywords_picks_random() {
        let model = trained_model(2, &[&["THE", " ", "CAT"]]);
        let kws: Vec<TSym> = vec![];
        let aux = HashSet::new();
        let mut rng = make_rng(42);
        let id = seed(&model, &kws, &aux, &mut rng);
        assert_ne!(id, ERROR_ID);
        assert_ne!(id, FIN_ID);
    }

    #[test]
    fn seed_returns_error_on_empty_model() {
        let model: BidirectionalModel<TSym> = BidirectionalModel::new(2);
        let kws: Vec<TSym> = vec![];
        let aux = HashSet::new();
        let mut rng = make_rng(42);
        let id = seed(&model, &kws, &aux, &mut rng);
        assert_eq!(id, ERROR_ID);
    }

    // seed visits keywords in input order, not sorted order.
    #[test]
    fn seed_visits_keywords_in_input_order() {
        // Train a model that knows ZEBRA and APPLE.
        let model = trained_model(2, &[&["ZEBRA", " ", "SAT"], &["APPLE", " ", "RAN"]]);
        // ZEBRA > APPLE alphabetically, so if seed sorted we would pick APPLE
        // first on many RNG seeds.  With input order [ZEBRA, APPLE], a start
        // index of 0 must land on ZEBRA first.
        let kws = vec![ts("ZEBRA"), ts("APPLE")];
        let aux: HashSet<TSym> = HashSet::new();

        // Force start index 0 by using a seeded RNG that produces 0 for
        // random_range(0..2).  Iterate a few seeds to find one that gives 0.
        // Both ZEBRA and APPLE are valid seeds (neither is aux), so whichever
        // index 0 points to must be returned.  We assert that at start=0 the
        // result is ZEBRA (input-order index 0), not APPLE (sorted index 0).
        let zebra_id = model.dictionary.find(&ts("ZEBRA")).unwrap();
        let apple_id = model.dictionary.find(&ts("APPLE")).unwrap();

        // With a fresh SmallRng(0), random_range(0..2) gives a deterministic
        // value; we just need to verify that the result is consistent with
        // input order (index 0 = ZEBRA), not sorted order (index 0 = APPLE).
        // Try many RNG seeds: for any seed that yields start==0, result must
        // be ZEBRA; for start==1, result must be APPLE.
        let mut found_start_zero = false;
        for seed_val in 0u64..200 {
            let mut rng = make_rng(seed_val);
            // Peek what start index would be chosen (same call as seed()).
            let mut rng_peek = make_rng(seed_val);
            let start_idx = rng_peek.random_range(0usize..2);
            let result = seed(&model, &kws, &aux, &mut rng);
            if start_idx == 0 {
                assert_eq!(
                    result, zebra_id,
                    "seed={seed_val}: start==0 should pick kws[0]=ZEBRA (input order), not APPLE (sorted order)"
                );
                found_start_zero = true;
            } else {
                assert_eq!(
                    result, apple_id,
                    "seed={seed_val}: start==1 should pick kws[1]=APPLE"
                );
            }
        }
        assert!(
            found_start_zero,
            "no RNG seed produced start index 0 in 200 tries"
        );
    }

    #[test]
    fn seed_returns_keyword_when_forward_root_has_no_children() {
        // A keyword in the dictionary seeds the reply even when the forward
        // root is childless, matching C seed() (megahal.c:2697-2706). This
        // state is reachable only via a degenerate loaded brain, since learning
        // couples dictionary and trie population.
        let mut model: BidirectionalModel<TSym> = BidirectionalModel::new(2);
        let hello_id = model.dictionary.intern(ts("HELLO"));
        assert!(model.forward.children(model.forward.root()).is_empty());

        let kws = vec![ts("HELLO")];
        let aux: HashSet<TSym> = HashSet::new();
        let mut rng = make_rng(0);
        assert_eq!(seed(&model, &kws, &aux, &mut rng), hello_id);
    }

    // --- evaluate_reply tests ---

    #[test]
    fn evaluate_empty_candidate_returns_zero() {
        let model = trained_model(2, &[&["A", "B", "C"]]);
        let kws = HashSet::new();
        let score = evaluate_reply(&model, &[], &kws);
        assert_eq!(score, 0.0);
    }

    #[test]
    fn evaluate_no_keywords_returns_zero() {
        let model = trained_model(2, &[&["A", "B", "C"]]);
        let kws = HashSet::new();
        let candidate = vec![ts("A"), ts("B"), ts("C")];
        let score = evaluate_reply(&model, &candidate, &kws);
        assert_eq!(score, 0.0);
    }

    #[test]
    fn evaluate_with_keywords_returns_positive() {
        let model = trained_model(2, &[&["A", "B", "C"], &["A", "B", "C"], &["A", "B", "C"]]);
        let mut kws = HashSet::new();
        kws.insert(ts("B"));
        let candidate = vec![ts("A"), ts("B"), ts("C")];
        let score = evaluate_reply(&model, &candidate, &kws);
        assert!(score > 0.0, "Expected positive surprise, got {score}");
    }

    #[test]
    fn evaluate_unknown_token_skipped() {
        let model = trained_model(2, &[&["A", "B", "C"]]);
        let mut kws = HashSet::new();
        kws.insert(ts("UNKNOWN"));
        // UNKNOWN is not in the dict → find returns None → skipped.
        let candidate = vec![ts("UNKNOWN")];
        let score = evaluate_reply(&model, &candidate, &kws);
        assert_eq!(score, 0.0);
    }

    // --- tokens_equal tests ---

    #[test]
    fn tokens_equal_same() {
        let a = vec![ts("A"), ts("B")];
        let b = vec![ts("A"), ts("B")];
        assert!(tokens_equal(&a, &b));
    }

    #[test]
    fn tokens_equal_different_length() {
        let a = vec![ts("A"), ts("B")];
        let b = vec![ts("A")];
        assert!(!tokens_equal(&a, &b));
    }

    #[test]
    fn tokens_equal_different_content() {
        let a = vec![ts("A"), ts("B")];
        let b = vec![ts("A"), ts("C")];
        assert!(!tokens_equal(&a, &b));
    }

    #[test]
    fn tokens_equal_both_empty() {
        let a: Vec<TSym> = vec![];
        let b: Vec<TSym> = vec![];
        assert!(tokens_equal(&a, &b));
    }

    // --- generate_reply integration tests ---

    #[test]
    fn generate_reply_empty_model() {
        let model: BidirectionalModel<TSym> = BidirectionalModel::new(2);
        let kws: Vec<TSym> = vec![];
        let aux = HashSet::new();
        let limit = GenerationLimit::Iterations(10);
        let mut rng = make_rng(42);
        let reply = generate_reply(&model, &[], &kws, &aux, &limit, &mut rng);
        assert!(reply.is_empty());
    }

    #[test]
    fn generate_reply_produces_output() {
        let model = trained_model(
            2,
            &[
                &["THE", " ", "CAT", " ", "SAT"],
                &["THE", " ", "DOG", " ", "RAN"],
                &["A", " ", "BIG", " ", "CAT"],
            ],
        );
        let kws = vec![ts("CAT")];
        let aux = HashSet::new();
        let limit = GenerationLimit::Iterations(10);
        let mut rng = make_rng(42);
        let reply = generate_reply(&model, &[], &kws, &aux, &limit, &mut rng);
        assert!(!reply.is_empty());
    }

    #[test]
    fn generate_reply_deterministic() {
        let build = || {
            let model = trained_model(
                2,
                &[
                    &["THE", " ", "CAT", " ", "SAT"],
                    &["THE", " ", "DOG", " ", "RAN"],
                ],
            );
            let kws = vec![ts("CAT")];
            let aux = HashSet::new();
            let limit = GenerationLimit::Iterations(50);
            let mut rng = make_rng(42);
            generate_reply(&model, &[], &kws, &aux, &limit, &mut rng)
        };
        assert_eq!(build(), build());
    }

    // Iterations(0) still generates one keyword-seeded candidate (C do/while).
    #[test]
    fn generate_reply_iterations_zero_still_evaluates_one_candidate() {
        let model = trained_model(
            2,
            &[
                &["THE", " ", "CAT", " ", "SAT"],
                &["THE", " ", "DOG", " ", "RAN"],
            ],
        );
        let kws = vec![ts("CAT")];
        let aux = HashSet::new();
        let limit = GenerationLimit::Iterations(0);
        let mut rng = make_rng(42);
        // C's do/while always runs the body once before checking the bound.
        // Iterations(0) means the limit check fires after iteration 1, so
        // exactly one keyword-seeded candidate is generated and evaluated.
        // The previous code checked the limit first and skipped generation.
        let reply = generate_reply(&model, &[], &kws, &aux, &limit, &mut rng);
        assert!(
            !reply.is_empty(),
            "Iterations(0) must still generate one keyword-seeded candidate per C do/while"
        );
    }

    #[test]
    fn generate_reply_with_timeout() {
        let model = trained_model(
            2,
            &[
                &["THE", " ", "CAT", " ", "SAT"],
                &["THE", " ", "DOG", " ", "RAN"],
                &["A", " ", "BIG", " ", "CAT"],
            ],
        );
        let kws = vec![ts("CAT")];
        let aux = HashSet::new();
        let limit = GenerationLimit::Timeout(Duration::from_millis(50));
        let mut rng = make_rng(42);
        let reply = generate_reply(&model, &[], &kws, &aux, &limit, &mut rng);
        assert!(!reply.is_empty());
    }

    #[test]
    fn generate_reply_with_both_limit() {
        let model = trained_model(
            2,
            &[
                &["THE", " ", "CAT", " ", "SAT"],
                &["THE", " ", "DOG", " ", "RAN"],
            ],
        );
        let kws = vec![ts("CAT")];
        let aux = HashSet::new();
        let limit = GenerationLimit::Both {
            timeout: Duration::from_millis(50),
            max_iterations: 10,
        };
        let mut rng = make_rng(42);
        let reply = generate_reply(&model, &[], &kws, &aux, &limit, &mut rng);
        assert!(!reply.is_empty());
    }

    #[test]
    fn generate_reply_with_aux_keywords() {
        let model = trained_model(
            2,
            &[
                &["MY", " ", "CAT", " ", "SAT"],
                &["YOUR", " ", "DOG", " ", "RAN"],
            ],
        );
        let kws = vec![ts("CAT"), ts("MY")];
        let mut aux = HashSet::new();
        aux.insert(ts("MY"));
        let limit = GenerationLimit::Iterations(20);
        let mut rng = make_rng(42);
        let reply = generate_reply(&model, &[], &kws, &aux, &limit, &mut rng);
        assert!(!reply.is_empty());
    }

    #[test]
    fn generate_reply_dissimilarity_test() {
        // If the candidate is identical to input, it should be rejected in favor
        // of a different candidate (when possible).
        let model = trained_model(
            2,
            &[
                &["A", " ", "B", " ", "C"],
                &["D", " ", "E", " ", "F"],
                &["A", " ", "B", " ", "C"],
            ],
        );
        let input = vec![ts("A"), ts(" "), ts("B"), ts(" "), ts("C")];
        let kws: Vec<TSym> = vec![];
        let aux = HashSet::new();
        let limit = GenerationLimit::Iterations(50);
        let mut rng = make_rng(42);
        let reply = generate_reply(&model, &input, &kws, &aux, &limit, &mut rng);
        // Verify that candidates identical to the input are filtered.
        let _ = reply;
    }
}