"""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, 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:
    primes = sieve_primes(hi + 200)
    correct = 0
    errors = []
    easy_total = 0
    easy_wrong = 0
    for n in range(lo, hi + 1):
        x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
        gen = greedy_decode(model, x, cfg)[0].tolist()
        pred = decode_tokens(gen, cfg)
        target = next_prime(n, primes)
        is_easy = (n % 2 == 0) or (n % 5 == 0)   # trivial composites (skip-evens / skip-5s)
        if is_easy:
            easy_total += 1
        if pred == target:
            correct += 1
        else:
            errors.append({"n": n, "target": target, "pred": pred})
            if is_easy:
                easy_wrong += 1
    total = hi - lo + 1
    acc = correct / total
    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 operationalization; P4 checked first
    if easy_total and easy_wrong / easy_total > 0.5:
        code = "P4"      # fails trivial evens/5-multiples -> pure memorization
    elif acc >= 0.85:
        code = "P3"      # surprising success beyond expectation
    elif len(errors) >= 3 and distinct_flagged >= 3 and flagged_frac >= 0.8:
        code = "P1"      # errors = sieves predicting no-small-factor composites -> learned {2,3,5,7} sieve
    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)."""
    primes = sieve_primes(300)
    gaps, steps = [], []
    for n in range(cfg.range_start, cfg.range_end + 1):
        x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
        h = model._encode(x)
        _, s = model._run_compute(h)
        gaps.append(next_prime(n, primes) - n)
        steps.append(float(s.mean()))
    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):
    model = build_model(cfg)
    model.load_state_dict(torch.load(os.path.join(out_dir, ckpt), map_location="cpu"))
    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":
            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()