"""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


def _sieve_rank_signature(k: int, lo: int, hi: int) -> set[int]:
    """Composites in [lo, hi] with all prime factors > p_k.
    A k-rank sieve (checks divisibility by first k primes) misses exactly these.
    Smallest composite identifies k: 1147→k=10, 1369→k=11, 1681→k=12, 1849→k=13,
    none→k≥14 (Addendum 6).
    """
    primes = sieve_primes(hi + 100)
    # k-prime sieve checks primes[0..k-1] = {2,3,...,p_k}; misses factors > p_k
    pk = primes[k - 1]  # 0-indexed: primes[0]=2, so k=4 → pk=primes[3]=7
    # small_primes: all primes ≤ pk (the ones the sieve checks)
    small_primes = [p for p in primes if p <= pk]
    out = set()
    for n in range(lo, hi + 1):
        if n < 2:
            continue
        # check if n is composite (divisible by any prime)
        is_comp = False
        for p in primes:
            if p * p > n:
                break
            if n % p == 0:
                is_comp = True
                break
        if not is_comp:
            continue
        # n is composite; check if ALL small_primes fail to divide it
        # (i.e. all prime factors are > pk)
        all_factors_large = True
        for p in small_primes:
            if n % p == 0:
                all_factors_large = False
                break
        if all_factors_large:
            out.add(n)
    return out


def _compute_sieve_rank(preds: list[int], lo: int, hi: int) -> int | None:
    """Find the most specific sieve rank that explains the model's error predictions.
    Returns k (the number of primes in the sieve, 1-indexed) or None.
    A rank-k sieve checks {2,3,...,p_k} and misses composites with all factors > p_k.
    We want the LARGEST k whose signature set contains the model's error predictions.
    """
    primes = sieve_primes(hi + 100)
    composites_in_range = set()
    for n in range(lo, hi + 1):
        if n < 2:
            continue
        is_comp = False
        for p in primes:
            if p * p > n:
                break
            if n % p == 0:
                is_comp = True
                break
        if is_comp:
            composites_in_range.add(n)
    pred_composites = sorted(composites_in_range & set(preds))
    if not pred_composites:
        return None  # no composites predicted — either exact or garbage
    # find the LARGEST k whose rank-k signature contains all error predictions
    for k in range(40, 0, -1):
        if k >= len(primes):
            continue
        sig = _sieve_rank_signature(k, lo, hi)
        if all(p in sig for p in pred_composites):
            return k
    return None


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
    # sieve rank (P5/P6 ladder, Addendum 6)
    error_preds = [e["pred"] for e in errors]
    rank = _compute_sieve_rank(error_preds, lo, hi) if not is_prime_task else None
    # P-ladder classification (prereg + Addenda 3/5/6)
    if not errors:
        code = "P6"      # exact — no probe misses
    elif easy_total and easy_wrong / easy_total > 0.5:
        code = "P4"      # fails trivial evens/5-multiples -> pure memorization
    elif rank is not None:
        code = f"P5({rank})"  # errors match rank-k sieve signature
    elif acc >= 0.85:
        code = "P3"      # surprising success beyond expectation
    else:
        code = "P2"      # scattered errors -> memorization / non-transferable heuristics
    return {
        "code": code, "acc": acc, "correct": correct, "total": total,
        "errors": errors, "sieve_rank": rank,
        "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)
    state = torch.load(os.path.join(out_dir, ckpt), map_location=device)
    # torch.compile wraps modules (model._orig_mod), so checkpoints saved from a
    # compiled model carry an "_orig_mod." key prefix; DataParallel adds "module.".
    # Strip both so the uncompiled eval build loads cleanly.
    clean = {}
    for k, v in state.items():
        for pfx in ("_orig_mod.", "module."):
            while k.startswith(pfx):
                k = k[len(pfx):]
        clean[k] = v
    model.load_state_dict(clean)
    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:
                # out-of-range probe: [range_end+1, range_end+1000] (E6+ uses wider window)
                p_lo = cfg.range_end + 1
                p_hi = cfg.range_end + 1000
                entry["probe"] = probe_report(model, cfg, lo=p_lo, hi=p_hi)
            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),
        "sieve_rank": p.get("sieve_rank"),
        "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()