"""Device benchmark: latency vs temperature under sustained load, plus accuracy, for the numpy runtime.
  python bench_device.py --model ../model --data <dir with *.jsonl eval sets> --out result.json [--threads-probe]
Phases: idle temperature (30 s) -> warm-up -> every prompt of every eval set once (sustained load; stats per 60 s
window: temperature, CPU clock, latency p50/p95, prompts/s) -> optional 1-thread vs default-thread probe.
Accuracy = runtime output (with snap) == target, i.e. what `run.sh test --full` measures, not the raw Colab number.
Temperature: Raspberry Pi = SoC sensor (vcgencmd); Windows = ACPI thermal zone (not the CPU core sensor).
"""
import argparse, json, os, platform, statistics, subprocess, sys, threading, time

SETS = ["curated", "wifi_holdout", "gap11_holdout", "B_golden_final", "A_practical-holdout"]

def is_pi(): return os.path.exists("/usr/bin/vcgencmd")
PS = os.path.join(os.environ.get("SystemRoot", r"C:\Windows"), "System32", "WindowsPowerShell", "v1.0", "powershell.exe")

def sample():
    """-> dict(temp_c, clock_mhz, throttled) for this device; missing values are None."""
    if is_pi():
        def v(*a): return subprocess.run(["vcgencmd", *a], capture_output=True, text=True).stdout.strip()
        t = float(v("measure_temp").split("=")[1].rstrip("'C")); f = int(v("measure_clock", "arm").split("=")[1]) / 1e6
        fan = open("/sys/class/thermal/cooling_device0/cur_state").read().strip() if os.path.exists("/sys/class/thermal/cooling_device0/cur_state") else None
        return {"temp_c": t, "clock_mhz": round(f), "throttled": v("get_throttled").split("=")[1], "fan_state": fan}
    if os.name == "nt":
        ps = ("$c=Get-Counter '\\Thermal Zone Information(*)\\Temperature','\\Processor Information(_Total)\\% Processor Performance','\\Processor Information(_Total)\\Processor Frequency';"
              "$c.CounterSamples | ForEach-Object { $_.Path + '=' + $_.CookedValue }")
        out = subprocess.run([PS, "-NoProfile", "-Command", ps], capture_output=True, text=True).stdout
        temps, perf, freq = [], None, None
        for line in out.splitlines():
            k, _, val = line.rpartition("=")
            if not val: continue
            if "temperature" in k.lower(): temps.append(float(val) - 273.15)
            elif "% processor performance" in k.lower(): perf = float(val)
            elif "processor frequency" in k.lower(): freq = float(val)
        clock = round(freq * perf / 100) if freq and perf else None
        return {"temp_c": round(max(temps), 1) if temps else None, "clock_mhz": clock, "throttled": None, "fan_state": None}
    return {"temp_c": None, "clock_mhz": None, "throttled": None, "fan_state": None}

class Sampler(threading.Thread):
    def __init__(self, every=10):
        super().__init__(daemon=True); self.every, self.rows, self.stop = every, [], False
    def run(self):
        while not self.stop:
            s = sample(); s["t"] = time.time(); self.rows.append(s); time.sleep(self.every)

def pct(xs, p):
    xs = sorted(xs); return xs[min(len(xs) - 1, int(round(p / 100 * (len(xs) - 1))))]

def device_info():
    info = {"platform": platform.platform(), "python": platform.python_version(), "machine": platform.machine()}
    import numpy; info["numpy"] = numpy.__version__
    if is_pi():
        info["model"] = open("/proc/device-tree/model").read().strip("\x00 \n"); info["cores"] = os.cpu_count()
        info["governor"] = open("/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor").read().strip()
    elif os.name == "nt":
        out = subprocess.run([PS, "-NoProfile", "-Command", "(Get-CimInstance Win32_Processor).Name"], capture_output=True, text=True).stdout.strip()
        info["model"] = out; info["cores"] = os.cpu_count()
    return info

def load_sets(data):
    rows = []
    for s in SETS:
        p = os.path.join(data, s + ".jsonl")
        if os.path.exists(p): rows += [(s, json.loads(l)) for l in open(p, encoding="utf-8")]
    return rows

