"""
Test module for the text generation functionality.
"""
import torch
from test_cases.test_utils import (
    load_test_data, 
    setup_language_model_with_test_data,
    compare_token_sequences,
    run_test,
    GENERATION_TEST_DATA_PATH
)

def test_generation():
    """Test the RNNLanguageModel generation functionality."""
    test_data = load_test_data(GENERATION_TEST_DATA_PATH)
    lm = setup_language_model_with_test_data(test_data)
    
    # Get input prefix tokens
    prefix_tokens = test_data["prefix_tokens"]

    # Test greedy generation
    torch.manual_seed(10601)  # Ensure deterministic output for testing
    greedy_tokens = lm.generate(prefix_tokens, max_tokens=5, temperature=0.0)
    compare_token_sequences(
        greedy_tokens,
        test_data["expected_greedy_tokens"],
        label="Greedy generation"
    )

    # Test temperature sampling
    torch.manual_seed(10601)  # Reset seed for reproducibility
    sampled_tokens = lm.generate(prefix_tokens, max_tokens=5, temperature=0.8)
    compare_token_sequences(
        sampled_tokens,
        test_data["expected_sampled_tokens"],
        label="Temperature sampling"
    )

    # Test longer generation
    torch.manual_seed(10601)  # Reset seed for reproducibility
    long_tokens = lm.generate(prefix_tokens, max_tokens=10, temperature=0.0)
    compare_token_sequences(
        long_tokens,
        test_data["expected_long_tokens"],
        label="Long generation"
    )

def test_text_generation():
    """Run text generation tests with formatting."""
    run_test(test_generation, "RNNLanguageModel.generate")