"""Joules per typed decision, metered with the same counter jcp-energy's rung 3 uses.

ABI constants are lifted from jcp-energy (sizeof/offsetof-verified against the
macOS 26, 26.5 and 27 SDKs), not re-derived here.
"""
import ctypes, os, time, statistics, json, warnings

FLAVOR, SIZE, ENERGY_NJ, PENERGY_NJ = 6, 464, 336, 344
_libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib")
_buf = ctypes.create_string_buffer(SIZE)
_pid = os.getpid()

def read_nj():
    rc = _libc.proc_pid_rusage(ctypes.c_int(_pid), ctypes.c_int(FLAVOR), ctypes.byref(_buf))
    if rc != 0:
        raise OSError(f"proc_pid_rusage failed: rc={rc}")
    e = int.from_bytes(_buf.raw[ENERGY_NJ:ENERGY_NJ+8], "little")
    p = int.from_bytes(_buf.raw[PENERGY_NJ:PENERGY_NJ+8], "little")
    return e, p

def liveness(window_s=0.5, interval_s=0.02):
    """jcp-energy's rule: >=5 advancing intervals AND >=80% of them, over >=300ms."""
    advanced = total = 0
    prev, _ = read_nj(); t_end = time.perf_counter() + window_s
    while time.perf_counter() < t_end:
        t1 = time.perf_counter() + interval_s
        x = 0
        while time.perf_counter() < t1:     # burn CPU
            x += 1
        cur, _ = read_nj(); total += 1
        if cur > prev: advanced += 1
        prev = cur
    ok = advanced >= 5 and total > 0 and advanced / total >= 0.80
    return ok, advanced, total

import laya_mlx as laya
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    agent = laya.load("aac6fef/laya-mlx", dtype="float16")

STATE = "I was billed twice. Please refund the duplicate today."
Q1 = {"refund": {"type":"noul","instructions":"Does the customer ask for money back?"}}
Q3 = {
 "department": {"type":"choice","instructions":"Which team should handle this request?",
                "criteria":{"billing":"invoices, payments, refunds","technical":"bugs and outages","sales":"new purchases"}},
 "urgency":    {"type":"score","instructions":"How urgent is this request?","criteria":["not urgent","soon","critical"]},
 "refund":     {"type":"noul","instructions":"Does the customer ask for money back?"}}

def bench(label, questions, n, device):
    a = laya.load("aac6fef/laya-mlx", dtype="float16", device=device) if device else agent
    for _ in range(3): a.predict(STATE, questions)          # warm
    e0, p0 = read_nj(); lat = []
    t0 = time.perf_counter()
    for _ in range(n):
        t = time.perf_counter(); r = a.predict(STATE, questions); lat.append((time.perf_counter()-t)*1e3)
    wall = time.perf_counter() - t0
    e1, p1 = read_nj()
    u = r["usage"]
    return {
        "label": label, "device": device or "default", "n": n,
        "p50_ms": round(statistics.median(lat), 2),
        "p95_ms": round(sorted(lat)[int(0.95*len(lat))-1], 2),
        "wall_s": round(wall, 3),
        "cpu_J_total": round((e1-e0)/1e9, 6),
        "cpu_J_per_decision": round((e1-e0)/1e9/n, 6),
        "pcore_J_per_decision": round((p1-p0)/1e9/n, 6),
        "input_tokens": u["input_tokens"], "output_tokens": u["output_tokens"],
        "state_tokens": u["state_tokens"],
    }

ok, adv, tot = liveness()
print(json.dumps({"counter":"ri_energy_nj via proc_pid_rusage(RUSAGE_INFO_V6)",
                  "live": ok, "advancing_intervals": adv, "intervals": tot,
                  "rule": "jcp-energy rung 3: >=5 advancing AND >=80%, >=300ms"}, indent=2))
rows = []
for dev in ("cpu", "gpu"):
    for lbl, qs, n in (("1 question (noul)", Q1, 30), ("3 questions (choice+score+noul)", Q3, 30)):
        try: rows.append(bench(lbl, qs, n, dev))
        except Exception as ex: rows.append({"label": lbl, "device": dev, "error": repr(ex)[:200]})
print(json.dumps(rows, indent=2))