"""Tests for CUDA device handling, chunked evaluation, and scaling to 1000 and 100_000."""
import torch

from src.config import Config, parse_args
from src.data import build_examples, get_splits, next_prime, sieve_primes
from src.model_api import build_model
from src.train import evaluate, resolve_device


def test_device_resolution():
    assert resolve_device("cpu") == torch.device("cpu")
    auto_dev = resolve_device("auto")
    expected = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    assert auto_dev == expected


def test_config_device_args():
    cfg = parse_args(["rnn", "0", "--device", "cpu", "--eval_batch_size", "256"])
    assert cfg.device == "cpu"
    assert cfg.eval_batch_size == 256


def test_bisect_next_prime_large_range():
    primes = sieve_primes(100_100)
    # Check boundary primes
    assert next_prime(2, primes) == 3
    assert next_prime(97, primes) == 101
    assert next_prime(997, primes) == 1009
    assert next_prime(99_991, primes) == 100_003


def test_chunked_evaluate_matches_single_batch():
    """Verify that chunked evaluation with eval_batch_size=4 gives identical results to full batch."""
    cfg = Config(model="rnn", eval_batch_size=4)
    model = build_model(cfg)
    train_in, val_in = get_splits(cfg)
    val_ex = build_examples(val_in, cfg)

    # Chunked
    tok_chunked, em_chunked, per_chunked = evaluate(model, val_ex, cfg)

    # Single full batch
    cfg_full = Config(model="rnn", eval_batch_size=len(val_ex))
    tok_full, em_full, per_full = evaluate(model, val_ex, cfg_full)

    assert tok_chunked == tok_full
    assert em_chunked == em_full
    assert per_chunked == per_full


def test_range_extension_1000():
    cfg = Config(range_start=2, range_end=1000, model="transformer")
    train_in, val_in = get_splits(cfg)
    assert len(train_in) + len(val_in) == 999
    train_ex = build_examples(train_in, cfg)
    assert len(train_ex) == len(train_in)
    model = build_model(cfg)
    tok, em, per = evaluate(model, build_examples(val_in[:20], cfg), cfg)
    assert 0.0 <= em <= 1.0


def test_range_extension_100k_data_shapes():
    """Test data generation and model forward at 100k range without OOM."""
    cfg = Config(range_start=2, range_end=100_000, model="rnn")
    # Take a small slice of inputs from 100k range
    sample_inputs = [2, 42, 999, 50_000, 99_991]
    examples = build_examples(sample_inputs, cfg)
    assert len(examples) == 5
    model = build_model(cfg)
    tok, em, per = evaluate(model, examples, cfg)
    assert 0.0 <= em <= 1.0