"""Figures + summary numbers for the report from pulled metrics (results/{ice,dm}/<run>/metrics.jsonl).

Run names: [<prefix>-]<init>-<freeze>-lr<lr>-<budget>[-s<seed>]
"""

import argparse
import json
import re
from collections import defaultdict
from pathlib import Path

import matplotlib

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

COLORS = {"scratch": "#2a78d6", "ivgpt": "#eb6834", "shuffled": "#1baf7a", "rescaled": "#eda100"}  # validated, light mode
STYLES = {"scratch": "-", "ivgpt": "-", "shuffled": "--", "rescaled": "-."}
LABELS = {"scratch": "scratch", "ivgpt": "iVideoGPT init", "shuffled": "shuffled iVideoGPT", "rescaled": "rescaled iVideoGPT"}
NAME = re.compile(r"(?:(?P<prefix>[^-]+)-)?(?P<init>scratch|ivgpt|shuffled|rescaled|sp[0-9.]+(?:attn|mlp|lower|upper)?|spshuf[0-9.]+|nobody)-(?P<freeze>none|body|lora|top4)-lr(?P<lr>[0-9.e-]+)-(?P<budget>\d+M)(?:-s(?P<seed>\d+))?$")

plt.rcParams.update({
    "font.size": 8, "axes.titlesize": 8.5, "axes.labelsize": 8, "legend.fontsize": 7,
    "axes.spines.top": False, "axes.spines.right": False, "axes.edgecolor": "#52514e",
    "axes.labelcolor": "#0b0b0b", "xtick.color": "#52514e", "ytick.color": "#52514e",
    "axes.grid": True, "grid.color": "#e4e3df", "grid.linewidth": 0.5, "lines.linewidth": 1.4,
    "figure.dpi": 200, "savefig.bbox": "tight", "pdf.fonttype": 42,
})


def load_runs(roots):
    runs = []
    for root in roots:
        for d in sorted(Path(root).glob("*/")):
            m = NAME.match(d.name)
            f = d / "metrics.jsonl"
            if not m or not f.exists():
                continue
            recs = [json.loads(l) for l in f.read_text().splitlines() if l.strip()]
            val = np.array([(r["tokens"], r["val_loss"]) for r in recs if "val_loss" in r])
            train = [r for r in recs if "tok_per_s" in r]
            if len(val) < 2:
                continue
            g = m.groupdict()
            runs.append({
                "name": d.name, "src": Path(root).name, "init": g["init"], "freeze": g["freeze"],
                "lr": float(g["lr"]), "budget": g["budget"], "seed": int(g["seed"] or 0),
                "done": (d / "DONE").exists(), "tokens": val[:, 0], "val": val[:, 1],
                "tok_per_s": float(np.median([r["tok_per_s"] for r in train[2:]])) if len(train) > 3 else None,
                "hardware": (d / "hardware.txt").read_text().strip() if (d / "hardware.txt").exists() else d.parent.name,
                "evals": load_evals(d / "eval.json"),
            })
    return runs


EVAL_TASKS = {"piqa": "acc_norm", "hellaswag": "acc_norm", "arc_easy": "acc_norm", "sciq": "acc_norm",
              "winogrande": "acc", "lambada_openai": "acc"}


def load_evals(path):
    if not path.exists():
        return None
    res = json.loads(path.read_text())["results"]
    return {t: {"v": res[t][f"{m},none"], "se": res[t].get(f"{m}_stderr,none")} for t, m in EVAL_TASKS.items() if t in res}


def tokens_to_reach(run, target):
    """Tokens at which the (log-linearly interpolated) val curve first hits target, else None."""
    t, v = run["tokens"], run["val"]
    idx = np.where(v <= target)[0]
    if len(idx) == 0:
        return None
    i = idx[0]
    if i == 0:
        return float(t[0])
    lt0, lt1 = np.log(max(t[i - 1], 1.0)), np.log(t[i])
    frac = (v[i - 1] - target) / (v[i - 1] - v[i])
    return float(np.exp(lt0 + frac * (lt1 - lt0)))


def best_by_arm(runs, budget, freeze):
    """Best-LR finished run per init (seed 0) at a given budget/freeze."""
    best = {}
    for r in runs:
        if r["budget"] != budget or r["freeze"] != freeze or r["seed"] != 0 or not r["done"]:
            continue
        if r["init"] not in best or r["val"][-1] < best[r["init"]]["val"][-1]:
            best[r["init"]] = r
    return best


