"""LaTeX tables for the report from pulled metrics -> report/tables.tex (one \\newcommand per table)."""

import json
import sys
from collections import defaultdict
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
from analyze import EVAL_TASKS, LABELS, load_runs  # noqa: E402

ARMS = ("scratch", "ivgpt", "shuffled", "rescaled", "sp0.1", "sp0.2", "sp0.4", "sp0.7", "spshuf0.4")
NICE = {"scratch": "scratch", "ivgpt": "pretrained", "shuffled": "shuffled", "rescaled": "rescaled",
        "sp0.1": r"S\&P $\lambda{=}0.1$", "sp0.2": r"S\&P $\lambda{=}0.2$", "sp0.4": r"S\&P $\lambda{=}0.4$",
        "sp0.7": r"S\&P $\lambda{=}0.7$", "spshuf0.4": r"shuffled S\&P $\lambda{=}0.4$"}


def lr_tex(lr):
    """Learning rates in the plain e-notation common in ML papers, e.g. 3e-3."""
    e = int(np.floor(np.log10(lr) + 1e-9))
    return f"{lr / 10 ** e:g}e-{abs(e)}"


def sweep_table(runs, budget, freeze):
    """Rows = arms, cols = LRs, final val loss of seed-0 runs; the LR used for all further runs in bold."""
    grid = defaultdict(dict)
    for r in runs:
        if r["budget"] == budget and r["freeze"] == freeze and r["done"] and r["seed"] == 0:
            grid[r["init"]][r["lr"]] = r["val"][-1]
    lrs = sorted({lr for g in grid.values() for lr in g})
    if not lrs:
        return ""
    lines = [r"\begin{tabular}{l" + "c" * len(lrs) + "}", r"\toprule",
             "initialization & " + " & ".join(lr_tex(lr) for lr in lrs) + r" \\", r"\midrule"]
    for a in ARMS:
        if a not in grid:
            continue
        chosen = chosen_lr(runs, budget, freeze, a)
        cells = []
        for lr in lrs:
            v = grid[a].get(lr)
            cells.append("--" if v is None else (rf"\textbf{{{v:.3f}}}" if lr == chosen else f"{v:.3f}"))
        lines.append(NICE[a] + " & " + " & ".join(cells) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    return "\n".join(lines)


def chosen_lr(runs, budget, freeze, init):
    """The LR with the most seeds, then the lowest mean final loss over seeds (same rule as replicate_table)."""
    groups = defaultdict(list)
    for r in runs:
        if r["budget"] == budget and r["freeze"] == freeze and r["init"] == init and r["done"]:
            groups[r["lr"]].append(r["val"][-1])
    return min(groups, key=lambda l: (-len(groups[l]), np.mean(groups[l]))) if groups else None


def best_lr(runs, budget, freeze, init):
    c = [(r["val"][-1], r["lr"]) for r in runs if r["budget"] == budget and r["freeze"] == freeze
         and r["init"] == init and r["done"] and r["seed"] == 0]
    return min(c)[1] if c else None


def replicate_table(runs, budget):
    """Best-LR runs with all seeds: per-seed finals, mean, and the gap to scratch."""
    rows, ref = [], None
    for a in ARMS:
        groups = defaultdict(list)
        for r in runs:
            if r["budget"] == budget and r["freeze"] == "none" and r["init"] == a and r["done"]:
                groups[r["lr"]].append((r["seed"], r["val"][-1]))
        if not groups:
            continue
        # prefer the LR with the most seeds, then the lowest mean
        lr = min(groups, key=lambda l: (-len(groups[l]), np.mean([v for _, v in groups[l]])))
        reps = sorted(groups[lr])
        vals = [v for _, v in reps]
        mean = float(np.mean(vals))
        if a == "scratch":
            ref = mean
        rows.append((a, lr, vals, mean))
    lines = [r"\begin{tabular}{llccc}", r"\toprule", r"initialization & chosen LR & seeds & mean & $\Delta$ vs scratch \\", r"\midrule"]
    for a, lr, vals, mean in rows:
        d = "--" if a == "scratch" or ref is None else f"{mean - ref:+.3f}"
        lines.append(f"{NICE[a]} & {lr_tex(lr)} & {', '.join(f'{v:.3f}' for v in vals)} & {mean:.3f} & {d}" + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    return "\n".join(lines), {a: {"lr": lr, "vals": vals, "mean": mean} for a, lr, vals, mean in rows}


def eval_table(runs):
    """Main conditions at their replicated best LR, accuracies averaged over seeds."""
    rows, ses = [], []
    for budget in ("500M", "2000M"):
        for a in ("scratch", "ivgpt", "sp0.4", "spshuf0.4"):
            groups = defaultdict(list)
            for r in runs:
                if r["budget"] == budget and r["freeze"] == "none" and r["init"] == a and r["done"] and r["evals"]:
                    groups[r["lr"]].append(r)
            if not groups:
                continue
            lr = min(groups, key=lambda l: (-len(groups[l]), np.mean([r["val"][-1] for r in groups[l]])))
            reps = groups[lr]
            acc = {t: np.mean([r["evals"][t]["v"] for r in reps]) for t in EVAL_TASKS}
            ses += [r["evals"][t]["se"] or 0 for r in reps for t in EVAL_TASKS]
            rows.append((budget, a, lr, len(reps), acc))
    if not rows:
        return ""
    short = {"piqa": "PIQA", "hellaswag": "HellaSwag", "arc_easy": "ARC-E", "sciq": "SciQ", "winogrande": "Wino.", "lambada_openai": "LAMBADA"}
    tasks = list(EVAL_TASKS)
    lines = [r"\begin{tabular}{llc" + "c" * (len(tasks) + 1) + "}", r"\toprule",
             "tokens & initialization & seeds & " + " & ".join(short[t] for t in tasks) + r" & mean \\", r"\midrule"]
    last = None
    for budget, a, lr, n, acc in rows:
        if last and budget != last:
            lines.append(r"\midrule")
        tok = budget.replace("2000M", "2B") if budget != last else ""
        last = budget
        vals = [100 * acc[t] for t in tasks]
        lines.append(f"{tok} & {NICE[a]} & {n} & " + " & ".join(f"{v:.1f}" for v in vals) + f" & {np.mean(vals):.1f}" + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    return "\n".join(lines), float(np.mean(ses))


def adapt_table(runs):
    """Frozen backbones at 500M tokens: rows = backbone, cols = adaptation, best LR per cell (seed 0)."""
    adapts = ("body", "lora", "top4")
    cell = {}
    for b in ("ivgpt", "scratch", "shuffled"):
        for a in adapts:
            v = [r["val"][-1] for r in runs if r["init"] == b and r["freeze"] == a and r["budget"] == "500M"
                 and r["done"] and r["seed"] == 0]
            cell[b, a] = min(v) if v else None
    best = {a: min(v for (b, aa), v in cell.items() if aa == a and v is not None) for a in adapts}
    fmt = lambda v, a: "--" if v is None else (rf"\textbf{{{v:.3f}}}" if v == best[a] else f"{v:.3f}")
    lines = [r"\begin{tabular}{lccc}", r"\toprule",
             r"frozen backbone & frozen & frozen + LoRA & frozen + 4 layers \\", r"\midrule"]
    for b in ("ivgpt", "scratch", "shuffled"):
        lines.append(f"{NICE[b]} & " + " & ".join(fmt(cell[b, a], a) for a in adapts) + r" \\")
    small = [r["val"][-1] for r in runs if r["init"] == "nobody" and r["budget"] == "500M" and r["done"]]
    if small:
        lines += [r"\midrule", f"none (small 4-layer model) & -- & -- & {min(small):.3f}" + r" \\"]
    lines += [r"\bottomrule", r"\end{tabular}"]
    return "\n".join(lines)


DOMAINS = [("ssv2", "SSv2 actions"), ("droid", "DROID instructions"), ("piqa", "PIQA"),
           ("fineweb", "FineWeb-Edu"), ("wiki", "WikiText-103"), ("gsm8k", "GSM8K")]


def domain_table(probe):
    """Per-domain validation loss (mean over seeds) after 500M and 2B tokens; best per row and budget in bold."""
    import re
    g = defaultdict(list)
    for name, r in probe.items():
        m = re.search(r"-(scratch|ivgpt|sp0\.4|spshuf0\.4)-none-lr[0-9e.-]+-(\d+M)", name)
        if m:
            g[m.group(2), m.group(1)].append(r)
    cols = [("500M", k) for k in ("scratch", "ivgpt", "sp0.4", "spshuf0.4")] + [("2000M", k) for k in ("scratch", "ivgpt", "sp0.4")]
    cols = [c for c in cols if g[c]]
    short = {"scratch": "scratch", "ivgpt": "pretrained", "sp0.4": r"S\&P", "spshuf0.4": r"shuf.\ S\&P"}
    n5 = sum(c[0] == "500M" for c in cols)
    lines = [r"\begin{tabular}{l" + "c" * n5 + "c" * (len(cols) - n5) + "}", r"\toprule",
             rf" & \multicolumn{{{n5}}}{{c}}{{after 500M tokens}} & \multicolumn{{{len(cols) - n5}}}{{c}}{{after 2B tokens}} \\",
             rf"\cmidrule(lr){{2-{n5 + 1}}}\cmidrule(lr){{{n5 + 2}-{len(cols) + 1}}}",
             "text & " + " & ".join(short[k] for _, k in cols) + r" \\", r"\midrule"]
    for d, lab in DOMAINS:
        vals = {c: float(np.mean([r[d] for r in g[c]])) for c in cols}
        best = {b: min(v for c, v in vals.items() if c[0] == b) for b in ("500M", "2000M")}
        lines.append(lab + " & " + " & ".join(
            rf"\textbf{{{vals[c]:.3f}}}" if vals[c] == best[c[0]] else f"{vals[c]:.3f}" for c in cols) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    return "\n".join(lines)


def main():
    runs = load_runs(["results/ice", "results/dm_backup"])
    out = {}
    tex = []
    for key, (budget, freeze) in {"SweepFull": ("500M", "none"), "PilotFull": ("100M", "none"),
                                  "PilotFrozen": ("100M", "body"), "SweepFrozen": ("500M", "body")}.items():
        t = sweep_table(runs, budget, freeze)
        if t:
            tex.append(rf"\newcommand{{\Table{key}}}{{%" + "\n" + t + "}")
    for budget, key in (("500M", "RepFive"), ("2000M", "RepTwoB")):
        t, data = replicate_table(runs, budget)
        if data:
            tex.append(rf"\newcommand{{\Table{key}}}{{%" + "\n" + t + "}")
            out[key] = data
    et = eval_table(runs)
    if et:
        tex.append(r"\newcommand{\TableEval}{%" + "\n" + et[0] + "}")
        out["eval_mean_stderr"] = et[1]
    tex.append(r"\newcommand{\TableAdapt}{%" + "\n" + adapt_table(runs) + "}")
    pf = Path("report/probe.json")
    if pf.exists():
        tex.append(r"\newcommand{\TableDomains}{%" + "\n" + domain_table(json.loads(pf.read_text())) + "}")
    Path("report/tables.tex").write_text("\n\n".join(tex) + "\n")
    Path("report/key_numbers.json").write_text(json.dumps(out, indent=1))
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()