def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True); ap.add_argument("--data", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--idle", type=int, default=30); ap.add_argument("--window", type=int, default=60)
    ap.add_argument("--probe-only", action="store_true"); ap.add_argument("--probe-n", type=int, default=200)
    ap.add_argument("--minutes", type=float, default=10, help="sustained-load duration (one full pass of all sets at minimum)")
    a = ap.parse_args()
    sys.path.insert(0, os.path.abspath(a.model)); from runtime import CommandModel
    t0 = time.time(); m = CommandModel(os.path.abspath(a.model)); load_s = time.time() - t0
    rows = load_sets(a.data)
    import random; random.Random(0).shuffle(rows)   # same mixed order everywhere, so every window/pass has the same workload
    if a.probe_only:   # one process per thread setting: env vars must be set before numpy loads
        lat = []
        for _, r in rows[:20]: m.predict_raw(r["text"])
        for _, r in rows[:a.probe_n]:
            t = time.perf_counter(); m.predict_raw(r["text"]); lat.append((time.perf_counter() - t) * 1000)
        json.dump({"n": len(lat), "p50_ms": round(pct(lat, 50), 1), "p95_ms": round(pct(lat, 95), 1), "mean_ms": round(statistics.mean(lat), 1)}, open(a.out, "w")); return
    res = {"device": device_info(), "model_load_s": round(load_s, 2), "params": int(sum(v.size for v in m.W.values()))}
    print(res["device"], "load", res["model_load_s"], "s", flush=True)
    idle = [sample() for _ in range(max(1, a.idle // 5)) if not time.sleep(5)]
    res["idle"] = {"temp_c": round(statistics.mean(s["temp_c"] for s in idle), 1) if idle[0]["temp_c"] is not None else None, "clock_mhz": idle[-1]["clock_mhz"]}
    print("idle", res["idle"], flush=True)
    for _, r in rows[:20]: m.predict_raw(r["text"])   # warm-up
    smp = Sampler(10); smp.start(); t_start = time.time(); recs = []; i = 0
    while i < len(rows) or time.time() - t_start < a.minutes * 60:   # at least one full pass, then loop until --minutes
        s, r = rows[i % len(rows)]
        t = time.perf_counter(); pred = m.predict_raw(r["text"])[0]; ms = (time.perf_counter() - t) * 1000
        recs.append({"set": s, "ok": pred == r["target"], "act_ok": pred.split()[0] == r["action"], "ms": ms, "t": time.time(), "first": i < len(rows), "pass": i // len(rows) + 1})
        if i % 500 == 0: print(f"  {i} {ms:.0f} ms {smp.rows[-1]['temp_c'] if smp.rows else ''} C  {time.time()-t_start:.0f}s", flush=True)
        i += 1
    smp.stop = True; wall = time.time() - t_start
    windows = []
    for w in range(int(wall // a.window) + 1):
        lo, hi = t_start + w * a.window, t_start + (w + 1) * a.window
        rs = [x for x in recs if lo <= x["t"] < hi]; ss = [x for x in smp.rows if lo <= x["t"] < hi]
        if len(rs) < 5: continue
        temps = [x["temp_c"] for x in ss if x["temp_c"] is not None]; clk = [x["clock_mhz"] for x in ss if x["clock_mhz"]]
        windows.append({"minute": w + 1, "n": len(rs), "temp_c": round(max(temps), 1) if temps else None, "clock_mhz": round(statistics.mean(clk)) if clk else None,
                        "mean_ms": round(statistics.mean(x["ms"] for x in rs), 1), "p50_ms": round(pct([x["ms"] for x in rs], 50), 1), "p95_ms": round(pct([x["ms"] for x in rs], 95), 1),
                        "prompts_per_s": round(len(rs) / min(a.window, hi - max(lo, recs[0]["t"])), 2),
                        "throttled": sorted({x["throttled"] for x in ss if x["throttled"]}), "fan_state": sorted({x["fan_state"] for x in ss if x["fan_state"]})})
    lat = [x["ms"] for x in recs]
    res["load"] = {"prompts": len(recs), "wall_s": round(wall, 1), "p50_ms": round(pct(lat, 50), 1), "p95_ms": round(pct(lat, 95), 1), "p99_ms": round(pct(lat, 99), 1),
                   "mean_ms": round(statistics.mean(lat), 1), "max_temp_c": max((w["temp_c"] for w in windows if w["temp_c"] is not None), default=None),
                   "throttled_any": sorted({f for x in smp.rows if x["throttled"] for f in [x["throttled"]]})}
    res["windows"] = windows
    passes = []   # identical workload per pass: the clean comparison of speed as the device heats up
    for pn in sorted({x["pass"] for x in recs}):
        rs = [x for x in recs if x["pass"] == pn]; lo, hi = rs[0]["t"], rs[-1]["t"]
        ss = [x for x in smp.rows if lo - 10 <= x["t"] <= hi]; temps = [x["temp_c"] for x in ss if x["temp_c"] is not None]
        clk = [x["clock_mhz"] for x in ss if x["clock_mhz"]]; ms = [x["ms"] for x in rs]
        passes.append({"pass": pn, "n": len(rs), "complete": len(rs) == len(rows), "temp_min_c": min(temps) if temps else None, "temp_max_c": max(temps) if temps else None,
                       "clock_mhz": round(statistics.mean(clk)) if clk else None, "mean_ms": round(statistics.mean(ms), 1), "p50_ms": round(pct(ms, 50), 1),
                       "p95_ms": round(pct(ms, 95), 1), "prompts_per_s": round(len(rs) / max(hi - lo, 1e-9), 2)})
    res["passes"] = passes
    first =[x for x in recs if x["first"]]
    res["accuracy"] = {s: {"n": sum(1 for x in first if x["set"] == s), "exact": round(100 * sum(x["ok"] for x in first if x["set"] == s) / max(1, sum(1 for x in first if x["set"] == s)), 2),
                           "action": round(100 * sum(x["act_ok"] for x in first if x["set"] == s) / max(1, sum(1 for x in first if x["set"] == s)), 2)} for s in SETS}
    res["accuracy_all"] = {"n": len(first), "exact": round(100 * sum(x["ok"] for x in first) / len(first), 2), "action": round(100 * sum(x["act_ok"] for x in first) / len(first), 2)}
    res["threads"] = {k: os.environ.get(k) for k in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS")}
    json.dump(res, open(a.out, "w"), indent=1); print(json.dumps(res["load"]), flush=True)

if __name__ == "__main__": main()