#!/usr/bin/env python3
"""Turn the logs from every stage into charts, once all three have finished.

    python3 scripts/plot_runs.py [--runs runs] [--out runs/graphs]

Reads two files per stage:

    runs/<stage>/metrics.jsonl   the cookbook's learning curves
    runs/<stage>/progress.jsonl  wall-clock and throughput (progress.py)

Stages that have not run yet are skipped with a note rather than an error, so
this is safe to run mid-pipeline -- you just get fewer panels.
"""

import argparse
import json
import os
import sys

STAGES = ("sft", "dpo", "rl")

# Candidate metric keys per panel, best first. The cookbook names the training
# loss differently per stage, so each panel takes the first key that exists.
LOSS_KEYS = ("train_mean_nll", "dpo_loss", "loss", "train_mean_bpb")
HOLDOUT_KEYS = ("test/nll", "test/bpb")
REWARD_KEYS = ("env/all/reward/total", "env/all/reward/mean", "reward/mean", "train_mean_reward")

# Palette roles, from the validated reference palette. Categorical slots are
# assigned in fixed order and never cycled; the diverging pair is blue<->red
# with a neutral gray midpoint. Text never wears a series colour.
C_SERIES_1 = "#2a78d6"   # blue   -- the default single series
C_SERIES_2 = "#008300"   # green  -- second series when two are genuinely needed
C_POS = "#2a78d6"        # diverging: improvement
C_NEG = "#e34948"        # diverging: regression
C_MID = "#b8b7b2"        # diverging midpoint / recessive rule
C_INK = "#0b0b0b"
C_INK_2 = "#52514e"
C_WARN = "#eb6834"       # status: marks the wasted region, always with a label

# The evaluator emits these per slice, per eval round.
SLICE_METRICS = ("parse_rate", "reward", "numbering", "dependency",
                 "n_claims_mae", "antecedent_gap")


def eval_rounds(records):
    """Records that carry evaluator output, in step order."""
    out = [r for r in records
           if any(k.startswith("overall/") for k in r)]
    return sorted(out, key=lambda r: r.get("step", 0))


def slice_names(records):
    names = set()
    for r in records:
        for k in r:
            if "/" in k and not k.startswith("overall/"):
                head, tail = k.rsplit("/", 1)
                if tail in SLICE_METRICS:
                    names.add(head)
    return sorted(names)


def read_jsonl(path):
    if not os.path.exists(path):
        return []
    out = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                out.append(json.loads(line))
            except json.JSONDecodeError:
                continue          # a run killed mid-write leaves a partial line
    return out


def series(records, key, step_key="step"):
    """(steps, values) for `key`, skipping records that lack it."""
    xs, ys = [], []
    for r in records:
        v = r.get(key)
        if isinstance(v, (int, float)) and v == v:
            xs.append(r.get(step_key, len(xs)))
            ys.append(float(v))
    return xs, ys


def first_series(records, keys):
    for k in keys:
        xs, ys = series(records, k)
        if ys:
            return k, xs, ys
    return None, [], []


def load(runs_dir):
    data = {}
    for stage in STAGES:
        d = os.path.join(runs_dir, stage)
        metrics = read_jsonl(os.path.join(d, "metrics.jsonl"))
        prog = read_jsonl(os.path.join(d, "progress.jsonl"))
        if metrics or prog:
            data[stage] = {"metrics": metrics, "progress": prog, "dir": d}
    return data


