"""Tests for the Phase-3 flags: task_mode=is_prime, train_frac, lr_decay."""
import torch

from src.config import Config, parse_args
from src.data import build_examples, decode_tokens, get_splits
from src.eval import probe_report
from src.model_api import build_model

from tests.test_eval_classification import StubModel


def _isprime_sieve(n):
    """Primality via {2,3,5,7} with the d<c guard — EXACT on [2,100] (every composite <=100
    has a prime factor <=10), wrong only on no-small-factor composites like 121."""
    if n < 2:
        return 0
    return int(all(n % d != 0 for d in (2, 3, 5, 7) if d < n))


def test_is_prime_examples():
    cfg = Config(task_mode="is_prime")
    ex = build_examples([2, 3, 4, 97, 100], cfg)
    got = [decode_tokens(y, cfg) for _, y in ex]
    assert got == [1, 1, 0, 1, 0], got


def test_is_prime_sieve_stub_perfect_in_range():
    """The {2,3,5,7} sieve is exact on [2,100] -> evaluate() must see EM 1.0."""
    from src.train import evaluate
    cfg = Config(task_mode="is_prime")
    train_in, val_in = get_splits(cfg)
    m = StubModel(_isprime_sieve)
    for split in (train_in, val_in):
        tok, em, _ = evaluate(m, build_examples(split, cfg), cfg)
        assert em == 1.0 and tok == 1.0


def test_is_prime_probe_classified_p1():
    """On [101,200] the sieve errs exactly on {121,143,169,187} (classified prime) -> P1."""
    cfg = Config(task_mode="is_prime")
    r = probe_report(StubModel(_isprime_sieve), cfg)
    assert sorted(e["n"] for e in r["errors"]) == [121, 143, 169, 187], r["errors"]
    assert r["code"] == "P1", r["code"]


def test_train_frac_subsamples_train_only():
    cfg = Config(train_frac=0.5)
    tr, va = get_splits(cfg)
    assert len(va) == 30                      # val untouched (locked size)
    assert len(tr) == 34                      # round(69*0.5) = 34 (banker's rounding)
    assert not set(tr) & set(va)
    # deterministic across calls
    tr2, va2 = get_splits(Config(train_frac=0.5))
    assert tr == tr2 and va == va2
    # subsample is a subset of the full train split
    tr_full, va_full = get_splits(Config(train_frac=1.0))
    assert set(tr) <= set(tr_full)


def test_new_flags_parse():
    a = parse_args(["rnn", "0", "--task_mode", "is_prime", "--train_frac", "0.4", "--lr_decay", "True"])
    assert a.task_mode == "is_prime" and a.train_frac == 0.4 and a.lr_decay is True
    b = parse_args(["transformer", "0", "--lr_decay", "False"])
    assert b.lr_decay is False and b.task_mode == "next_prime" and b.train_frac == 1.0