"""CSV -> LaTeX fragments for the paper (booktabs). Numbers are formatted here only; values come from results/tables CSVs."""
from __future__ import annotations
import pandas as pd, numpy as np
from pathlib import Path
from ..util import ROOT
TAB = ROOT / "results/tables"; OUT = ROOT / "paper/tables"

def esc(s):
    return str(s).replace("_", "\\_").replace("%", "\\%").replace("&", "\\&").replace("→", "$\\to$").replace("≥", "$\\ge$").replace("×", "$\\times$").replace("—", "--")

def df_to_tex(df: pd.DataFrame, path: Path, colfmt: str | None = None, floatfmt="{:.2f}"):
    cols = list(df.columns); colfmt = colfmt or "l" * len(cols)
    lines = ["\\begin{tabular}{@{}" + colfmt + "@{}}", "\\toprule", " & ".join(esc(c) for c in cols) + " \\\\", "\\midrule"]
    for _, r in df.iterrows():
        cells = []
        for c in cols:
            v = r[c]
            if isinstance(v, float):
                cells.append("--" if np.isnan(v) else (floatfmt.format(v) if abs(v) < 1e4 else f"{v:,.0f}"))
            else:
                cells.append(esc(v))
        lines.append(" & ".join(cells) + " \\\\")
    lines += ["\\bottomrule", "\\end{tabular}"]
    OUT.mkdir(parents=True, exist_ok=True); path.write_text("\n".join(lines) + "\n")

def main():
    t1 = pd.read_csv(TAB / "table1_taxonomy.csv"); df_to_tex(t1, OUT / "table1_taxonomy.tex", "lp{4.2cm}p{2cm}p{2.3cm}p{2.3cm}p{1.6cm}p{2.6cm}")
    t2 = pd.read_csv(TAB / "table2_models.csv")[["alias", "repo", "revision", "layers", "kv_heads", "q_heads", "head_dim", "max_pos", "cache_bytes_per_token_bf16", "prefill_ms_8k", "prefill_ms_16k"]]
    t2.columns = ["alias", "checkpoint", "rev.", "L", "H_kv", "H_q", "d_h", "max pos", "cache B/token", "prefill 8k", "prefill 16k"]; df_to_tex(t2, OUT / "table2_models.tex", "llllllllrrr")
    t3 = pd.read_csv(TAB / "table3_datasets.csv"); df_to_tex(t3, OUT / "table3_datasets.tex", "llrrr", "{:.0f}")
    print("latex tables written")

if __name__ == "__main__":
    main()