#!/usr/bin/env python
"""Build the self-contained HTML experiment report: report/index.html

Reads results.json + curves.png from both model run dirs, renders extra figures
(probe errors, halting scatter) inline, embeds everything as base64 — one file,
viewable on a phone over the tailnet.

Usage: .venv/bin/python -m scripts.build_report [--runs-dir DIR]
"""
import argparse
import base64
import io
import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from src.eval import FLAGGED_SIEVE_PREDS

CODE_LEGEND = {
    "O1": "Sharp transition — grokking-like",
    "O2": "Memorization won (train 100%, val low)",
    "O3": "Gradual rise — smooth heuristic, NOT grokking by our definition",
    "O4": "Train never saturated — setup/optimization failure, no interpretation",
    "O-PARTIAL": "Train saturated, val ended mid-range — partial in-range generalization",
    "H1": "Halt gate collapsed to the min_steps floor",
    "H2": "Halt gate pinned at K — never learned to halt",
    "H3": "Intermediate steps, positive gap correlation — learned compute budget",
    "H4": "Intermediate steps, no gap correlation — noisy/unused",
    "P1": "Errors = predictions of 121/143/169/187/209 — learned {2,3,5,7} sieve",
    "P2": "Scattered errors — memorization / non-transferable heuristics",
    "P3": "Probe accuracy ≥85% — surprising, needs seed verification",
    "P4": "Fails trivial evens/5-multiples — pure memorization",
}

INTERP = {
    "O1": "The tied cell + heavy weight decay found a structured in-range solution; memorization was repelled. NOT yet evidence of the full sieve — may be skip-evens/5s + divisibility heuristics.",
    "O2": "The lookup table is the lower-norm solution under these hyperparameters — consistent with the spec's 'memorization is a stronger attractor'.",
    "O3": "A 'fast' generalizing solution exists that gradient descent finds directly — still publishable, but not grokking.",
    "O4": "No scientific interpretation until the training setup is fixed.",
    "O-PARTIAL": "The model learned part of the mapping but never completed the in-range transition within budget.",
    "H1": "ACT failed — read architecture conclusions with a collapsed gate; the fixed-K ablation becomes the informative run.",
    "H2": "Never learned to halt — penalty too weak.",
    "H3": "Evidence of a learned computation budget: more compute steps for larger gaps.",
    "H4": "The halt signal is not being used meaningfully.",
    "P1": "Strongest positive result available at this scale: definitive evidence of a learned sieve with the training-range divisor set.",
    "P2": "No evidence of a divisibility-based algorithm transferring out-of-range.",
    "P3": "Richer algorithm than the {2,3,5,7} sieve — treat with suspicion, verify across seeds.",
    "P4": "The in-range solution didn't even transfer the trivial heuristics — strong memorization evidence.",
}


def b64(png_bytes: bytes) -> str:
    return "data:image/png;base64," + base64.b64encode(png_bytes).decode()


def img_b64(path: str) -> str:
    with open(path, "rb") as fh:
        return b64(fh.read())


def probe_fig(result: dict, model: str) -> str:
    errors = result.get("errors", [])
    flagged = [e for e in errors if e["pred"] in FLAGGED_SIEVE_PREDS]
    fig, ax = plt.subplots(figsize=(10, 3.2))
    xs, ys, cs, ls = [], [], [], []
    for e in errors:
        xs.append(e["n"]); ys.append(1.0)
        cs.append("tab:red" if e["pred"] in FLAGGED_SIEVE_PREDS else "tab:blue")
        ls.append(f"n={e['n']}\npred {e['pred']}\ntarget {e['target']}")
    ax.scatter(xs, ys, c=cs, s=28, zorder=3)
    ax.set_xlim(100, 201)
    ax.set_ylim(0.6, 1.4)
    ax.set_yticks([1.0]); ax.set_yticklabels(["wrong"])
    ax.set_xlabel("input n")
    ax.set_title(f"{model}: probe errors (red = predicted a no-small-divisor composite "
                 f"{sorted(FLAGGED_SIEVE_PREDS)})")
    for x, y, l in zip(xs, ys, ls):
        ax.annotate(l, (x, y), textcoords="offset points", xytext=(0, 10), fontsize=6,
                    ha="center", rotation=90, va="bottom")
    plt.tight_layout()
    buf = io.BytesIO(); plt.savefig(buf, dpi=110); plt.close(fig)
    return b64(buf.getvalue())