def plot_learning_curves(plt, data, out):
    stages = [s for s in STAGES if data.get(s)]
    if not stages:
        return None
    fig, axes = plt.subplots(1, len(stages), figsize=(5.5 * len(stages), 4),
                             squeeze=False)
    for ax, stage in zip(axes[0], stages):
        m = data[stage]["metrics"]
        # RL optimises reward, the other two minimise a loss -- but fall through
        # to the loss keys so a stage still charts if the primary key is absent.
        wanted = (REWARD_KEYS + LOSS_KEYS) if stage == "rl" else (LOSS_KEYS + REWARD_KEYS)
        key, xs, ys = first_series(m, wanted)
        if ys:
            ax.plot(xs, ys, lw=1.2, label=key)
        hkey, hxs, hys = first_series(m, HOLDOUT_KEYS)
        if hys:
            # Held-out on the same axes: the gap opening up is the overfitting
            # signal, and it only reads as a gap if both lines share a scale.
            ax.plot(hxs, hys, lw=1.6, marker="o", ms=3, label=hkey)
        ax.set_title("%s" % stage.upper())
        ax.set_xlabel("step")
        ax.grid(alpha=0.3)
        if ys or hys:
            ax.legend(fontsize=8)
        else:
            ax.text(0.5, 0.5, "no loss series", ha="center", transform=ax.transAxes)
    fig.suptitle("Learning curves")
    return save(fig, out, "01_learning_curves.png")


def plot_eval_quality(plt, data, out):
    """overall/* from the slice evaluator, across every stage on one timeline."""
    panels = ["overall/reward", "overall/parse_rate", "overall/antecedent_gap"]
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    any_data = False
    offset = 0
    for stage in STAGES:
        if stage not in data:
            continue
        m = data[stage]["metrics"]
        for ax, key in zip(axes, panels):
            xs, ys = series(m, key)
            if ys:
                any_data = True
                ax.plot([x + offset for x in xs], ys, marker="o", ms=4, label=stage)
        steps = [r.get("step", 0) for r in m]
        offset += (max(steps) + 1) if steps else 0
    for ax, key in zip(axes, panels):
        ax.set_title(key)
        ax.set_xlabel("cumulative step (SFT -> DPO -> RL)")
        ax.grid(alpha=0.3)
        if any_data:
            ax.legend(fontsize=8)
    if not any_data:
        axes[0].text(0.5, 0.5, "no slice-evaluator rounds logged yet",
                     ha="center", transform=axes[0].transAxes)
    fig.suptitle("Claim quality across the pipeline")
    return save(fig, out, "02_eval_quality.png")


def plot_slice_breakdown(plt, data, out):
    """Per-slice reward at the last eval of each stage -- where the model is weak."""
    rows = {}
    for stage in STAGES:
        if stage not in data:
            continue
        for r in data[stage]["metrics"]:
            slices = {k.rsplit("/", 1)[0]: v for k, v in r.items()
                      if k.endswith("/reward") and not k.startswith("overall/")
                      and isinstance(v, (int, float))}
            if slices:
                rows[stage] = slices        # keep overwriting -> ends up the last
    if not rows:
        return None
    names = sorted(next(iter(rows.values())).keys())
    fig, ax = plt.subplots(figsize=(max(9, 0.75 * len(names) * len(rows)), 4.5))
    width = 0.8 / len(rows)
    for i, (stage, slices) in enumerate(rows.items()):
        ax.bar([j + i * width for j in range(len(names))],
               [slices.get(n, 0.0) for n in names], width=width, label=stage.upper())
    ax.set_xticks([j + 0.4 - width / 2 for j in range(len(names))])
    ax.set_xticklabels(names, rotation=30, ha="right", fontsize=8)
    ax.set_ylabel("mean reward")
    ax.set_title("Final reward by validation slice")
    ax.grid(alpha=0.3, axis="y")
    ax.legend()
    return save(fig, out, "03_slice_breakdown.png")


def plot_timing(plt, data, out):
    fig, axes = plt.subplots(1, 2, figsize=(11, 4))
    ax_rate, ax_wall = axes

    for stage in STAGES:
        if stage not in data:
            continue
        p = data[stage]["progress"]
        xs, ys = series(p, "sec_per_step")
        if ys:
            ax_rate.plot(xs, ys, lw=1, label=stage)
        exs, eys = series(p, "elapsed_s")
        if eys:
            ax_wall.plot(exs, [y / 60.0 for y in eys], lw=1.2, label=stage)

    ax_rate.set_title("Seconds per step")
    ax_rate.set_xlabel("step")
    ax_wall.set_title("Wall clock (minutes)")
    ax_wall.set_xlabel("step")
    for ax in axes:
        ax.grid(alpha=0.3)
        if ax.get_legend_handles_labels()[0]:
            ax.legend(fontsize=8)
    fig.suptitle("Throughput and wall clock")
    return save(fig, out, "04_timing.png")


