"""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_p3():
    """On [101,200] the sieve errs on {121,143,169,187} (96% acc) -> P3 (is_prime has no sieve-rank)."""
    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"] == "P3", r["code"]
    assert r["acc"] >= 0.85


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_integers_mode_eval_does_not_crash_probe(tmp_path):
    """Regression: integers-mode eval crashed in the probe (out-of-vocab inputs >= 104).
    Train a tiny integers model, then run the REAL eval main against it — must not crash,
    and the probe must be N/A (Addendum 5: D2 is scored on in-range val EM only)."""
    import json
    import os
    import subprocess
    import sys
    import torch
    from src.data import build_examples, get_splits, make_batch
    from src.model_api import build_model

    cfg = Config(vocab_mode="integers")
    m = build_model(cfg)
    train_in, val_in = get_splits(cfg)
    train_ex = build_examples(train_in, cfg)
    opt = torch.optim.AdamW(m.parameters(), lr=1e-3)
    for _ in range(10):
        ex = train_ex[:16]
        b = make_batch(ex, cfg)
        out = m(b["x"], b["y_in"])
        loss = torch.nn.functional.cross_entropy(
            out["logits"].reshape(-1, cfg.vocab), b["y"].clamp(max=cfg.vocab - 1).reshape(-1), reduction="none")
        mask = b["y_mask"].float().reshape(-1)
        loss = (loss * mask).sum() / mask.sum()
        opt.zero_grad(); loss.backward(); opt.step()

    od = tmp_path / "rnn" / "seed0"
    os.makedirs(od, exist_ok=True)
    cfg.save(str(od / "config.json"))
    torch.save(m.state_dict(), str(od / "last.pt"))
    torch.save(m.state_dict(), str(od / "best.pt"))
    import csv
    with open(str(od / "metrics.csv"), "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["step", "train_loss", "train_token_acc", "train_em", "val_token_acc",
                    "val_em", "mean_halt_steps", "log_examples", "param_count"])
        w.writerow([200, 0.1, 1.0, 1.0, 1.0, 0.5, 3.0, "", 1000])

    r = subprocess.run([sys.executable, "-m", "src.eval", "rnn", "0", "--runs-dir", str(tmp_path)],
                       capture_output=True, text=True, cwd=os.path.join(os.path.dirname(__file__), ".."))
    assert r.returncode == 0, r.stderr[-600:]
    res = json.load(open(str(od / "results.json")))
    assert res["final"]["probe"]["code"] == "N/A"


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