"""
Test module for the RNNLanguageModel implementation.
"""
import torch
from test_cases.test_utils import (
    compare_tensors, 
    load_test_data, 
    setup_language_model_with_test_data,
    run_test,
    LM_TEST_DATA_PATH
)

def test_rnn_language_model_forward():
    """Test the complete RNNLanguageModel forward pass implementation."""
    test_data = load_test_data(LM_TEST_DATA_PATH)
    lm = setup_language_model_with_test_data(test_data)
    
    # Get input tokens and run forward pass
    tokens = test_data["tokens"]
    token_logits, hidden_states, attn_inputs = lm(tokens)

    # Compare with expected outputs
    compare_tensors(
        token_logits,
        test_data["expected_token_logits"],
        label="RNNLanguageModel token_logits",
        atol=1e-4,
    )
    compare_tensors(
        hidden_states,
        test_data["expected_hidden_states"],
        label="RNNLanguageModel hidden_states",
        atol=1e-4,
    )
    compare_tensors(
        attn_inputs,
        test_data["expected_attn_inputs"],
        label="RNNLanguageModel attn_inputs",
        atol=1e-4,
    )

    # Test token selection
    token_logits_sample = test_data["token_logits_sample"]
    
    # Test greedy selection
    torch.manual_seed(42)  # Reset seed for reproducibility
    greedy_token = lm.select_token(token_logits_sample, temperature=0.0)
    if greedy_token != test_data["expected_greedy_token"]:
        raise ValueError(
            f"Greedy token selection mismatch! Expected {test_data['expected_greedy_token']}, got {greedy_token}"
        )

    # Test temperature sampling
    torch.manual_seed(42)  # Reset seed for reproducibility
    sampled_token = lm.select_token(token_logits_sample, temperature=0.8)
    if sampled_token != test_data["expected_sampled_token"]:
        raise ValueError(
            f"Sampled token selection mismatch! Expected {test_data['expected_sampled_token']}, got {sampled_token}"
        )

def test_rnn_language_model_step():
    """Test the step-by-step execution of RNNLanguageModel."""
    test_data = load_test_data(LM_TEST_DATA_PATH)
    lm = setup_language_model_with_test_data(test_data)
    
    # Test individual steps
    step_token = test_data["step_token"]
    
    # Step 1: Get embeddings
    step_embeddings = lm.embeddings(step_token)
    compare_tensors(
        step_embeddings, 
        test_data["expected_step_embeddings"], 
        label="Step embeddings",
        atol=1e-5
    )
    
    # Step 2: Process with RNN
    step_hidden_state, step_attn_input = lm.rnn.step(step_embeddings, None)
    compare_tensors(
        step_hidden_state, 
        test_data["expected_step_hidden_state"], 
        label="Step hidden state",
        atol=1e-5
    )
    compare_tensors(
        step_attn_input, 
        test_data["expected_step_attn_input"], 
        label="Step attention input",
        atol=1e-5
    )
    
    # Step 3: Process with attention
    step_output_state = lm.attention.step(step_attn_input.unsqueeze(1))
    compare_tensors(
        step_output_state, 
        test_data["expected_step_output_state"], 
        label="Step output state",
        atol=1e-5
    )
    
    # Step 4: Get token logits
    step_token_logits = lm.lm_head(step_output_state)
    compare_tensors(
        step_token_logits, 
        test_data["expected_step_token_logits"], 
        label="Step token logits",
        atol=1e-5
    )

def test_rnn_language_model():
    """Run all RNNLanguageModel tests."""
    run_test(test_rnn_language_model_forward, "RNNLanguageModel forward pass")
    run_test(test_rnn_language_model_step, "RNNLanguageModel step-by-step execution")