def plot_slice_trajectories(plt, data, out, metric="parse_rate"):
    """One small multiple per slice. 11 slices is past any categorical palette's
    capacity, so this facets instead of drawing 11 lines in one axes."""
    if "sft" not in data:
        return None
    rounds = eval_rounds(data["sft"]["metrics"])
    names = slice_names(data["sft"]["metrics"])
    if len(rounds) < 2 or not names:
        return None

    cols = 4
    rows = -(-len(names) // cols)
    fig, axes = plt.subplots(rows, cols, figsize=(3.1 * cols, 2.4 * rows),
                             sharex=True, sharey=True)
    axes = axes.ravel() if hasattr(axes, "ravel") else [axes]
    for ax, name in zip(axes, names):
        xs = [r.get("step", 0) for r in rounds if (name + "/" + metric) in r]
        ys = [r[name + "/" + metric] for r in rounds if (name + "/" + metric) in r]
        # A single series per panel: the panel title carries identity, so no legend.
        ax.plot(xs, ys, color=C_SERIES_1, linewidth=2, marker="o", markersize=5)
        if ys:
            # Direct-label the end point only, never every point.
            ax.annotate("%.2f" % ys[-1], (xs[-1], ys[-1]), textcoords="offset points",
                        xytext=(4, 2), fontsize=8, color=C_INK_2)
        ax.set_title(name.replace("domain_", "").replace("_", " "), fontsize=9, color=C_INK)
        ax.set_ylim(-0.05, 1.08)
        ax.grid(alpha=0.25)
        ax.tick_params(labelsize=8, colors=C_INK_2)
    for ax in axes[len(names):]:
        ax.set_visible(False)
    fig.suptitle("%s per validation slice (1.0 = every generation parsed)"
                 % metric.replace("_", " "), color=C_INK)
    fig.supxlabel("training step", fontsize=9, color=C_INK_2)
    return save(fig, out, "05_slice_trajectories.png")


def plot_convergence(plt, data, out):
    """Where the held-out loss stopped improving, and how much run remained after.

    This is the panel that answers 'should I have trained this long', which the
    aggregate learning curve does not. The second panel measures the run after
    convergence in wall-clock time.
    """
    if "sft" not in data:
        return None
    m = data["sft"]["metrics"]
    _, xs, ys = first_series(m, HOLDOUT_KEYS)
    if len(ys) < 3:
        return None

    # Convergence = first point within 2% of the best value achieved.
    best = min(ys)
    spread = max(ys) - best
    thresh = best + 0.02 * spread if spread else best
    conv_i = next((i for i, v in enumerate(ys) if v <= thresh), len(ys) - 1)
    conv_step = xs[conv_i]

    prog = data["sft"]["progress"]
    t_x = [r.get("step", 0) for r in prog if r.get("elapsed_s") is not None]
    t_y = [r["elapsed_s"] / 60.0 for r in prog if r.get("elapsed_s") is not None]
    total_min = t_y[-1] if t_y else 0.0
    min_at_conv = 0.0
    for sx, sy in zip(t_x, t_y):
        if sx <= conv_step:
            min_at_conv = sy

    fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.2))

    ax = axes[0]
    ax.plot(xs, ys, color=C_SERIES_1, linewidth=2, marker="o", markersize=6)
    if conv_i < len(xs) - 1:
        ax.axvspan(conv_step, xs[-1], color=C_WARN, alpha=0.10)
        ax.annotate("no further gain\nafter step %d" % conv_step,
                    xy=(conv_step + (xs[-1] - conv_step) * 0.45,
                        best + (max(ys) - best) * 0.55),
                    fontsize=9, color=C_WARN, ha="center")
    ax.axvline(conv_step, color=C_MID, linewidth=2, linestyle="--")
    ax.set_title("Held-out loss: converged at step %d of %d" % (conv_step, xs[-1]), color=C_INK)
    ax.set_xlabel("training step", color=C_INK_2)
    ax.set_ylabel("held-out NLL", color=C_INK_2)
    ax.grid(alpha=0.25)

    ax = axes[1]
    if t_y:
        ax.plot(t_x, t_y, color=C_SERIES_1, linewidth=2)
        ax.axvline(conv_step, color=C_MID, linewidth=2, linestyle="--")
        wasted = max(0.0, total_min - min_at_conv)
        ax.fill_between([x for x in t_x if x >= conv_step],
                        [min_at_conv] * sum(1 for x in t_x if x >= conv_step),
                        [y for x, y in zip(t_x, t_y) if x >= conv_step],
                        color=C_WARN, alpha=0.18)
        ax.annotate("%.0f min after convergence\n(of %.0f min total)" % (wasted, total_min),
                    xy=(0.42, 0.18), xycoords="axes fraction",
                    fontsize=10, color=C_WARN)
    ax.set_title("Cumulative wall clock", color=C_INK)
    ax.set_xlabel("training step", color=C_INK_2)
    ax.set_ylabel("minutes", color=C_INK_2)
    ax.grid(alpha=0.25)

    fig.suptitle("Did the run earn its length?", color=C_INK)
    return save(fig, out, "06_convergence.png")