def halting_fig(result: dict, cfg: dict) -> str:
    gaps = result.get("gaps", [])
    steps = result.get("steps_by_n", [])
    n = list(range(cfg["range_start"], cfg["range_end"] + 1))
    fig, ax = plt.subplots(figsize=(8, 3))
    ax.scatter(gaps, steps, s=14)
    ax.set_xlabel("gap to next prime")
    ax.set_ylabel("mean halt steps")
    ax.set_title(f"RNN halting vs gap-to-next-prime (rho={result.get('corr_gap', 0):.3f})")
    plt.tight_layout()
    buf = io.BytesIO(); plt.savefig(buf, dpi=110); plt.close(fig)
    return b64(buf.getvalue())


def meta_table(run_dir: str) -> str:
    rows = []
    try:
        with open(os.path.join(run_dir, "run_meta.json")) as fh:
            m = json.load(fh)
        for k, v in m.items():
            rows.append(f"<tr><td>{k}</td><td>{v}</td></tr>")
    except FileNotFoundError:
        rows.append("<tr><td colspan=2>run_meta.json missing</td></tr>")
    return "".join(rows)


def code_badge(code: str) -> str:
    color = {"O1": "#2e7d32", "O2": "#c62828", "O3": "#f9a825", "O4": "#757575",
             "O-PARTIAL": "#f9a825", "H1": "#c62828", "H2": "#c62828", "H3": "#2e7d32",
             "H4": "#757575", "P1": "#2e7d32", "P2": "#c62828", "P3": "#f9a825",
             "P4": "#c62828"}.get(code, "#757575")
    return (f'<span class="badge" style="background:{color}">{code}</span> '
            f'<span class="dim">{CODE_LEGEND.get(code, "")}</span>')


