"""Tables for the probe and representation results -> report/tables_extra.tex.

TableDomainsTwo: per-domain validation loss after 2B tokens (report/probe2.json), mean over seeds.
TableRank: effective rank of final-layer features and dormant MLP units (report/rank.json, rank2.json).
"""

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

import numpy as np

REPORT = Path(__file__).resolve().parents[1] / "report"
DOMAINS = [("ssv2", "SSv2 action templates"), ("ssv2_abs", "same templates, abstract nouns"),
           ("abs_tmpl", "other templates, abstract nouns"), ("droid", "DROID robot instructions"),
           ("piqa", "PIQA"), ("fineweb", "FineWeb-Edu"), ("wiki", "WikiText-103"), ("gsm8k", "GSM8K")]
MEDIUM = [("scratch", r"(?:^|-)scratch-none-lr2e-3"), ("pretrained", r"-ivgpt-none-lr3e-3"), ("S\\&P", r"-sp0\.4-none")]
SMALL = [("scratch", r"^act-sm-scratch"), ("shuffled", r"^act-base-shuffled"), ("pretrained", r"^act-base-ivgpt")]


def group(probe, pattern, budget):
    rows = [v for k, v in probe.items() if re.search(pattern, k) and k.split("-")[-1 if "-s" not in k[-3:] else -2] == budget]
    return rows


def domain_table(probe):
    cols = [(lab, group(probe, p, "2000M")) for lab, p in MEDIUM] + [(lab, group(probe, p, "2000M")) for lab, p in SMALL]
    for lab, rs in cols:
        assert rs, lab
    lines = [r"\newcommand{\TableDomainsTwo}{", r"\begin{tabular}{lcccccc}", r"\toprule",
             r"& \multicolumn{3}{c}{medium (400M backbone)} & \multicolumn{3}{c}{small (113M backbone)} \\",
             r"\cmidrule(lr){2-4}\cmidrule(lr){5-7}",
             "text & " + " & ".join(lab for lab, _ in cols) + r" \\", r"\midrule"]
    for key, name in DOMAINS:
        vals = [float(np.mean([r[key] for r in rs])) for _, rs in cols]
        cells = []
        for block in (vals[:3], vals[3:]):
            lo = min(block)
            cells += [rf"\textbf{{{v:.3f}}}" if v == lo else f"{v:.3f}" for v in block]
        lines.append(f"{name} & " + " & ".join(cells) + r" \\")
        if key == "abs_tmpl":
            lines.append(r"\midrule")
    lines.append(r"\midrule")
    for key, name in [("copy_second", "copying, second copy"), ("icl_score", "in-context score")]:
        vals = [float(np.mean([r[key] for r in rs])) for _, rs in cols]
        lines.append(f"{name} & " + " & ".join(f"{v:.2f}" if key == "copy_second" else f"{v:.3f}" for v in vals) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}}"]
    print("seeds per column:", [(lab, len(rs)) for lab, rs in cols])
    return lines


def rank_table(rank):
    def mean_of(prefixes):
        rs = [v for k, v in rank.items() if any(k == p or (p.startswith("trained/") and k.startswith(p)) for p in prefixes)]
        return (float(np.mean([r["eff_rank"] for r in rs])), float(np.mean([r["dormant"] for r in rs])),
                float(np.mean([np.mean(r["per_layer_dormant"][:2]) for r in rs])), len(rs))

    rows = [
        ("medium", "scratch", ["init/m_scratch"], ["trained/10006a-scratch", "trained/050000-scratch"]),
        ("", "pretrained", ["init/m_ivgpt"], ["trained/10007-ivgpt", "trained/050000-ivgpt", "trained/z0723-ivgpt"]),
        ("", "shuffled", ["init/m_shuffled"], ["trained/z0700-shuffled"]),
        ("", "rescaled", ["init/m_rescaled"], []),
        ("", "S\\&P", ["init/m_sp"], ["trained/0600-sp0.4", "trained/0551-sp0.4"]),
        ("", "shuffled S\\&P", ["init/m_spshuf"], ["trained/z0702-spshuf"]),
        ("small", "scratch", ["init/s_scratch"], ["trained/act-sm-scratch"]),
        ("", "pretrained", ["init/s_base"], ["trained/act-base-ivgpt"]),
        ("", "shuffled", ["init/s_shuffled"], ["trained/act-base-shuffled"]),
        ("", "S\\&P", ["init/s_sp"], []),
    ]
    lines = [r"\newcommand{\TableRank}{", r"\begin{tabular}{llcccccc}", r"\toprule",
             r"& & \multicolumn{3}{c}{at initialization} & \multicolumn{3}{c}{after 500M tokens} \\",
             r"\cmidrule(lr){3-5}\cmidrule(lr){6-8}",
             r"model & initialization & eff.\ rank & dormant & dormant, layers 1--2 & eff.\ rank & dormant & dormant, layers 1--2 \\",
             r"\midrule"]
    for size, lab, init, trained in rows:
        if size == "small":
            lines.append(r"\midrule")
        a = mean_of(init)
        cells = [f"{a[0]:.0f}", f"{100 * a[1]:.1f}\\%", f"{100 * a[2]:.1f}\\%"]
        if trained:
            b = mean_of(trained)
            cells += [f"{b[0]:.0f}", f"{100 * b[1]:.1f}\\%", f"{100 * b[2]:.1f}\\%"]
            print(size, lab, "trained seeds", b[3])
        else:
            cells += ["--", "--", "--"]
        lines.append(f"{size} & {lab} & " + " & ".join(cells) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}}"]
    return lines


def main():
    probe = json.loads((REPORT / "probe2.json").read_text())
    rank = json.loads((REPORT / "rank.json").read_text()) | json.loads((REPORT / "rank2.json").read_text())
    out = domain_table(probe) + rank_table(rank)
    (REPORT / "tables_extra.tex").write_text("\n".join(out) + "\n")
    print("wrote", REPORT / "tables_extra.tex")


if __name__ == "__main__":
    main()