def plot_slice_deltas(plt, data, out, metric="reward"):
    """Change from the first eval to the last, per slice.

    Change has polarity, so this is the diverging case: blue for improvement,
    red for regression, sorted so the worst regression reads first.
    """
    if "sft" not in data:
        return None
    rounds = eval_rounds(data["sft"]["metrics"])
    names = slice_names(data["sft"]["metrics"])
    if len(rounds) < 2 or not names:
        return None

    deltas = []
    for n in names:
        key = n + "/" + metric
        vals = [r[key] for r in rounds if key in r]
        if len(vals) >= 2:
            deltas.append((n, vals[0], vals[-1], vals[-1] - vals[0]))
    if not deltas:
        return None
    deltas.sort(key=lambda t: t[3])

    fig, ax = plt.subplots(figsize=(9.5, 0.42 * len(deltas) + 2.2))
    labels = [d[0].replace("domain_", "").replace("_", " ") for d in deltas]
    vals = [d[3] for d in deltas]
    colors = [C_NEG if v < 0 else C_POS for v in vals]
    ax.barh(labels, vals, color=colors, height=0.62)
    ax.axvline(0, color=C_MID, linewidth=2)
    for i, (n, a, b, d) in enumerate(deltas):
        ax.annotate("%.2f -> %.2f" % (a, b),
                    (d, i), textcoords="offset points",
                    xytext=(6 if d >= 0 else -6, 0), va="center",
                    ha="left" if d >= 0 else "right",
                    fontsize=8, color=C_INK_2)
    ax.set_title("Change in %s from first eval to last (blue = better, red = worse)"
                 % metric, color=C_INK)
    ax.set_xlabel("change", color=C_INK_2)
    ax.grid(alpha=0.25, axis="x")
    pad = max(abs(min(vals)), abs(max(vals))) * 0.35 + 0.02
    ax.set_xlim(min(vals) - pad, max(vals) + pad)
    return save(fig, out, "07_slice_deltas.png")


