#!/usr/bin/env python3
"""Joules per generated token from a powermetrics sample file and the benchmark records that carry epoch timestamps.
Energy of a request = trapezoid integral of combined package power over [t_first_token, t_end] (decode phase) minus the idle
baseline power measured before the first run (so that only the incremental energy of decoding is attributed to tokens);
both gross and net values are reported. Writes benchmarks/processed/energy_<tag>.json and paper/tables/energy.tex."""
import json, os, statistics, sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "benchmarks", "scripts"))
from power_sampler import energy_joules
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

def main(d="benchmarks/raw/gateG/energy_laptop", tag="laptop"):
    d = os.path.join(ROOT, d); pw = os.path.join(d, "power.jsonl"); bj = os.path.join(d, "bench.jsonl")
    pts = [json.loads(l) for l in open(pw)]; recs = [json.loads(l) for l in open(bj) if l.strip()]
    t_first_run = min(r["t_start_epoch"] for r in recs)
    idle = [p["combined_mW"] for p in pts if p["t"] < t_first_run - 1 and p.get("combined_mW") is not None]
    idle_mW = statistics.median(idle) if idle else None
    groups = {}
    for r in recs:
        if r.get("warmup"): continue
        k = (r["backend"], r.get("capacity"), r.get("predictor") or "none", bool(r.get("nocache")))
        e = energy_joules(pw, r["t_first_epoch"], r["t_end_epoch"])
        if e is None: continue
        dur = r["t_end_epoch"] - r["t_first_epoch"]; n = r["gen_tokens"]
        groups.setdefault(k, []).append({"J_per_token_gross": e / n, "J_per_token_net": (e - idle_mW * 1e-3 * dur) / n if idle_mW else None,
                                         "mean_W": e / dur, "tok_s": r["gen_tps"], "n": n})
    out = {"idle_mW": idle_mW, "samples": len(pts), "configs": []}
    for k, v in groups.items():
        out["configs"].append({"backend": k[0], "capacity": k[1], "predictor": k[2], "cold": k[3], "runs": len(v),
                               "J_per_token_gross_median": statistics.median(x["J_per_token_gross"] for x in v),
                               "J_per_token_net_median": statistics.median(x["J_per_token_net"] for x in v) if idle_mW else None,
                               "mean_W_median": statistics.median(x["mean_W"] for x in v), "tok_s_median": statistics.median(x["tok_s"] for x in v)})
    os.makedirs(os.path.join(ROOT, "benchmarks/processed"), exist_ok=True)
    json.dump(out, open(os.path.join(ROOT, f"benchmarks/processed/energy_{tag}.json"), "w"), indent=1)
    L = ["\\begin{table}[t]", "\\centering\\small", f"\\caption{{Energy per generated token on the laptop (M5 Max), OLMoE 4-bit, 256-token greedy decode, held-out general/code prompts; package power from \\code{{powermetrics}} at 1 s; idle baseline {idle_mW/1e3:.1f} W. Source: \\code{{benchmarks/raw/gateG/energy\\_laptop/}}.}}", "\\label{tab:energy}",
         "\\begin{tabular}{llrrrrr}", "\\toprule", "Backend & residency & runs & tok/s & mean W & J/token (gross) & J/token (net of idle) \\\\", "\\midrule"]
    for c in sorted(out["configs"], key=lambda c: -c["tok_s_median"]):
        res = "all resident" if c["backend"] == "mlx_lm" or (c["capacity"] or 0) >= 64 else f"{c['capacity']} slots{', cold NVMe' if c['cold'] else ''}{', ' + c['predictor'] if c['predictor'] != 'none' else ''}"
        L.append(f"{c['backend'].replace('_', '-')} & {res} & {c['runs']} & {c['tok_s_median']:.1f} & {c['mean_W_median']:.1f} & {c['J_per_token_gross_median']:.3f} & {c['J_per_token_net_median']:.3f} \\\\")
    L += ["\\bottomrule", "\\end{tabular}", "\\end{table}"]
    open(os.path.join(ROOT, "paper/tables/energy.tex"), "w").write("\n".join(L) + "\n")
    print(json.dumps(out, indent=1))

if __name__ == "__main__":
    main(*sys.argv[1:])