def direct_label(ax, runs_by_init, xmax=None):
    ends = sorted(((r["val"][-1], init) for init, r in runs_by_init.items()), reverse=True)
    lo, hi = ax.get_ylim()
    gap = 0.045 * (hi - lo)
    last_y = None
    for y, init in ends:  # top to bottom, nudge to avoid collisions
        ly = y if last_y is None else min(y, last_y - gap)
        r = runs_by_init[init]
        ax.annotate(LABELS[init], (r["tokens"][-1], y), xytext=(4, (ly - y) / (hi - lo) * 100), textcoords="offset points",
                    fontsize=6.5, color="#0b0b0b", va="center")
        last_y = ly


def fig_curves(runs, budget, freeze, out, title, ylim=None, logx=True):
    best = best_by_arm(runs, budget, freeze)
    if not best:
        return None
    fig, ax = plt.subplots(figsize=(3.3, 2.3))
    for init in ("scratch", "shuffled", "rescaled", "ivgpt"):
        if init in best:
            r = best[init]
            m = r["tokens"] > 0
            ax.plot(r["tokens"][m] / 1e6, r["val"][m], STYLES[init], color=COLORS[init],
                    label=f"{LABELS[init]} (lr {r['lr']:g})")
    if logx:
        ax.set_xscale("log")
    if ylim:
        ax.set_ylim(*ylim)
    ax.set_xlabel("text tokens (M)")
    ax.set_ylabel("validation loss (nats/token)")
    ax.set_title(title)
    ax.legend(frameon=False, loc="upper right")
    fig.savefig(out)
    plt.close(fig)
    return {k: {"lr": v["lr"], "final": float(v["val"][-1]), "name": v["name"]} for k, v in best.items()}


def fig_lr(runs, budget, freeze, out):
    """Final val loss vs peak LR, one line per arm (seed 0)."""
    table = defaultdict(dict)
    fig, ax = plt.subplots(figsize=(3.3, 2.3))
    for init in ("scratch", "shuffled", "rescaled", "ivgpt"):
        pts = sorted((r["lr"], float(r["val"][-1])) for r in runs
                     if r["budget"] == budget and r["freeze"] == freeze and r["init"] == init and r["done"] and r["seed"] == 0)
        for lr, v in pts:
            table[init][lr] = v
        if pts:
            x, y = zip(*pts)
            ax.plot(x, y, STYLES[init], marker="o", markersize=3.5, color=COLORS[init], label=LABELS[init])
    if not table:
        plt.close(fig)
        return None
    ax.set_xscale("log")
    ax.set_xlabel("peak learning rate")
    ax.set_ylabel("final validation loss")
    ax.set_title(f"{'full fine-tune' if freeze == 'none' else 'frozen body'}, {budget} tokens")
    ax.legend(frameon=False)
    fig.savefig(out)
    plt.close(fig)
    return dict(table)


def fig_multiplier(runs, budget, freeze, out, seeds=False):
    """Data multiplier: tokens scratch needs / tokens arm needs to reach the same loss."""
    groups = defaultdict(list)
    best = best_by_arm(runs, budget, freeze)
    if "scratch" not in best:
        return None
    for r in runs:
        if r["budget"] == budget and r["freeze"] == freeze and r["done"] and r["lr"] == best.get(r["init"], {}).get("lr"):
            groups[r["init"]].append(r)
    scratch_runs = groups["scratch"]
    lo = max(min(r["val"].min() for rs in groups.values() for r in rs), 0)
    hi = min(r["val"][1] for rs in groups.values() for r in rs)
    targets = np.linspace(hi, lo + 0.02, 40)
    fig, ax = plt.subplots(figsize=(3.3, 2.3))
    summary = {}
    for init in ("shuffled", "rescaled", "ivgpt"):
        if init not in groups:
            continue
        ratios = []
        for a in groups[init]:
            for s in scratch_runs:
                row = []
                for tg in targets:
                    ta, ts = tokens_to_reach(a, tg), tokens_to_reach(s, tg)
                    row.append(ts / ta if ta and ts else np.nan)
                ratios.append(row)
        ratios = np.array(ratios)
        med = np.nanmedian(ratios, axis=0)
        ax.plot(targets, med, STYLES[init], color=COLORS[init], label=LABELS[init])
        if len(ratios) > 1:
            ax.fill_between(targets, np.nanmin(ratios, 0), np.nanmax(ratios, 0), color=COLORS[init], alpha=0.15, lw=0)
        summary[init] = {f"{t:.2f}": float(m) for t, m in zip(targets[::8], med[::8])}
    ax.axhline(1.0, color="#52514e", lw=0.8)
    ax.invert_xaxis()
    ax.set_xlabel("target validation loss (lower = later in training)")
    ax.set_ylabel("data multiplier vs scratch")
    ax.set_title(f"tokens(scratch) / tokens(arm), {budget}")
    ax.legend(frameon=False)
    fig.savefig(out)
    plt.close(fig)
    return summary


