"""Classification correctness for src/eval.py — the preregistered O/H/P codes.

Ground truth is synthetic: stub models with KNOWN behavior feed the classifier,
and synthetic metrics CSVs feed grokking_signature. (Original: ad-hoc verifier
/tmp/hermes-verify-eval.py, promoted to a permanent regression suite.)
"""
import csv
import os
import tempfile

import torch

from src.config import Config
from src.data import sieve_primes
from src.eval import FLAGGED_SIEVE_PREDS, grokking_signature, halting_report, probe_report

DIGITS_CFG = Config(vocab_mode="digits")
EOS = DIGITS_CFG.eos_id
PAD = DIGITS_CFG.pad_id


def _decode_row(row) -> int:
    toks = [int(t) for t in row.tolist() if int(t) != PAD]
    return int("".join(map(str, toks))) if toks else -1


class StubModel(torch.nn.Module):
    """One-hot logits forcing greedy_decode to emit exactly pred_fn(n)."""

    def __init__(self, pred_fn):
        super().__init__()
        self.pred_fn = pred_fn

    def forward(self, x, y_in):
        B, T_out = x.shape[0], y_in.shape[1]
        logits = torch.full((B, T_out, 11), -100.0)
        for i in range(B):
            n = _decode_row(x[i])
            toks = [int(d) for d in str(self.pred_fn(n))] + [EOS]
            for t in range(T_out):
                tok = toks[t] if t < len(toks) else EOS
                logits[i, t, tok] = 100.0
        return {"logits": logits, "halt_steps": None}


def _sieve35(n):
    c = n + 1
    while any(c % d == 0 for d in (2, 3, 5, 7)):
        c += 1
    return c


def _perfect(n):
    return next(p for p in sieve_primes(500) if p > n)


def _easy_only(n):
    return _perfect(n) if (n % 2 == 0 or n % 5 == 0) else 199


def test_probe_sieve35_classified_p1():
    """A pure {2,3,5,7} sieve must classify P1: errors are predictions of 121/143/169/187/209."""
    r = probe_report(StubModel(_sieve35), DIGITS_CFG)
    preds = sorted({e["pred"] for e in r["errors"]})
    assert preds == sorted(FLAGGED_SIEVE_PREDS), r["errors"]
    assert len(r["errors"]) == 22  # n in 113..120, 139..142, 167..168, 181..186, 199..200
    assert r["code"] == "P1", r


def test_probe_perfect_classified_p3():
    r = probe_report(StubModel(_perfect), DIGITS_CFG)
    assert r["acc"] == 1.0 and r["code"] == "P3", r


def test_probe_easy_only_classified_p2():
    r = probe_report(StubModel(_easy_only), DIGITS_CFG)
    assert r["code"] == "P2", r


def test_probe_always_wrong_classified_p4():
    r = probe_report(StubModel(lambda n: n), DIGITS_CFG)
    assert r["code"] == "P4", r


def _write_csv(path, train_ems, val_ems):
    with open(path, "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"])
        for i, (t, v) in enumerate(zip(train_ems, val_ems)):
            w.writerow([i * 200, 0.1, 1.0, t, 1.0, v, 10.0, "", 1000])


def _sig_of(tmp, tr, va):
    p = os.path.join(tmp, "m.csv")
    _write_csv(p, tr, va)
    return grokking_signature(p)


def test_signature_o1_sharp_transition(tmp_path):
    tr = [0.99] * 20
    va = [0.1] * 15 + [0.2, 0.95, 0.99, 1.0, 1.0]
    r = _sig_of(str(tmp_path), tr, va)
    assert r["code"] == "O1", r


def test_signature_o2_memorization(tmp_path):
    assert _sig_of(str(tmp_path), [1.0] * 20, [0.1] * 20)["code"] == "O2"


def test_signature_o3_gradual(tmp_path):
    tr = [1.0] * 20
    va = [0.3, 0.45, 0.6, 0.75, 0.9, 0.95] + [1.0] * 14
    assert _sig_of(str(tmp_path), tr, va)["code"] == "O3"


def test_signature_o4_train_fails(tmp_path):
    assert _sig_of(str(tmp_path), [0.6] * 20, [0.2] * 20)["code"] == "O4"


def test_signature_opartial(tmp_path):
    assert _sig_of(str(tmp_path), [1.0] * 20, [0.6] * 20)["code"] == "O-PARTIAL"


def test_halting_h1_collapse_and_h2_pin():
    from src.model_api import build_model
    cfg = Config(model="rnn")
    m = build_model(cfg)
    with torch.no_grad():
        m.halt_head.bias.fill_(50.0)
    assert halting_report(m, cfg)["code"] == "H1"
    m2 = build_model(cfg)
    with torch.no_grad():
        m2.halt_head.bias.fill_(-50.0)
    assert halting_report(m2, cfg)["code"] == "H2"