#!/usr/bin/env python3
"""Sample package power with powermetrics (requires sudo) at 1 s intervals into a JSONL file; a companion function integrates
energy over a benchmark's wall-clock window to produce joules per token. Usage:
  sudo powermetrics --samplers cpu_power,gpu_power,thermal -i 1000 -f plist | python3 power_sampler.py --out power.jsonl
"""
import argparse, json, plistlib, sys, time

def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--out", required=True); a = ap.parse_args()
    buf = b""; out = open(a.out, "a")
    for chunk in iter(lambda: sys.stdin.buffer.read(65536), b""):
        buf += chunk
        while b"</plist>" in buf:
            i = buf.index(b"</plist>") + len(b"</plist>")
            doc, buf = buf[:i], buf[i:]
            s = doc.find(b"<?xml")
            if s < 0: continue
            try:
                d = plistlib.loads(doc[s:])
            except Exception:
                continue
            rec = {"t": time.time(), "elapsed_ns": d.get("elapsed_ns"),
                   "cpu_mW": (d.get("processor") or {}).get("cpu_power"), "gpu_mW": (d.get("processor") or {}).get("gpu_power"),
                   "ane_mW": (d.get("processor") or {}).get("ane_power"), "combined_mW": (d.get("processor") or {}).get("combined_power"),
                   "thermal_pressure": (d.get("thermal_pressure"))}
            out.write(json.dumps(rec) + "\n"); out.flush()

def energy_joules(power_jsonl, t0, t1, field="combined_mW"):
    """Integrate power (mW) over [t0, t1] (epoch seconds) with the trapezoid rule; returns joules."""
    pts = [json.loads(l) for l in open(power_jsonl)]
    pts = [(p["t"], p[field]) for p in pts if p.get(field) is not None and t0 - 1.5 <= p["t"] <= t1 + 1.5]
    if len(pts) < 2: return None
    e = 0.0
    for (ta, pa), (tb, pb) in zip(pts, pts[1:]):
        lo, hi = max(ta, t0), min(tb, t1)
        if hi > lo: e += (pa + pb) / 2 * 1e-3 * (hi - lo)
    return e

if __name__ == "__main__":
    main()