"""Print Tables 2, 5 and 6 of the paper from records/<system>/units.json. No data, no GPU.

    python -m reproduce.tables [table2 | table5 | table6 | all]

`python -m reproduce.table2` regenerates any entry of Table 2 or Table 6 from a checkpoint (or rescores the archived
samples) and compares it with these records.
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

RECORDS = Path(__file__).resolve().parents[1] / "records"
ORDER = ("euler", "sqg", "waves", "nlkg", "passive_scalar")


def _rec(system):
    return json.loads((RECORDS / system / "units.json").read_text())


def _unit(rec, key):
    return next(u for u in rec["units"] if u["key"] == key)


def _fmt(e):
    return f"{e['mean']:.2f} +- {e['sem']:.2f}"


def table2():
    print("Table 2: test csFID-10k, mean +- SEM over the ten test conditions (+: a selected hyperparameter lies on the boundary of its searched range)")
    print(f"{'system':16s} {'Direct SI':>18s} {'EDM':>18s} {'Cascaded SI':>18s} {'S*':>4s}")
    for s in ORDER:
        r = _rec(s)
        cells = []
        for m in ("direct_si", "edm", "cascaded_si"):
            u = _unit(r, r["table2"][m])
            flag = " +" if u.get("selection_flags") else "  "
            cells.append(_fmt(u["expected_csfid"]) + flag)
        print(f"{s:16s} {cells[0]:>18s} {cells[1]:>18s} {cells[2]:>18s} {r['s_star']:4d}")


def table5():
    print("Table 5: selected configurations (S, learning rate, noise parameter, selected update, sampler steps per segment, NFE)")
    for s in ORDER:
        r = _rec(s)
        for m in ("direct_si", "edm", "cascaded_si"):
            if s == "passive_scalar" and m == "cascaded_si":
                continue
            u = _unit(r, r["table2"][m]); sel = u["selection"]
            noise = f"P_mean={sel['P_mean']:g}" if "P_mean" in sel else f"sigma={sel['sigma']:g}"
            print(f"{s:16s} {m:12s} S={u['S']:<3d} lr={sel['lr']:<8g} {noise:16s} update={sel['selected_update']:<6d} "
                  f"steps={u['sampler']['n_steps']:<4d} NFE={u['sampler']['nfe']}")


def table6():
    print("Table 6 / Fig. 3: test csFID-10k of the selected Cascaded SI model at each depth S (S = 1 is Direct SI), and EDM")
    for s in ORDER:
        r = _rec(s)
        row = [f"S={u['S']}: {_fmt(u['expected_csfid'])}" for u in sorted((u for u in r["units"] if u["key"] != "edm"), key=lambda u: u["S"])]
        print(f"{s:16s} " + " | ".join(row) + f" | EDM: {_fmt(_unit(r, 'edm')['expected_csfid'])}")


if __name__ == "__main__":
    what = sys.argv[1] if len(sys.argv) > 1 else "all"
    for name, fn in (("table2", table2), ("table5", table5), ("table6", table6)):
        if what in (name, "all"):
            fn(); print()