"""Batched evaluate() correctness — the speedup rewrite must preserve exact semantics."""
import torch

from src.config import Config
from src.data import build_examples, get_splits
from src.train import evaluate

from tests.test_eval_classification import StubModel, _sieve35


def test_batched_evaluate_known_model():
    """The {2,3,5,7} sieve is CORRECT on the whole [2,100] range (any composite <= 101 has a
    factor in {2,3,5,7}, and every in-range next prime has no small divisor) — so evaluate()
    against the sieve stub must return token acc 1.0 and exact-match 1.0 on both splits."""
    cfg = Config()
    train_in, val_in = get_splits(cfg)
    m = StubModel(_sieve35)
    for split_in in (train_in, val_in):
        ex = build_examples(split_in, cfg)
        tok, em, per = evaluate(m, ex, cfg)
        assert tok == 1.0, f"token acc {tok} (split size {len(ex)})"
        assert em == 1.0, f"exact-match {em} (split size {len(ex)})"
        assert len(per) == len(ex)
        assert all(ok for _, _, _, ok in per)


def test_batched_evaluate_matches_singleton_greedy():
    """For a REAL model, batched evaluate's greedy decode must equal greedy_decode per row
    (batch-invariance guarantees this — belt and braces against a regressed layout fix)."""
    from src.model_api import build_model, greedy_decode
    from src.data import make_batch, decode_tokens
    cfg = Config(model="rnn")
    torch.manual_seed(0)
    m = build_model(cfg)
    _, val_in = get_splits(cfg)
    val_ex = build_examples(val_in, cfg)
    _, em, per = evaluate(m, val_ex, cfg)
    # independent per-example greedy via the public helper
    for i, (x, target) in enumerate(val_ex):
        xi = torch.tensor(x, dtype=torch.long).unsqueeze(0)
        gen = greedy_decode(m, xi, cfg)[0].tolist()
        pred_alone = decode_tokens(gen, cfg)
        assert per[i][2] == pred_alone, f"example {i}: batched {per[i][2]} vs singleton {pred_alone}"