"""Static LaTeX tables that do not depend on model results: universe strata, instrument list, baseline specs, benchmark."""
from __future__ import annotations

import json
from pathlib import Path

import pandas as pd
import yaml

ROOT = Path(__file__).resolve().parents[1]
T = ROOT / "paper" / "tables"; T.mkdir(parents=True, exist_ok=True)


def esc(s) -> str:
    return str(s).replace("_", "\\_").replace("&", "\\&").replace("%", "\\%").replace("$", "\\$")


def universe_tables() -> None:
    man = json.loads((ROOT / "data" / "manifests" / "universe_main.json").read_text())
    df = pd.DataFrame(man["instruments"])
    piv = df.pivot_table(index=["asset_class", "adv_tercile"], columns="vol_tercile", values="symbol", aggfunc="count", fill_value=0)
    piv = piv.reindex(columns=["low", "mid", "high"])
    lines = ["\\begin{table}[htbp]\\centering\\small", "\\caption{Universe allocation: instruments per asset class $\\times$ dollar-volume tercile $\\times$ daily-volatility tercile (selection date 2017-12-29; 60 of the 300 instruments are held out)}",
             "\\label{tab:universe_strata}", "\\begin{tabular}{llccc}", "\\toprule", "Asset class & ADV tercile & Vol. low & Vol. mid & Vol. high \\\\", "\\midrule"]
    for (ac, adv), r in piv.iterrows():
        lines.append(f"{ac} & {adv} & {r['low']} & {r['mid']} & {r['high']} \\\\")
    held = df.groupby("asset_class")["role"].apply(lambda s: (s == "heldout_instrument").sum())
    lines += ["\\midrule", f"held out & & \\multicolumn{{3}}{{c}}{{stocks {held.get('stock', 0)}, ETFs {held.get('etf', 0)}}} \\\\", "\\bottomrule", "\\end{tabular}", "\\end{table}"]
    (T / "universe_strata.tex").write_text("\n".join(lines))
    # compact instrument list (tickers only, allowed by the license)
    pooled = sorted(df[df.role == "pooled"].symbol); held = sorted(df[df.role != "pooled"].symbol)
    def block(name, syms):
        return [f"\\paragraph{{{name} ({len(syms)}).}}", "\\small " + esc(", ".join(syms)) + "\\normalsize"]
    (T / "universe_list.tex").write_text("\n".join(block("Pooled instruments", pooled) + block("Held-out instruments", held)))
    # group counts
    g = df.groupby(["asset_class", "group"]).size().reset_index(name="n")
    lines = ["\\begin{table}[htbp]\\centering\\small", "\\caption{Sector (SIC division) and ETF category counts in the universe}", "\\label{tab:universe_groups}",
             "\\begin{tabular}{llr}", "\\toprule", "Asset class & Group & $n$ \\\\", "\\midrule"]
    lines += [f"{r.asset_class} & {esc(r.group)} & {r.n} \\\\" for r in g.itertuples()]
    lines += ["\\bottomrule", "\\end{tabular}", "\\end{table}"]
    (T / "universe_groups.tex").write_text("\n".join(lines))


def baseline_spec_table() -> None:
    spec = yaml.safe_load((ROOT / "configs" / "baselines.yaml").read_text())
    names = {"naive_rw": "Random walk", "seasonal_vol": "Seasonal volatility", "ewma": "EWMA-$t$", "har_rv": "HAR-RV", "garch_t": "GARCH(1,1)-$t$",
             "ridge": "Ridge", "elastic_net": "Elastic net", "lightgbm": "LightGBM", "mlp": "MLP", "tcn": "TCN"}
    lines = ["\\begin{table}[htbp]\\centering\\scriptsize", "\\caption{Baseline specifications (information set = completed bars at the origin; tuned on training/validation periods only)}",
             "\\label{tab:baselines}", "\\begin{tabular}{p{2.6cm}p{12.5cm}}", "\\toprule", "Model & Specification \\\\", "\\midrule"]
    for k, v in spec.items():
        if k in ("name", "description", "not_run"):
            continue
        desc = "; ".join(f"{esc(a)}: {esc(b)}" for a, b in v.items())
        lines.append(f"{names.get(k, esc(k))} & {desc} \\\\")
    lines += ["\\midrule", "Not run & " + "; ".join(f"\\emph{{{esc(a)}}}: {esc(b)}" for a, b in spec["not_run"].items()) + " \\\\", "\\bottomrule", "\\end{tabular}", "\\end{table}"]
    (T / "baselines_spec.tex").write_text("\n".join(lines))


def benchmark_table() -> None:
    f = ROOT / "results" / "metrics" / "benchmark_M3U96b.json"
    if not f.exists():
        return
    b = json.loads(f.read_text())
    lines = ["\\begin{table}[htbp]\\centering\\small", f"\\caption{{Measured training throughput and memory on the compute node ({esc(b['chip'])}, {esc(b['node'])}, torch {esc(b['torch'])}, {esc(b['device'])}; FP32)}}",
             "\\label{tab:benchmark}", "\\begin{tabular}{lrrrrrr}", "\\toprule", "Preset & Batch & Parameters & Steps/s & Examples/s & Peak memory (MB) & Latency/origin (ms) \\\\", "\\midrule"]
    for k, r in b["train"].items():
        preset, bs = k.rsplit("_bs", 1)
        inf = b["inference"].get(k, {})
        lines.append(f"{preset} & {bs} & {r['params']:,} & {r['steps_per_s']:.2f} & {r['ex_per_s']:.0f} & {r['peak_mem_mb']:.0f} & {inf.get('latency_ms_per_origin_batched', float('nan')):.2f} \\\\")
    lines += ["\\midrule", "Data loader (64/batch) & " + "; ".join(f"{k.replace('workers_', '')} workers: {v['ex_per_s']:.0f} ex/s" for k, v in b["loader"].items()) + " & & & & & \\\\",
              "\\bottomrule", "\\end{tabular}", "\\end{table}"]
    (T / "benchmark.tex").write_text("\n".join(lines).replace(",", "{,}").replace("{,}{,}", ","))


if __name__ == "__main__":
    universe_tables(); baseline_spec_table(); benchmark_table()
    print("static tables written:", sorted(p.name for p in T.glob("*.tex")))