def fig_delta(runs, budget, freeze, out):
    """Val loss minus best-LR scratch (seed 0) at matched eval steps; seeds at each arm's best LR, mean +/- range."""
    best = best_by_arm(runs, budget, freeze)
    if "scratch" not in best:
        return None
    ref = dict(zip(best["scratch"]["tokens"], best["scratch"]["val"]))
    fig, ax = plt.subplots(figsize=(3.3, 2.3))
    summary = {}
    for init in ("scratch", "shuffled", "rescaled", "ivgpt"):
        if init not in best:
            continue
        reps = [r for r in runs if r["init"] == init and r["budget"] == budget and r["freeze"] == freeze
                and r["lr"] == best[init]["lr"] and r["done"]]
        toks = sorted(set.intersection(*[set(r["tokens"]) for r in reps]) & set(ref))
        toks = [t for t in toks if t > 0]
        d = np.array([[dict(zip(r["tokens"], r["val"]))[t] - ref[t] for t in toks] for r in reps])
        x = np.array(toks) / 1e6
        ax.plot(x, d.mean(0), STYLES[init], color=COLORS[init], label=f"{LABELS[init]} (n={len(reps)})")
        if len(reps) > 1:
            ax.fill_between(x, d.min(0), d.max(0), color=COLORS[init], alpha=0.18, lw=0)
        summary[init] = {"n": len(reps), "final_mean": float(d[:, -1].mean()), "final_min": float(d[:, -1].min()),
                         "final_max": float(d[:, -1].max()), "lr": best[init]["lr"]}
    ax.axhline(0.0, color="#52514e", lw=0.8)
    ax.set_xscale("log")
    ax.set_xlabel("text tokens (M)")
    ax.set_ylabel("val loss minus scratch (nats/token)")
    ax.set_title(f"relative to best-LR scratch, {budget}")
    ax.legend(frameon=False, fontsize=6.5)
    fig.savefig(out)
    plt.close(fig)
    return summary


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--roots", nargs="+", default=["results/ice", "results/dm_backup"])
    p.add_argument("--out", default="report/figs")
    args = p.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    runs = load_runs(args.roots)
    s = {"runs": [{k: (v if not isinstance(v, np.ndarray) else None) for k, v in r.items()}
                  | {"final_val": float(r["val"][-1]), "final_tokens": float(r["tokens"][-1])} for r in runs]}
    budgets = sorted({r["budget"] for r in runs}, key=lambda b: int(b[:-1]))
    for b in budgets:
        s[f"best_full_{b}"] = fig_curves(runs, b, "none", out / f"curves_full_{b}.pdf", f"full fine-tune, best LR per arm ({b})")
        s[f"best_frozen_{b}"] = fig_curves(runs, b, "body", out / f"curves_frozen_{b}.pdf", f"frozen body, best LR per arm ({b})")
        s[f"lr_full_{b}"] = fig_lr(runs, b, "none", out / f"lr_full_{b}.pdf")
        s[f"lr_frozen_{b}"] = fig_lr(runs, b, "body", out / f"lr_frozen_{b}.pdf")
        s[f"mult_full_{b}"] = fig_multiplier(runs, b, "none", out / f"multiplier_full_{b}.pdf")
        s[f"mult_frozen_{b}"] = fig_multiplier(runs, b, "body", out / f"multiplier_frozen_{b}.pdf")
        s[f"delta_full_{b}"] = fig_delta(runs, b, "none", out / f"delta_full_{b}.pdf")
        s[f"delta_frozen_{b}"] = fig_delta(runs, b, "body", out / f"delta_frozen_{b}.pdf")
    (out.parent / "summary.json").write_text(json.dumps(s, indent=1, default=str))
    for r in sorted(runs, key=lambda r: r["name"]):
        print(f"{r['name']:45s} {'done' if r['done'] else 'run '} {r['tokens'][-1] / 1e6:>7.0f}M "
              f"val={r['val'][-1]:.3f} tok/s={r['tok_per_s']}")


if __name__ == "__main__":
    main()