#!/usr/bin/env python3
"""Time and energy of a decode, for research/determinism-cost.md: run a decoder as a
child, let it decode a fixed workload in rounds, and read the energy counters around
each round. Standard library only (ctypes for the counters).

    wai_det_energy.py --label L --rounds R --k K [--idle-s S] [--counter C ...]
                      [--meta JSON] -- <child ...>

The child speaks the protocol of `wai_det_cost serve` and `wai_determinism_cost.py
fserve`: it prints `ready <n>` once loaded and warmed up, then for each `go <k>` decodes
its n jobs k times and prints `done <decodes> <wall-ns>`, and blocks on its input in
between. Counters (`--counter`, repeatable):

- `task` (macOS): the child's own CPU energy, `ri_energy_nj` of `proc_pid_rusage`
  (RUSAGE_INFO_V6), with its CPU time and the energy billed to it by system services.
  Read while the child blocks, after its CPU time has stopped moving (two reads 50 ms
  apart agree), so every thread's energy has folded into the count. Scope: this process,
  CPU only (the count excludes the GPU and the neural engine).
- `ioreport-gpu` (macOS): the `GPU Energy` rail of the IOReport `Energy Model` group, the
  whole GPU. Its share of the work is attributable only where the child held at least 99%
  of the GPU time every client accumulated over the round (`accumulatedGPUTime` in the I/O
  Registry), which is recorded.
- `rapl` (Linux): powercap `intel-rapl:0` (`package-0`), the whole package.

A whole-device counter (`ioreport-gpu`, `rapl`) is read every `--sample-ms` (default 100)
through an idle bracket of `--idle-s` seconds (default 5) before each round and after the
last, and through the round. Per round it gives the idle and busy mean power and their
sample standard deviations, the intervals in which the counter advanced, the marginal
energy (busy minus the mean of the brackets on either side, times the window) and the
ratio of the marginal power to the idle standard deviation (energy-measurement §3, §4).
The output is one JSON object on standard output.
"""
import ctypes
import json
import os
import platform
import statistics
import subprocess
import sys
import threading
import time


# ---------------------------------------------------------------- macOS: task energy

class Task:
    """`proc_pid_rusage` of another process of this user (RUSAGE_INFO_V6)."""
    OFF = {"user": 16, "system": 24, "billed_energy": 264, "serviced_energy": 272, "energy_nj": 336}

    def __init__(self, pid):
        self.pid = pid
        self.lib = ctypes.CDLL("/usr/lib/libproc.dylib")
        self.lib.proc_pid_rusage.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p]
        self.lib.proc_pid_rusage.restype = ctypes.c_int
        tb = (ctypes.c_uint32 * 2)()
        ctypes.CDLL("/usr/lib/libSystem.dylib").mach_timebase_info(ctypes.byref(tb))
        self.tick_ns = tb[0] / tb[1]

    def read(self):
        buf = (ctypes.c_uint8 * 512)()
        if self.lib.proc_pid_rusage(self.pid, 6, buf) != 0:
            raise OSError("proc_pid_rusage failed")
        u = lambda o: int.from_bytes(bytes(buf[o:o + 8]), "little")
        return {"energy_nj": u(336), "cpu_ns": (u(16) + u(24)) * self.tick_ns,
                "billed_nj": u(264), "serviced_nj": u(272)}

    def settled(self):
        """A read once the CPU time has stopped moving: two reads 50 ms apart agree."""
        prev = self.read()
        for _ in range(200):
            time.sleep(0.05)
            cur = self.read()
            if cur["cpu_ns"] == prev["cpu_ns"] and cur["energy_nj"] == prev["energy_nj"]:
                return cur
            prev = cur
        raise RuntimeError("the child's CPU time never settled")


# ---------------------------------------------------------------- macOS: IOReport GPU rail