def model_section(model: str, runs_dir: str) -> str:
    rd = os.path.join(runs_dir, model, "seed0")
    with open(os.path.join(rd, "results.json")) as fh:
        res = json.load(fh)
    with open(os.path.join(rd, "config.json")) as fh:
        cfg = json.load(fh)
    sig = res["signature"]
    final = res["final"]
    probe = final.get("probe", {})
    hlt = final.get("halting")
    curves = img_b64(os.path.join(rd, "curves.png"))
    parts = [f"<h2>{model} — seed 0</h2>"]
    parts.append(f'<p><b>Signature:</b> {code_badge(sig["code"])}</p>')
    if hlt:
        parts.append(f'<p><b>Halting:</b> {code_badge(hlt["code"])} '
                     f'(mean {hlt["mean_steps"]:.2f} steps, ρ(gap, steps) = {hlt["corr_gap"]:.3f})</p>')
    parts.append(f'<p><b>Probe [101,200]:</b> {code_badge(probe.get("code", "?"))} '
                 f'— accuracy {probe.get("acc", float("nan")):.3f} '
                 f'({probe.get("correct", 0)}/{probe.get("total", 0)}), '
                 f'{len(probe.get("errors", []))} errors</p>')
    parts.append(f'<p><b>Val exact-match</b> (last.pt, unselected): {final["val_exact_match"]:.3f} · '
                 f'(best.pt, val-selected): {res["val_selected"]["val_exact_match"]:.3f} '
                 f'<span class="dim">— best.pt is selection-holed by design</span></p>')
    parts.append(f'<p><b>Params:</b> {final["params"]:,} · vocab={cfg["vocab"]} · '
                 f'wd={cfg["weight_decay"]} · lr={cfg["lr"]} · K={cfg["max_steps"]} '
                 f'· halting={cfg["halting"]}</p>')
    parts.append(f'<img class="fig" src="{curves}" alt="curves">')
    if probe.get("errors"):
        parts.append(f'<img class="fig" src="{probe_fig(probe, model)}" alt="probe errors">')
    if hlt and "gaps" in hlt:
        parts.append(f'<img class="fig" src="{halting_fig(hlt, cfg)}" alt="halting">')
    parts.append(f"<pre class='dim'>{json.dumps({'signature': sig, 'probe': {k: v for k, v in probe.items() if k != 'errors'}, 'probe_errors': probe.get('errors', []), 'halting': {k: v for k, v in (hlt or {}).items() if k not in ('gaps', 'steps_by_n')}}, indent=2)}</pre>")
    return "".join(parts)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs-dir", default="runs-take2")
    ap.add_argument("--out", default="report/index.html")
    a = ap.parse_args()

    rnn_sec = model_section("rnn", a.runs_dir)
    tf_sec = model_section("transformer", a.runs_dir)
    meta_rnn = meta_table(os.path.join(a.runs_dir, "rnn", "seed0"))
    meta_tf = meta_table(os.path.join(a.runs_dir, "transformer", "seed0"))

    html = f"""<!DOCTYPE html>
<html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>prime-grokking — Experiment 1</title>
<style>
:root {{ color-scheme: dark; }}
body {{ font-family: -apple-system, system-ui, sans-serif; margin: 0 auto; max-width: 860px;
       padding: 16px; background: #0f1115; color: #d7dbe0; line-height: 1.55; }}
h1 {{ font-size: 1.5rem; }} h2 {{ font-size: 1.2rem; border-bottom: 1px solid #2a2f38; padding-bottom: 4px; }}
a {{ color: #7fb3ff; }}
.dim {{ color: #8b929c; font-size: .85rem; }}
.badge {{ display: inline-block; padding: 1px 8px; border-radius: 10px; color: #fff; font-weight: 600; font-size: .8rem; }}
.fig {{ width: 100%; height: auto; margin: 8px 0; border-radius: 8px; }}
table {{ border-collapse: collapse; width: 100%; font-size: .85rem; }}
td, th {{ border: 1px solid #2a2f38; padding: 4px 8px; text-align: left; }}
pre {{ overflow-x: auto; background: #15181d; padding: 10px; border-radius: 8px; font-size: .72rem; }}
blockquote {{ border-left: 3px solid #3a4150; margin-left: 0; padding-left: 12px; color: #a9b0ba; }}
</style></head><body>
<h1>prime-grokking — Experiment 1 (seed 0)</h1>
<p class="dim">Can a minimal architecture <i>grok</i> the next-prime function? Pre-registered experiment ·
2026-08-14 · runs on voidlaptop CPU · repo: ssh://meru/~/projects/prime-grokking.git</p>

<h2>Why this exists</h2>
<p>Grokking — the delayed jump from memorization to generalization — is one of the cleanest windows
into the <b>memorization → generalization transition</b> (Q3 of the research agenda). The canonical
grokking task is modular addition, which has smooth group structure (a low-complexity solution:
rotate on a circle). <b>Next-prime strips that away:</b> it looks smooth but is structurally a
combinatorial search (skip candidates, test divisibility, stop at the first hit). The hypothesis,
from the spec: a <b>weight-tied</b> recurrent cell should be able to represent the algorithmic
solution (one divisibility operation reused K times) with <i>fewer</i> effective parameters than a
lookup table — so under heavy weight decay, the algorithm should beat memorization. If it can't
happen here, in this maximally simple setting, that's evidence architectures need something
fundamentally different to cross the gap.</p>

<h2>What we did</h2>
<ul>
<li><b>Task:</b> n → next prime, n ∈ [2,100], 30% random holdout (seed 0), tokens = decimal digits + EOS.</li>
<li><b>Arm A (main):</b> weight-tied RNN — ONE 2-layer cell reused for input read-in, K=20 compute
steps, and output decoding; ACT learned halting gate with penalty ramp.</li>
<li><b>Arm B (baseline):</b> GPT-style transformer, d_model=128, 2 layers, 4 heads. Fixed d_model;
param counts logged, NOT matched (weight sharing is the studied variable).</li>
<li><b>Pre-registered:</b> the full outcome→interpretation matrix (codes O1–O4, H1–H4, P1–P4) was
committed to the research repo BEFORE any runs. Lock chain:
<code>00c696d → b3fe898 → 3fe3935 → 3015189</code>.</li>
<li><b>Multi-model review:</b> Gemini 3.6 Flash design review (removed the un-tied GRU decoder that
would have masked the weight-tying claim; added λ ramp) and an OpenAI Codex hostile code review
that caught two BLOCKERs before launch (batch-layout dependence of representations; misclassified
prereg codes) plus the off-by-one in the halting floor and a probe-diagnostic semantics error
(errors are flagged by <i>prediction</i>, not input — set extended to include 209=11·19).</li>
<li><b>Verification:</b> 34 tests green, including a classifier regression suite with ground-truth
stub models (a pure {{2,3,5,7}} sieve must classify P1 — it does).</li>
</ul>

<h2>Results</h2>
{rnn_sec}
{tf_sec}

<h2>Interpretation (locked before the runs)</h2>
<p>Every code below has a pre-registered meaning; the report only matches measurements against it.</p>
<table><tr><th>Code</th><th>Meaning (from preregistration.md)</th></tr>
{"".join(f"<tr><td><b>{k}</b></td><td>{v}</td></tr>" for k, v in INTERP.items())}
</table>
<blockquote><b>Seed-0 caveat:</b> everything here is one seed. Seeds {{1,2}} are required before any
claim. Val EM on best.pt is selection-holed (early stopping uses it); probe + halting use the
unselected last checkpoint. Training-range caveat: for n ≤ 100 the sieve only needs divisors
{{2,3,5,7}} — in-range generalization does not imply the general algorithm.</blockquote>

<h2>Next steps (jayrup's call)</h2>
<ul>
<li>Seeds 1, 2 (required before any claim)</li>
<li>wd sweep {{0.01, 0.1, 0.3, 1.0, 3.0}}</li>
<li>Ablations: halting=False (fixed-K), vocab_mode=integers, range [2,200]/[2,1000]</li>
<li>If O2: smaller train fractions (40%/50%) to sharpen the memorization attractor</li>
</ul>

<h2>Environment</h2>
<table><tr><th>rnn</th><th></th></tr>{meta_rnn}
<tr><th>transformer</th><th></th></tr>{meta_tf}</table>

<h2>Reproduce</h2>
<pre>git clone ssh://meru/~/projects/prime-grokking.git && cd prime-grokking
python3 -m venv .venv
.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu
.venv/bin/pip install numpy matplotlib pytest
.venv/bin/pytest tests/ -q
scripts/run_experiment.sh rnn 0
scripts/run_experiment.sh transformer 0
.venv/bin/python -m src.eval rnn 0; .venv/bin/python -m src.eval transformer 0
.venv/bin/python -m scripts.plot rnn 0; .venv/bin/python -m scripts.plot transformer 0</pre>

<p class="dim">Generated by scripts/build_report.py · design docs: design/experiment-spec.md,
design/preregistration.md, design/reviews/ · honest-limitations policy: results are compared
against the pre-registered matrix verbatim; nothing is interpreted post-hoc.</p>
</body></html>"""

    os.makedirs(os.path.dirname(a.out), exist_ok=True)
    with open(a.out, "w") as fh:
        fh.write(html)
    print(f"saved {a.out} ({len(html)//1024} KB)")


if __name__ == "__main__":
    main()