"""Action-conditioning experiment (small iVideoGPT): best-LR loss per starting point -> report/tables_act.tex.

Run names: act-<src>-<init>-none-lr<lr>-<budget>[-s<seed>], with budget 500M, 1000M, 2000M or 3000M and src in
  sm     small config, random init (scratch)
  base   thuml/ivideogpt-oxe-64-act-free (small pretrained model)
  bfree  thuml/ivideogpt-bair-64-act-free (base fine-tuned on BAIR without actions)
  bact   thuml/ivideogpt-bair-64-act-cond (base fine-tuned on BAIR with actions)
"""

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

import numpy as np

from report_figs import effective_tokens

NAME = re.compile(r"act-(?P<src>sm|base|bfree|bact)-(?P<init>scratch|ivgpt|shuffled|sp[0-9.]+)-none-lr(?P<lr>[0-9.e-]+)-(?P<budget>500M|1000M|2000M|3000M)(?:-s(?P<seed>\d+))?$")
ROWS = [("base", "ivgpt", "pretrained"), ("bfree", "ivgpt", "+ BAIR, no actions"),
        ("bact", "ivgpt", "+ BAIR, with actions"), ("base", "shuffled", "shuffled pretrained")]


def load(root, budget="500M"):
    runs = defaultdict(list)  # (src, init, lr) -> runs over seeds
    for d in sorted(Path(root).glob("act-*/")):
        m, f = NAME.match(d.name), d / "metrics.jsonl"
        if not m or m["budget"] != budget or not f.exists() or not (d / "DONE").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])
        runs[(m["src"], m["init"], float(m["lr"]))].append({"tokens": val[:, 0], "val": val[:, 1]})
    return runs


def best(runs, src, init):
    """(lr, mean final loss over seeds, runs) at the LR with the lowest mean final loss."""
    cands = [(lr, np.mean([r["val"][-1] for r in rs]), rs) for (s, i, lr), rs in runs.items() if s == src and i == init]
    return min(cands, key=lambda c: c[1]) if cands else None


def scratch_curve(root, lr=3e-3):
    """(budget, mean final loss, None) for scratch at one LR, as report_figs.effective_tokens expects."""
    pts = []
    for name, b in [("500M", 5e8), ("1000M", 1e9), ("2000M", 2e9), ("3000M", 3e9)]:
        rs = load(root, name).get(("sm", "scratch", lr))
        if rs:
            pts.append((b, float(np.mean([r["val"][-1] for r in rs])), None))
    return pts


def savings(pts, loss, budget):
    """Tokens scratch needs to match `loss`, as a multiple of `budget` (the metric of the medium-model study)."""
    t = effective_tokens(pts, loss)
    if t is not None:
        return rf"{t / budget:.2f}$\times$"
    return rf"$>${pts[-1][0] / budget:.2f}$\times$" if loss < pts[-1][1] else r"$<$1$\times$"


def main():
    root = Path(__file__).resolve().parents[1] / "results/ice"
    runs = load(root)
    for k, rs in sorted(runs.items()):
        print(k, [round(float(r["val"][-1]), 4) for r in rs])
    sc = best(runs, "sm", "scratch")
    pts = scratch_curve(root)
    print("scratch best", sc[:2], "curve", pts)
    lines = [r"\newcommand{\TableAct}{", r"\begin{tabular}{lcccc}", r"\toprule",
             r"& \multicolumn{2}{c}{direct transfer} & \multicolumn{2}{c}{S\&P ($\lambda=0.4$)} \\",
             r"\cmidrule(lr){2-3}\cmidrule(lr){4-5}",
             r"initialization & loss & token multiple & loss & token multiple \\", r"\midrule",
             rf"scratch & {sc[1]:.3f} & 1.00$\times$ & & \\"]
    for src, direct, label in ROWS:
        cells = []
        for init in (direct, "sp0.4" if direct == "ivgpt" else None):
            b = best(runs, src, init) if init else None
            if b is None:
                cells += ["", ""] if init is None else ["--", "--"]
                continue
            tok = savings(pts, float(b[1]), 5e8)
            cells += [f"{b[1]:.3f}", tok]
            print(src, init, "lr", b[0], round(float(b[1]), 4), tok, "seeds", len(b[2]),
                  [round(float(r["val"][-1]), 4) for r in b[2]])
        lines.append(f"{label} & " + " & ".join(cells) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}}"]
    long = load(root, "2000M")
    if long:
        lines += [r"\newcommand{\TableActLong}{", r"\begin{tabular}{lcccc}", r"\toprule",
                  r"& \multicolumn{2}{c}{after 500M tokens} & \multicolumn{2}{c}{after 2B tokens} \\",
                  r"\cmidrule(lr){2-3}\cmidrule(lr){4-5}",
                  r"initialization & loss & token multiple & loss & token multiple \\", r"\midrule"]
        for src, init, label in [("sm", "scratch", "scratch"), ("base", "shuffled", "shuffled pretrained"),
                                 ("base", "ivgpt", "pretrained")]:
            cells = []
            for rs, budget in [(runs, 5e8), (long, 2e9)]:
                b = best(rs, src, init)
                if b is None:
                    cells += ["--", "--"]
                    continue
                sav = r"1.00$\times$" if init == "scratch" else savings(pts, float(b[1]), budget)
                cells += [f"{b[1]:.3f}", sav]
                print(int(budget), src, init, "lr", b[0], round(float(b[1]), 4), sav, "seeds", len(b[2]),
                      [round(float(r["val"][-1]), 4) for r in b[2]])
            lines.append(f"{label} & " + " & ".join(cells) + r" \\")
        lines += [r"\bottomrule", r"\end{tabular}}"]
    for (s, i, lr), rs in sorted(runs.items()):
        if s == "base" and i.startswith("sp"):
            print("lambda", i[2:], "lr", lr, [round(float(r["val"][-1]), 4) for r in rs])
    out = root.parents[1] / "report/tables_act.tex"
    out.write_text("\n".join(lines) + "\n")
    print("wrote", out)


if __name__ == "__main__":
    main()