class IOReportGPU:
    UNIT = {"nJ": 1, "uJ": 1_000, "mJ": 1_000_000, "J": 1_000_000_000, "pJ": 0.001}

    def __init__(self):
        c, v = ctypes.c_void_p, ctypes
        self.cf = cf = v.CDLL("/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation")
        self.io = io = v.CDLL("/usr/lib/libIOReport.dylib")
        for f, res, args in (
            (cf.CFStringCreateWithCString, c, [c, v.c_char_p, v.c_uint32]),
            (cf.CFDictionaryGetCount, v.c_long, [c]),
            (cf.CFDictionaryCreateMutableCopy, c, [c, v.c_long, c]),
            (cf.CFDictionaryGetValue, c, [c, c]),
            (cf.CFArrayGetCount, v.c_long, [c]),
            (cf.CFArrayGetValueAtIndex, c, [c, v.c_long]),
            (cf.CFStringGetCString, v.c_bool, [c, v.c_char_p, v.c_long, v.c_uint32]),
            (cf.CFRelease, None, [c]),
            (io.IOReportCopyChannelsInGroup, c, [c, c, v.c_uint64, v.c_uint64, v.c_uint64]),
            (io.IOReportCreateSubscription, c, [c, c, v.POINTER(c), v.c_uint64, c]),
            (io.IOReportCreateSamples, c, [c, c, c]),
            (io.IOReportCreateSamplesDelta, c, [c, c, c]),
            (io.IOReportChannelGetChannelName, c, [c]),
            (io.IOReportChannelGetUnitLabel, c, [c]),
            (io.IOReportSimpleGetIntegerValue, v.c_int64, [c, v.c_int32]),
        ):
            f.restype, f.argtypes = res, args
        chans = io.IOReportCopyChannelsInGroup(self._s("Energy Model"), None, 0, 0, 0)
        self.chan = cf.CFDictionaryCreateMutableCopy(None, cf.CFDictionaryGetCount(chans), chans)
        subbed = c()
        self.sub = io.IOReportCreateSubscription(None, self.chan, ctypes.byref(subbed), 0, None)
        self.key = self._s("IOReportChannels")

    def _s(self, s):
        return self.cf.CFStringCreateWithCString(None, s.encode(), 0x08000100)

    def _str(self, ref):
        b = ctypes.create_string_buffer(256)
        return b.value.decode() if ref and self.cf.CFStringGetCString(ref, b, 256, 0x08000100) else ""

    def sample(self):
        return self.io.IOReportCreateSamples(self.sub, self.chan, None)

    def delta_nj(self, a, b):
        """The GPU Energy rail's advance from sample a to sample b, in nJ."""
        d = self.io.IOReportCreateSamplesDelta(a, b, None)
        arr = self.cf.CFDictionaryGetValue(d, self.key)
        tot = 0
        for i in range(self.cf.CFArrayGetCount(arr)):
            ch = self.cf.CFArrayGetValueAtIndex(arr, i)
            if self._str(self.io.IOReportChannelGetChannelName(ch)) == "GPU Energy":
                unit = self._str(self.io.IOReportChannelGetUnitLabel(ch)).strip()
                tot += self.io.IOReportSimpleGetIntegerValue(ch, 0) * self.UNIT[unit]
        self.cf.CFRelease(d)
        return tot

    def release(self, s):
        self.cf.CFRelease(s)


def gpu_times():
    """accumulatedGPUTime (ns) summed per process over every I/O Registry accelerator
    client, from `ioreg`: each client node's `AppUsage` entries go to the pid its
    `IOUserClientCreator` names, whichever of the two the node lists first."""
    import re
    out = subprocess.run(["ioreg", "-r", "-c", "IOAccelerator", "-l", "-w0"], capture_output=True, text=True).stdout
    per = {}
    for node in re.split(r"\n\s*[|\s]*\+-o ", out):
        m = re.search(r'"IOUserClientCreator" = "pid (\d+),', node)
        if not m:
            continue
        t = sum(int(x) for x in re.findall(r'"accumulatedGPUTime"=(\d+)', node))
        per[int(m.group(1))] = per.get(int(m.group(1)), 0) + t
    return per


# ---------------------------------------------------------------- Linux: RAPL

class Rapl:
    def __init__(self, zone="/sys/class/powercap/intel-rapl:0"):
        self.f = os.path.join(zone, "energy_uj")
        self.name = open(os.path.join(zone, "name")).read().strip()
        self.wrap = int(open(os.path.join(zone, "max_energy_range_uj")).read())

    def read_uj(self):
        return int(open(self.f).read())

    def delta_nj(self, a, b):
        d = b - a
        if d < 0:
            d += self.wrap
        return d * 1000


# ---------------------------------------------------------------- sampling

class Sampler:
    """Reads one whole-device counter every period: (monotonic ns, interval nJ) pairs."""

    def __init__(self, kind, period):
        self.kind, self.period = kind, period
        self.dev = IOReportGPU() if kind == "ioreport-gpu" else Rapl()
        self.lock = threading.Lock()
        self.samples = []
        self.stop = False
        self.t = threading.Thread(target=self.loop, daemon=True)
        self.t.start()

    def _read(self):
        return self.dev.sample() if self.kind == "ioreport-gpu" else self.dev.read_uj()

    def loop(self):
        prev, tp = self._read(), time.monotonic_ns()
        while not self.stop:
            time.sleep(self.period)
            cur, tc = self._read(), time.monotonic_ns()
            nj = self.dev.delta_nj(prev, cur)
            if self.kind == "ioreport-gpu":
                self.dev.release(prev)
            with self.lock:
                self.samples.append((tp, tc, nj))
            prev, tp = cur, tc

    def between(self, t0, t1):
        """The intervals that lie inside [t0, t1]."""
        with self.lock:
            return [s for s in self.samples if s[0] >= t0 and s[1] <= t1]

    def covering(self, t0, t1):
        """The intervals that overlap [t0, t1]: the window's energy, its edges included."""
        with self.lock:
            return [s for s in self.samples if s[1] > t0 and s[0] < t1]


