#!/usr/bin/env python3
"""Generate LaTeX table fragments under paper/tables/ from processed/raw logs. Every number in the paper's tables comes from here."""
import json, os, sys, glob, statistics
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "paper", "tables"); os.makedirs(OUT, exist_ok=True)

def sim_table(path, label, caption):
    d = json.load(open(path)); rows = d["results"]
    caps = sorted({r["capacity"] for r in rows}); preds = []
    for r in rows:
        if r["predictor"] not in preds: preds.append(r["predictor"])
    lines = ["\\begin{table}[t]", "\\centering\\small", f"\\caption{{{caption}}}", f"\\label{{{label}}}",
             "\\begin{tabular}{ll" + "rr" * len(caps) + "}", "\\toprule",
             "Predictor & split & " + " & ".join(f"\\multicolumn{{2}}{{c}}{{$C={c}$}}" for c in caps) + " \\\\",
             " & & " + " & ".join("hit & stall" for _ in caps) + " \\\\", "\\midrule"]
    for p in preds:
        for split in ("fit", "heldout"):
            cells = []
            for c in caps:
                m = [r for r in rows if r["predictor"] == p and r["split"] == split and r["capacity"] == c and r["policy"] == "lru"]
                cells.append(f"{m[0]['hit_rate']:.3f} & {m[0]['stall_ms_per_token']:.1f}" if m else "-- & --")
            lines.append(f"{p.replace('_', '\\_')} & {split} & " + " & ".join(cells) + " \\\\")
    lines += ["\\bottomrule", "\\end{tabular}", "\\begin{tablenotes}\\footnotesize", f"\\item Source: \\code{{{os.path.relpath(path, ROOT).replace('_', '\\_')}}}; t\\_read {d['t_read_ms']} ms, t\\_layer {d['t_layer_ms']} ms, {d['loaders']} loader(s); hit = per-layer hit rate, stall in ms per token (modelled).", "\\end{tablenotes}", "\\end{table}"]
    open(os.path.join(OUT, f"{label.split(':')[-1]}.tex"), "w").write("\n".join(lines) + "\n")
    print("wrote", label)

def bench_summary(paths, label, caption, keys=("backend", "capacity", "predictor", "m", "nocache", "split")):
    """Median decode tok/s, hit rate and stall per configuration from bench_decode JSONL files (warm-up excluded)."""
    recs = []
    for p in paths:
        recs += [json.loads(l) for l in open(p) if l.strip()]
    recs = [r for r in recs if not r.get("warmup")]
    groups = {}
    for r in recs:
        k = tuple(r.get(x) for x in keys); groups.setdefault(k, []).append(r)
    lines = ["\\begin{table}[t]", "\\centering\\small", f"\\caption{{{caption}}}", f"\\label{{{label}}}",
             "\\begin{tabular}{llrrrrrr}", "\\toprule", "Backend & cfg & $n$ & decode tok/s (median) & p95 ITL ms & hit & stall s & TTFT ms \\\\", "\\midrule"]
    for k in sorted(groups, key=lambda x: tuple(str(v) for v in x)):
        g = groups[k]
        med = lambda f: statistics.median([r[f] for r in g if r.get(f) is not None]) if any(r.get(f) is not None for r in g) else None
        cfg = f"C={k[1]}, {k[2]}, m={k[3]}" + (", cold" if k[4] else "") if k[0] == "titan" else "--"
        f = lambda v, fmt: (fmt % v) if v is not None else "--"
        lines.append(f"{k[0].replace('_', '\\_')} & {cfg.replace('_', '\\_')} & {len(g)} & {f(med('gen_tps'), '%.1f')} & {f(med('itl_ms_p95'), '%.1f')} & {f(med('cache_hit_rate'), '%.3f')} & {f(med('stall_s'), '%.2f')} & {f(med('ttft_ms'), '%.0f')} \\\\")
    lines += ["\\bottomrule", "\\end{tabular}", "\\begin{tablenotes}\\footnotesize", "\\item Sources: " + ", ".join(f"\\code{{{os.path.relpath(p, ROOT).replace('_', '\\_')}}}" for p in paths) + ".", "\\end{tablenotes}", "\\end{table}"]
    open(os.path.join(OUT, f"{label.split(':')[-1]}.tex"), "w").write("\n".join(lines) + "\n")
    print("wrote", label)

if __name__ == "__main__":
    v2 = os.path.join(ROOT, "benchmarks/processed/gateC/sim_olmoe_v2_ssd043_m8.json")
    v1 = os.path.join(ROOT, "benchmarks/processed/gateC/sim_olmoe_laptop_ssd_v1.json")
    sim_table(v2 if os.path.exists(v2) else v1, "tab:sim", "Trace-driven simulation (OLMoE, per-layer LRU tables, cold per prompt): hit rate and modelled stall per token by predictor, split and capacity $C$.")
    cold = os.path.join(ROOT, "benchmarks/raw/gateC/cold_nvme_laptop/sweep.jsonl")
    if os.path.exists(cold):
        bench_summary([cold], "tab:cold", "Cold-NVMe tier on the M5 Max laptop (purge + F\\_NOCACHE), OLMoE 4-bit, held-out prompts, 64 tokens.")
    q = os.path.join(ROOT, "benchmarks/raw/gateC/qwen3_m3u96a/baseline.jsonl")
    if os.path.exists(q):
        bench_summary([q], "tab:qwen3", "Qwen3-30B-A3B-Instruct-2507 4-bit on M3U96a (M3 Ultra, 96 GB): mlx-lm vs TITAN all-resident (128 slots, preloaded), 64 tokens, 3 repeats.")