def plot_generation_health(plt, data, out):
    """Is the model still producing well-shaped claim sets?

    n_claims_mae is the early-warning signal for runaway generation: it moves
    long before the reward aggregate does, and a value far above the reference
    claim count means the model is not terminating.
    """
    if "sft" not in data:
        return None
    rounds = eval_rounds(data["sft"]["metrics"])
    names = slice_names(data["sft"]["metrics"])
    if len(rounds) < 2 or not names:
        return None

    fig, axes = plt.subplots(1, 2, figsize=(13, 4.3))

    ax = axes[0]
    xs, ys = [], []
    for r in rounds:
        vals = [r[n + "/n_claims_mae"] for n in names if (n + "/n_claims_mae") in r]
        if vals:
            xs.append(r.get("step", 0))
            ys.append(sum(vals) / len(vals))
    ax.plot(xs, ys, color=C_SERIES_1, linewidth=2, marker="o", markersize=6)
    ax.axhline(16, color=C_MID, linewidth=2, linestyle="--")
    ax.annotate("reference set is ~16 claims;\nan error this size means the\nmodel is not terminating",
                xy=(0.04, 0.62), xycoords="axes fraction", fontsize=9, color=C_INK_2)
    if ys:
        ax.annotate("%.1f" % ys[-1], (xs[-1], ys[-1]), textcoords="offset points",
                    xytext=(5, 0), fontsize=9, color=C_INK_2)
    ax.set_title("Claim-count error (mean absolute)", color=C_INK)
    ax.set_xlabel("training step", color=C_INK_2)
    ax.set_ylabel("claims off vs reference", color=C_INK_2)
    ax.grid(alpha=0.25)

    ax = axes[1]
    last = rounds[-1]
    pr = [(n, last.get(n + "/parse_rate")) for n in names if (n + "/parse_rate") in last]
    pr.sort(key=lambda t: t[1])
    lbl = [p[0].replace("domain_", "").replace("_", " ") for p in pr]
    val = [p[1] for p in pr]
    # Sequential magnitude: one hue, light->dark. Low parse rate reads lightest.
    ramp = ["#cde2fb", "#9ec5f4", "#6da7ec", "#3987e5", "#256abf", "#184f95"]
    cols = [ramp[min(len(ramp) - 1, int(v * len(ramp)))] for v in val]
    ax.barh(lbl, val, color=cols, height=0.62)
    for i, v in enumerate(val):
        ax.annotate("%.2f" % v, (v, i), textcoords="offset points", xytext=(5, 0),
                    va="center", fontsize=8, color=C_INK_2)
    ax.set_xlim(0, 1.12)
    ax.set_title("parse_rate by slice at the final eval", color=C_INK)
    ax.set_xlabel("fraction of generations that parsed", color=C_INK_2)
    ax.grid(alpha=0.25, axis="x")

    fig.suptitle("Generation health", color=C_INK)
    return save(fig, out, "08_generation_health.png")


def plot_rl_health(plt, data, out):
    """GRPO's own signals: did the training reward climb, and did the policy move?

    Three single-series panels (different scales, so never one axes):
      - env reward: the objective GRPO maximises. A jump then a plateau near the
        ceiling means the reward saturated and advantages went to ~0 -- there was
        little left for RL to optimise.
      - KL from the reference: how far the policy actually moved. Near-zero means
        the updates were tiny (the flip side of a saturated reward).
      - entropy: a collapse toward 0 would signal mode collapse / repetition.
    """
    if "rl" not in data:
        return None
    m = data["rl"]["metrics"]
    rx, reward = first_series(m, REWARD_KEYS)[1:]
    if len(reward) < 3:
        return None
    kx, kl = series(m, "optim/kl_sample_train_v1")
    ex, ent = series(m, "optim/entropy")

    panels = [("env reward (GRPO objective)", rx, reward, C_SERIES_1, None),
              ("KL from reference policy", kx, kl, C_SERIES_1, "near 0 = policy barely moved"),
              ("policy entropy", ex, ent, C_SERIES_1, "stable = no mode collapse")]
    panels = [(t, x, y, c, n) for t, x, y, c, n in panels if y]
    fig, axes = plt.subplots(1, len(panels), figsize=(4.4 * len(panels), 4))
    axes = axes if hasattr(axes, "__len__") else [axes]
    for ax, (title, xs, ys, col, note) in zip(axes, panels):
        ax.plot(xs, ys, color=col, linewidth=2, marker="o", markersize=5)
        if note:
            ax.annotate(note, xy=(0.05, 0.06), xycoords="axes fraction",
                        fontsize=9, color=C_INK_2)
        ax.set_title(title, color=C_INK)
        ax.set_xlabel("training step", color=C_INK_2)
        ax.grid(alpha=0.25)
    fig.suptitle("GRPO: reward saturated, so the policy moved little", color=C_INK)
    return save(fig, out, "10_rl_health.png")


