"""Paper figures from the per-step training logs (ckpt/<arm>/log.jsonl, one JSON line per optimiser step).

  python plot.py runlogs/A.jsonl runlogs/B.jsonl runlogs/C.jsonl [runlogs/D.jsonl] --out figures

Figure 1: training loss vs tokens seen (smoothed over 50 steps, raw curve faint behind it), all arms.
Figure 2: validation loss vs tokens (every 200 steps), all arms.
Figure 3: the last 25% of training zoomed in, where the arms separate.
Also writes figures/curves.csv (step, tokens, loss per arm) for any plotting tool.
"""
import argparse, json, pathlib
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

NAMES = {"A": "A: learned table (268M trainable)", "B": "B: frozen table, float body (101M)",
         "C": "C: frozen table, ternary body (101M)", "D": "D: frozen random codes (101M)"}
COLORS = {"A": "#4E5A63", "B": "#3DE0C2", "C": "#FF5C7A", "D": "#C9A227"}


def load(p):
    rows = [json.loads(l) for l in open(p)]
    by_step = {r["step"]: r for r in rows}                     # a resumed run may repeat steps: keep the last
    return [by_step[s] for s in sorted(by_step)]


def smooth(y, k=50):
    if len(y) < k:
        return np.asarray(y)
    c = np.cumsum(np.insert(y, 0, 0.0))
    s = (c[k:] - c[:-k]) / k                                  # centred moving average: no lag against the raw curve
    h = k // 2
    return np.concatenate([np.full(h, np.nan), s, np.full(len(y) - len(s) - h, np.nan)])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("logs", nargs="+")
    ap.add_argument("--out", default="figures")
    a = ap.parse_args()
    out = pathlib.Path(a.out); out.mkdir(exist_ok=True)
    runs = {}
    for p in a.logs:
        arm = pathlib.Path(p).stem[0].upper()
        runs[arm] = load(p)

    fig, ax = plt.subplots(figsize=(8, 5))
    for arm, R in runs.items():
        t = np.array([r["tokens"] for r in R]) / 1e9; y = np.array([r["loss"] for r in R])
        ax.plot(t, y, color=COLORS[arm], alpha=0.15, lw=0.6)
        ax.plot(t, smooth(y), color=COLORS[arm], lw=1.8, label=NAMES[arm])
    ax.set_xlabel("training tokens (billions)"); ax.set_ylabel("training loss (nats/token)")
    ax.set_ylim(top=min(8, ax.get_ylim()[1])); ax.grid(alpha=0.3); ax.legend()
    fig.tight_layout(); fig.savefig(out / "fig1_train_loss.png", dpi=200); fig.savefig(out / "fig1_train_loss.pdf")

    fig, ax = plt.subplots(figsize=(8, 5))
    for arm, R in runs.items():
        V = [r for r in R if "val_loss" in r]
        ax.plot([r["tokens"] / 1e9 for r in V], [r["val_loss"] for r in V], "o-", ms=3, color=COLORS[arm], label=NAMES[arm])
    ax.set_xlabel("training tokens (billions)"); ax.set_ylabel("validation loss (nats/token)")
    ax.set_ylim(top=min(8, ax.get_ylim()[1])); ax.grid(alpha=0.3); ax.legend()
    fig.tight_layout(); fig.savefig(out / "fig2_val_loss.png", dpi=200); fig.savefig(out / "fig2_val_loss.pdf")

    fig, ax = plt.subplots(figsize=(8, 5))
    for arm, R in runs.items():
        n = len(R); R2 = R[int(n * 0.75):]
        t = np.array([r["tokens"] for r in R2]) / 1e9; y = smooth(np.array([r["loss"] for r in R2]))
        ax.plot(t, y, color=COLORS[arm], lw=1.8, label=NAMES[arm])
    ax.set_xlabel("training tokens (billions)"); ax.set_ylabel("training loss, last 25% (smoothed)")
    ax.grid(alpha=0.3); ax.legend()
    fig.tight_layout(); fig.savefig(out / "fig3_last_quarter.png", dpi=200); fig.savefig(out / "fig3_last_quarter.pdf")

    steps = sorted(set().union(*[{r["step"] for r in R} for R in runs.values()]))
    idx = {arm: {r["step"]: r for r in R} for arm, R in runs.items()}
    with open(out / "curves.csv", "w") as f:
        f.write("step,tokens," + ",".join(f"loss_{a},val_{a}" for a in runs) + "\n")
        for s in steps:
            tok = next((idx[a][s]["tokens"] for a in runs if s in idx[a]), "")
            cells = []
            for a_ in runs:
                r = idx[a_].get(s, {})
                cells += [str(r.get("loss", "")), str(r.get("val_loss", ""))]
            f.write(f"{s},{tok}," + ",".join(cells) + "\n")
    print("saved", sorted(p.name for p in out.iterdir()))


if __name__ == "__main__":
    main()