"""Post-run analysis: final metrics, [101,200] probe with sieve diagnostic, grokking signature.

Interpretation codes are locked in design/preregistration.md (with operationalization
thresholds in Addendum 2). This module MEASURES and classifies against those definitions.
Where the prereg prose supplies no testable boundary, the code emits the measurements
plus "unclassified" rather than inventing a label.
"""
import argparse
import csv
import json
import os

import numpy as np
import torch

from src.config import Config
from src.data import build_examples, decode_tokens, encode_int, get_splits, is_prime_n, next_prime, sieve_primes
from src.model_api import build_model, greedy_decode
from src.train import evaluate

# Composites with no prime factor <= 7 inside the probe's candidate window [102, 211].
# A model that learned only the {2,3,5,7} sieve predicts THESE as "next primes" (errors on
# n = 113..120, 139..142, 167..168, 181..186, 199..200 — 22 errors total). Includes 209 = 11*19,
# which the original {121,143,169,187} set (composites <= 200) missed — see Addendum 3.
FLAGGED_SIEVE_PREDS = {121, 143, 169, 187, 209}


def probe_report(model, cfg: Config, lo: int = 101, hi: int = 200) -> dict:
    try:
        device = next(model.parameters()).device
    except StopIteration:
        device = torch.device("cpu")
    margin = max(200, int(hi * 0.1) + 50)
    primes = sieve_primes(hi + margin)
    correct = 0
    errors = []
    easy_total = 0
    easy_wrong = 0
    is_prime_task = cfg.task_mode == "is_prime"
    eval_bs = getattr(cfg, "eval_batch_size", 512)
    inputs = list(range(lo, hi + 1))

    for bi in range(0, len(inputs), eval_bs):
        chunk = inputs[bi: bi + eval_bs]
        xs = [encode_int(n, cfg) for n in chunk]
        in_max = max(len(xi) for xi in xs)
        x_tensor = torch.full((len(chunk), in_max), cfg.pad_id, dtype=torch.long, device=device)
        for i, xi in enumerate(xs):
            x_tensor[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long, device=device)

        gen = greedy_decode(model, x_tensor, cfg).cpu()
        for i, n in enumerate(chunk):
            pred = decode_tokens(gen[i].tolist(), cfg)
            if is_prime_task:
                target = 1 if is_prime_n(n) else 0
                ok = (pred == 1) == (target == 1)      # any non-"1" output reads as "composite"
            else:
                target = next_prime(n, primes)
                ok = pred == target
            is_easy = (n % 2 == 0) or (n % 5 == 0)     # trivial composites (skip-evens / skip-5s)
            if is_easy:
                easy_total += 1
            if ok:
                correct += 1
            else:
                errors.append({"n": n, "target": target, "pred": pred})
                if is_easy:
                    easy_wrong += 1
    total = hi - lo + 1
    acc = correct / total
    if is_prime_task:
        flagged = [e for e in errors if e["n"] in FLAGGED_SIEVE_PREDS]   # input IS the classified number
        distinct_flagged = len({e["n"] for e in flagged})
    else:
        flagged = [e for e in errors if e["pred"] in FLAGGED_SIEVE_PREDS]
        distinct_flagged = len({e["pred"] for e in flagged})
    flagged_frac = len(flagged) / len(errors) if errors else 0.0
    # classification per prereg + Addendum 3/5 operationalization
    if easy_total and easy_wrong / easy_total > 0.5:
        code = "P4"      # fails trivial evens/5-multiples -> pure memorization
    elif is_prime_task and len(errors) >= 3 and distinct_flagged >= 3 and flagged_frac >= 0.8:
        code = "P1"      # is-prime: errors concentrated on no-small-factor composites -> learned sieve
    elif acc >= 0.85:
        code = "P3"      # surprising success beyond expectation (next_prime semantics)
    elif len(errors) >= 3 and distinct_flagged >= 3 and flagged_frac >= 0.8:
        code = "P1"      # next_prime: errors = sieves predicting no-small-factor composites
    else:
        code = "P2"      # scattered errors -> memorization / non-transferable heuristics
    return {
        "code": code, "acc": acc, "correct": correct, "total": total,
        "errors": errors, "flagged_errors": flagged,
        "flagged_pred_fraction": flagged_frac, "distinct_flagged_preds": distinct_flagged,
        "easy_total": easy_total, "easy_wrong": easy_wrong,
        "easy_err_rate": (easy_wrong / easy_total) if easy_total else None,
    }