def plot_dpo_health(plt, data, out):
    """DPO's own objective: does the chosen/rejected reward gap widen?

    accuracy (fraction of pairs the policy already prefers correctly) and the
    implicit chosen/rejected rewards are the signals that show whether DPO
    learned the preference at all -- the sampled-output panels cannot, because
    they measure form, not scope. chosen_reward and rejected_reward share a
    scale (both are beta * log-prob ratios), so they belong on ONE axes; the gap
    between them IS the margin. accuracy is a different scale, so it is its own
    panel rather than a second y-axis.
    """
    if "dpo" not in data:
        return None
    m = data["dpo"]["metrics"]
    ax_xs, acc = series(m, "accuracy")
    if len(acc) < 3:
        return None

    fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.2))

    ax = axes[0]
    ax.plot(ax_xs, acc, color=C_SERIES_1, linewidth=2)
    ax.axhline(0.5, color=C_MID, linewidth=1.5, linestyle="--")
    ax.annotate("0.5 = no preference learned", (ax_xs[len(ax_xs)//2], 0.5),
                textcoords="offset points", xytext=(0, 6), fontsize=8, color=C_INK_2)
    ax.set_ylim(0.45, 1.02)
    ax.set_title("Preference accuracy: %.2f -> %.2f" % (acc[0], acc[-1]), color=C_INK)
    ax.set_xlabel("training step", color=C_INK_2)
    ax.set_ylabel("fraction of pairs ranked correctly", color=C_INK_2)
    ax.grid(alpha=0.25)

    ax = axes[1]
    cx, cr = series(m, "chosen_reward")
    rx, rr = series(m, "rejected_reward")
    if cr and rr:
        # Two series -> a legend is present; both carry text-token labels, not
        # the series colour.
        ax.plot(cx, cr, color=C_SERIES_1, linewidth=2, label="chosen (granted)")
        ax.plot(rx, rr, color=C_NEG, linewidth=2, label="rejected (as-filed)")
        ax.fill_between(cx, cr, rr[:len(cr)], color=C_SERIES_1, alpha=0.08)
        ax.legend(fontsize=9, loc="upper right")
        ax.annotate("the gap is the margin;\nwider = stronger preference",
                    xy=(0.05, 0.4), xycoords="axes fraction", fontsize=9, color=C_INK_2)
    ax.set_title("Implicit rewards pull apart", color=C_INK)
    ax.set_xlabel("training step", color=C_INK_2)
    ax.set_ylabel("implicit reward (beta x logprob ratio)", color=C_INK_2)
    ax.grid(alpha=0.25)

    fig.suptitle("DPO learned the examiner preference", color=C_INK)
    return save(fig, out, "09_dpo_health.png")


def save(fig, out, name):
    fig.tight_layout()
    path = os.path.join(out, name)
    fig.savefig(path, dpi=130)
    import matplotlib.pyplot as plt
    plt.close(fig)
    return path


def verdict(data):
    """State what the numbers say, so the summary is readable without the charts.

    Deliberately mechanical: it reports movements and thresholds, and stops short
    of claiming the model is good or bad. Reward here measures claim FORM, not
    drafting quality -- an untuned base model that emits generic well-formed
    claims scores near 1.0 -- so a fall in reward means "worse at emitting clean
    parseable claim sets", not "worse at patent drafting".
    """
    if "sft" not in data:
        return []
    m = data["sft"]["metrics"]
    rounds = eval_rounds(m)
    lines = ["", "## Read of this run", ""]
    if len(rounds) < 2:
        return lines + ["Only %d eval round(s) -- not enough to read a trend." % len(rounds)]

    def overall(r, k):
        return r.get("overall/" + k)

    first, last = rounds[0], rounds[-1]
    notes = []

    # 1. Did held-out loss converge early?
    _, xs, ys = first_series(m, HOLDOUT_KEYS)
    if len(ys) >= 3:
        best = min(ys)
        spread = max(ys) - best
        thresh = best + 0.02 * spread if spread else best
        conv_i = next((i for i, v in enumerate(ys) if v <= thresh), len(ys) - 1)
        frac = xs[conv_i] / float(xs[-1]) if xs[-1] else 1.0
        if frac < 0.5:
            notes.append(
                "- **Held-out loss converged at step %d of %d** (%.0f%% of the run). "
                "The remaining steps did not improve it, so a shorter run would "
                "have reached the same place." % (xs[conv_i], xs[-1], frac * 100))

    # 2. Did sampled-generation quality move?
    for key, label in (("parse_rate", "parse rate"), ("reward", "reward")):
        a, b = overall(first, key), overall(last, key)
        if a is None or b is None:
            continue
        if b < a - 0.05:
            notes.append("- **overall/%s fell %.2f -> %.2f.** Sampled output got worse "
                         "even where the loss did not." % (key, a, b))
        elif b > a + 0.05:
            notes.append("- overall/%s rose %.2f -> %.2f." % (key, a, b))

    # 3. Runaway generation?
    names = slice_names(m)
    maes = [last[n + "/n_claims_mae"] for n in names if (n + "/n_claims_mae") in last]
    if maes:
        mae = sum(maes) / len(maes)
        if mae > 16:
            notes.append(
                "- **Claim-count error averages %.1f against a ~16-claim reference.** "
                "That is the signature of a model that is not terminating -- most "
                "likely repetition under greedy decoding, since eval samples at "
                "temperature 0. Check actual generations before concluding the "
                "model regressed." % mae)

    # 4. Baseline comparison -- was the first eval already good?
    a = overall(first, "reward")
    if a is not None and a > 0.95:
        notes.append(
            "- The **first eval already scored %.2f**, before training had "
            "meaningfully changed the model. This metric cannot see what SFT "
            "added; it measures form, and the base model already had the form." % a)

    # 5. Worst slices at the end
    prs = sorted(((n, last[n + "/parse_rate"]) for n in names if (n + "/parse_rate") in last),
                 key=lambda t: t[1])[:3]
    if prs and prs[0][1] < 0.8:
        notes.append("- Weakest slices at the final eval: %s."
                     % ", ".join("%s %.2f" % (n.replace("domain_", ""), v) for n, v in prs))

    if "dpo" in data:
        dm = data["dpo"]["metrics"]
        _, acc = series(dm, "accuracy")
        _, marg = series(dm, "margin")
        drounds = eval_rounds(dm)
        if acc:
            notes.append("- **DPO preference accuracy reached %.2f** (margin %.1f -> %.1f). "
                         "By its own objective the preference was learned cleanly."
                         % (acc[-1], marg[0] if marg else 0, marg[-1] if marg else 0))
        if len(drounds) >= 2:
            a = drounds[0].get("overall/parse_rate")
            b = drounds[-1].get("overall/parse_rate")
            if a is not None and b is not None:
                verb = "barely moved" if abs(b - a) < 0.05 else ("rose" if b > a else "fell")
                notes.append("- On sampled output, DPO's parse_rate %s (%.2f -> %.2f): the "
                             "preference signal did not fix the decoding-time parse failures, "
                             "which are a separate problem." % (verb, a, b))
            # per-domain winners/losers under DPO
            dom_deltas = []
            for n in slice_names(dm):
                key = n + "/reward"
                vals = [r[key] for r in drounds if key in r]
                if len(vals) >= 2:
                    dom_deltas.append((n.replace("domain_", ""), vals[-1] - vals[0]))
            dom_deltas.sort(key=lambda t: t[1])
            if dom_deltas:
                worst = dom_deltas[0]; best = dom_deltas[-1]
                notes.append("- DPO was uneven across domains: %s %+.2f (worst), "
                             "%s %+.2f (best)." % (worst[0], worst[1], best[0], best[1]))

    if "rl" in data:
        rm = data["rl"]["metrics"]
        rx, reward = first_series(rm, REWARD_KEYS)[1:]
        _, kl = series(rm, "optim/kl_sample_train_v1")
        rr = eval_rounds(rm)
        if len(reward) >= 3:
            jump = reward[1] - reward[0] if len(reward) > 1 else 0
            notes.append("- **GRPO reward jumped %.2f -> %.2f in the first step, then held ~%.2f.** "
                         "The reward saturated immediately, so group advantages went to ~0 and "
                         "there was little left for RL to optimise -- as expected when SFT+DPO "
                         "already scored near the ceiling." % (reward[0], reward[0] + jump, reward[-1]))
        if kl:
            notes.append("- KL from the reference stayed ~%.4f: the policy barely moved, "
                         "the flip side of a saturated reward." % (sum(kl) / len(kl)))
        if len(rr) >= 2:
            a = rr[0].get("overall/parse_rate"); b = rr[-1].get("overall/parse_rate")
            if a is not None and b is not None:
                notes.append("- On held-out slices during RL, parse_rate went %.2f -> %.2f "
                             "(steps %s-%s)." % (a, b, rr[0].get("step"), rr[-1].get("step")))

    return lines + (notes or ["Nothing anomalous in the logged metrics."])


def summarise(data):
    lines = ["# Run summary", ""]
    total_time = 0.0
    lines.append("| stage | steps | wall clock | eval rounds |")
    lines.append("|---|---|---|---|")
    for stage in STAGES:
        if stage not in data:
            lines.append("| %s | _not run_ | | |" % stage.upper())
            continue
        p = data[stage]["progress"]
        if not p:
            lines.append("| %s | ? | ? | ? |" % stage.upper())
            continue
        last = p[-1]
        secs = last.get("elapsed_s") or 0.0
        total_time += secs
        lines.append("| %s | %d | %dh%02dm | %d |"
                     % (stage.upper(), last.get("step", 0) + 1,
                        int(secs // 3600), int(secs % 3600 // 60),
                        sum(1 for r in p if r.get("is_eval"))))
    lines += ["", "**Total wall clock** %dh%02dm"
              % (int(total_time // 3600), int(total_time % 3600 // 60))]
    lines += verdict(data)
    return "\n".join(lines) + "\n"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", default="runs")
    ap.add_argument("--out", default="runs/graphs")
    args = ap.parse_args()

    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        sys.exit("matplotlib is not installed. Run: uv pip install -r requirements.txt")

    data = load(args.runs)
    if not data:
        sys.exit("No logs under %s/. Run a training stage first." % args.runs)

    missing = [s for s in STAGES if s not in data]
    if missing:
        print("note: no logs for %s -- plotting what exists" % ", ".join(missing))

    os.makedirs(args.out, exist_ok=True)
    written = [
        plot_learning_curves(plt, data, args.out),
        plot_eval_quality(plt, data, args.out),
        plot_slice_breakdown(plt, data, args.out),
        plot_timing(plt, data, args.out),
        plot_slice_trajectories(plt, data, args.out),
        plot_convergence(plt, data, args.out),
        plot_slice_deltas(plt, data, args.out),
        plot_generation_health(plt, data, args.out),
        plot_dpo_health(plt, data, args.out),
        plot_rl_health(plt, data, args.out),
    ]

    summary_path = os.path.join(args.out, "summary.md")
    with open(summary_path, "w") as f:
        f.write(summarise(data))

    print("\nWrote %d charts to %s/" % (sum(1 for w in written if w), args.out))
    for w in written:
        if w:
            print("  " + w)
    print("  " + summary_path)
    print()
    print(summarise(data))


if __name__ == "__main__":
    main()