def bracket(iv):
    """Mean and sample standard deviation of the power (W) over intervals, and the
    liveness facts of energy-measurement §3.1."""
    if len(iv) < 2:
        return None
    p = [nj / (t1 - t0) for (t0, t1, nj) in iv]  # nJ/ns = W
    adv = [nj > 0 for (_, _, nj) in iv]
    run = best = 0
    for a in adv:
        run = run + 1 if a else 0
        best = max(best, run)
    return {"intervals": len(iv), "advanced": sum(adv), "longest_run": best,
            "span_ns": iv[-1][1] - iv[0][0], "mean_w": statistics.mean(p), "sd_w": statistics.stdev(p)}


def verdict(ratio):
    if ratio is None:
        return None
    return "clean" if ratio >= 10 else "relative-only" if ratio >= 3 else "noisy" if ratio >= 1 else "unusable"


def main(argv):
    a = argv[1:]
    sep = a.index("--")
    opts, child = a[:sep], a[sep + 1:]
    def opt(name, default=None, many=False):
        vals = [opts[i + 1] for i, o in enumerate(opts) if o == name]
        return vals if many else (vals[-1] if vals else default)
    label, rounds, k = opt("--label"), int(opt("--rounds", "7")), int(opt("--k", "1"))
    idle_s, period = float(opt("--idle-s", "5")), float(opt("--sample-ms", "100")) / 1000
    counters = opt("--counter", many=True)
    samplers = {c: Sampler(c, period) for c in counters if c in ("ioreport-gpu", "rapl")}
    p = subprocess.Popen(child, stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True, bufsize=1)
    line = p.stdout.readline().split()
    if not line or line[0] != "ready":
        sys.exit(f"the child did not get ready: {line}")
    jobs = int(line[1])
    task = Task(p.pid) if "task" in counters else None
    out = {"label": label, "meta": json.loads(opt("--meta", "{}")), "rounds": [], "jobs": jobs, "k": k,
           "counters": counters, "idle_s": idle_s, "sample_s": period,
           "host": {"system": platform.system(), "machine": platform.machine(), "node_cpu": os.cpu_count()},
           "child": child}
    def idle():
        t0 = time.monotonic_ns()
        time.sleep(idle_s)
        return t0, time.monotonic_ns()
    brackets = [idle()]
    for r in range(rounds):
        rec = {"load_avg_before": os.getloadavg()}
        t_before = task.settled() if task else None
        g_before = gpu_times() if "ioreport-gpu" in counters else None
        t0 = time.monotonic_ns()
        p.stdin.write(f"go {k}\n")
        p.stdin.flush()
        done = p.stdout.readline().split()
        t1 = time.monotonic_ns()
        if not done or done[0] != "done":
            sys.exit(f"round {r}: the child answered {done}")
        rec.update({"decodes": int(done[1]), "child_wall_ns": int(done[2]), "window_ns": t1 - t0})
        if len(done) > 3:
            rec["child_note"] = " ".join(done[3:])
        if task:
            t_after = task.settled()
            rec["task"] = {kk: t_after[kk] - t_before[kk] for kk in t_after}
        if g_before is not None:
            g_after = gpu_times()
            dt = {pid: g_after.get(pid, 0) - g_before.get(pid, 0) for pid in set(g_after) | set(g_before)}
            total = sum(v for v in dt.values() if v > 0)
            mine = dt.get(p.pid, 0)
            rec["gpu_time"] = {"child_ns": mine, "all_clients_ns": total,
                               "share": (mine / total) if total else None,
                               "vanished": sorted(pid for pid in g_before if pid not in g_after)}
        rec["load_avg_after"] = os.getloadavg()
        brackets.append(idle())
        for c, s in samplers.items():
            time.sleep(2 * period)
            busy = bracket(s.between(t0, t1))
            cov = s.covering(t0, t1)
            pre, post = bracket(s.between(*brackets[-2])), bracket(s.between(*brackets[-1]))
            e = {"busy": busy, "idle_before": pre, "idle_after": post,
                 "window_nj": sum(nj for (_, _, nj) in cov),
                 "covered_ns": (cov[-1][1] - cov[0][0]) if cov else 0}
            if busy and pre and post:
                idle_w = (pre["mean_w"] + post["mean_w"]) / 2
                idle_sd = max(pre["sd_w"], post["sd_w"])
                marg_w = busy["mean_w"] - idle_w
                e["marginal_nj"] = e["window_nj"] - idle_w * e["covered_ns"]
                e["ratio"] = (marg_w / idle_sd) if idle_sd > 0 else None
                e["verdict"] = verdict(e["ratio"])
            rec[c] = e
        out["rounds"].append(rec)
        print(f"{label} round {r + 1}/{rounds}: {rec['decodes']} decodes in {rec['child_wall_ns'] / 1e9:.2f} s"
              + (f", task {rec['task']['energy_nj'] / 1e9:.2f} J" if task else "")
              + "".join(f", {c} marginal {rec[c].get('marginal_nj', 0) / 1e9:.2f} J ({rec[c].get('verdict')})" for c in samplers),
              file=sys.stderr, flush=True)
    p.stdin.write("exit\n")
    p.stdin.flush()
    p.wait(timeout=60)
    for s in samplers.values():
        s.stop = True
    json.dump(out, sys.stdout, indent=1)
    print()


if __name__ == "__main__":
    main(sys.argv)