def grokking_signature(metrics_path: str) -> dict:
    """Classify the training curve against preregistered codes O1-O4 (Addendum 2 operationalization)."""
    with open(metrics_path) as fh:
        rows = list(csv.DictReader(fh))
    if not rows:
        return {"code": "O4", "note": "no eval rows"}
    train = [float(r["train_em"]) for r in rows]
    val = [float(r["val_em"]) for r in rows]
    n = len(rows)
    # first index where train EM >= 0.95 for 10 CONSECUTIVE evals
    sat_start = next((i for i in range(n - 9) if all(t >= 0.95 for t in train[i:i + 10])), None)
    hi = next((i for i, v in enumerate(val) if v >= 0.9), None)
    # transition: last eval with val <= 0.2 strictly before hi, and after saturation window
    lo = None
    if hi is not None and sat_start is not None:
        lo_cands = [i for i in range(sat_start + 10, hi) if val[i] <= 0.2]
        if lo_cands:
            lo = max(lo_cands)
    width = (hi - lo) if (hi is not None and lo is not None) else None

    max_train = max(train)
    if max_train < 0.95:
        code = "O4"          # train never reached 0.95 -> setup/optimization failure
    elif sat_start is None:
        code = "O4"          # reached 0.95 but never sustained 10 evals within budget
    elif hi is not None and width is not None and width <= 5:
        code = "O1"          # sharp transition AFTER sustained train saturation
    elif hi is not None:
        code = "O3"          # reached 0.9+ but not via the sharp O1 pattern (gradual)
    elif val[-1] <= 0.3:
        code = "O2"          # train memorized, val stayed low
    else:
        code = "O-PARTIAL"   # val ended in (0.3, 0.9) with no 0.9 reach — prereg has no boundary
    return {
        "code": code, "sat_start_eval": sat_start, "val_hi_eval_idx": hi,
        "transition_width_evals": width, "n_evals": n,
        "final_train_em": train[-1], "final_val_em": val[-1],
        "max_train_em": max_train,
    }


@torch.no_grad()
def halting_report(model, cfg: Config) -> dict:
    """RNN halting structure: mean steps at run end + correlation with gap-to-next-prime (H1-H4)."""
    try:
        device = next(model.parameters()).device
    except StopIteration:
        device = torch.device("cpu")
    margin = max(100, int(cfg.range_end * 0.05) + 50)
    primes = sieve_primes(max(300, cfg.range_end + margin))
    gaps, steps = [], []
    inputs = list(range(cfg.range_start, cfg.range_end + 1))
    eval_bs = getattr(cfg, "eval_batch_size", 512)
    from src.data import pad_inputs
    for bi in range(0, len(inputs), eval_bs):
        chunk = inputs[bi: bi + eval_bs]
        xs = [encode_int(n, cfg) for n in chunk]
        in_max = max(len(xi) for xi in xs)
        x_tensor = torch.full((len(chunk), in_max), cfg.pad_id, dtype=torch.long, device=device)
        for i, xi in enumerate(xs):
            x_tensor[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long, device=device)
        x_padded = pad_inputs(x_tensor, cfg)
        h = model._encode(x_padded)
        _, s = model._run_compute(h)
        s_cpu = s.cpu().tolist()
        for i, n in enumerate(chunk):
            gaps.append(next_prime(n, primes) - n)
            steps.append(float(s_cpu[i]))
    mean = float(np.mean(steps))
    rho = float(np.corrcoef(gaps, steps)[0, 1]) if len(set(gaps)) > 1 else 0.0
    lo, hi = cfg.min_steps + 0.5, cfg.max_steps - 0.5
    if mean <= lo:
        code = "H1"          # collapse to the min_steps floor
    elif mean >= hi:
        code = "H2"          # pinned at K — never learned to halt
    elif rho >= 0.3:
        code = "H3"          # intermediate + positive gap correlation -> computation budget
    else:
        code = "H4"          # intermediate but no positive gap correlation — noisy/unused
    return {"code": code, "mean_steps": mean, "corr_gap": rho,
            "min_steps": cfg.min_steps, "max_steps": cfg.max_steps,
            "gaps": gaps, "steps_by_n": steps}


def _load(out_dir: str, ckpt: str, cfg: Config, device: torch.device | None = None):
    if device is None:
        device_str = getattr(cfg, "device", "auto")
        device = torch.device("cuda" if (device_str == "cuda" or (device_str == "auto" and torch.cuda.is_available())) else "cpu")
    model = build_model(cfg).to(device)
    model.load_state_dict(torch.load(os.path.join(out_dir, ckpt), map_location=device))
    model.eval()
    return model


def main() -> None:
    ap = argparse.ArgumentParser(description="prime-grokking eval")
    ap.add_argument("model")
    ap.add_argument("seed")
    ap.add_argument("--runs-dir", default="runs")
    a = ap.parse_args()
    out_dir = os.path.join(a.runs_dir, a.model, f"seed{a.seed}")
    cfg = Config.load(os.path.join(out_dir, "config.json"))

    _, val_in = get_splits(cfg)
    val_ex = build_examples(val_in, cfg)

    report = {"model": cfg.model, "seed": cfg.seed, "params": None, "val_selected": {}, "final": {}}
    for ckpt in ("best.pt", "last.pt"):
        path = os.path.join(out_dir, ckpt)
        if not os.path.exists(path):
            continue
        model = _load(out_dir, ckpt, cfg)
        tok, em, per = evaluate(model, val_ex, cfg)
        entry = {
            "ckpt": ckpt,
            "params": model.param_count(),
            "val_token_acc": tok,
            "val_exact_match": em,
            "note": ("val-selected checkpoint (early stop / best-on-val per prereg — "
                     "val EM here is selection-holed, not an untouched test estimate" if ckpt == "best.pt"
                     else "final checkpoint, no val selection"),
        }
        if ckpt == "last.pt":
            if cfg.vocab_mode == "integers":
                entry["probe"] = {
                    "code": "N/A",
                    "note": "atomic integer tokens make out-of-range inputs out-of-vocabulary by construction; "
                            "D2 is scored on in-range val EM only (Addendum 5)",
                }
            else:
                entry["probe"] = probe_report(model, cfg)
            if cfg.model == "rnn":
                entry["halting"] = halting_report(model, cfg)
            entry["per_example"] = [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok}
                                    for x, t, p, ok in per]
        if ckpt == "best.pt":
            report["val_selected"] = entry
        else:
            report["final"] = entry

    report["signature"] = grokking_signature(os.path.join(out_dir, "metrics.csv"))
    with open(os.path.join(out_dir, "results.json"), "w") as fh:
        json.dump(report, fh, indent=2)

    f = report["final"]
    p = f.get("probe", {})
    print(json.dumps({
        "val_em_best": round(report["val_selected"].get("val_exact_match", float("nan")), 4),
        "val_em_last": round(f.get("val_exact_match", float("nan")), 4),
        "probe_code": p.get("code"), "probe_acc": round(p.get("acc", float("nan")), 3),
        "flagged_errors": [e["n"] for e in p.get("flagged_errors", [])],
        "signature": report["signature"]["code"],
        "halting": (f.get("halting") or {}).get("code"),
        "results": os.path.join(out_dir, "results.json"),
    }, indent=2))


if __name__ == "__main__":
    main()