#!/usr/bin/env python3
"""Read-only low-rate CPU/NPU sampling of an already running inference process."""
import argparse
import json
import os
from pathlib import Path
import re
import statistics
import time

p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--pid", type=int, required=True)
p.add_argument("--duration", type=float, default=30)
p.add_argument("--interval", type=float, default=0.1)
p.add_argument("--output", type=Path, required=True)
a = p.parse_args()
if a.duration <= 0 or a.interval <= 0:
    p.error("duration and interval must be positive")
hz = os.sysconf("SC_CLK_TCK")
base = Path(f"/proc/{a.pid}")


def ticks(path):
    text = path.read_text()
    fields = text[text.rfind(")") + 2:].split()
    return {"cpu": int(fields[11]) + int(fields[12]),
            "minflt": int(fields[7]), "majflt": int(fields[9])}


start = time.monotonic()
before = ticks(base / "stat")
thread_before = {x.name: ticks(x / "stat")["cpu"] for x in (base / "task").iterdir()}
samples = []
deadline = start
while time.monotonic() - start < a.duration:
    sample = {"elapsed_s": time.monotonic() - start, "timestamp": time.time()}
    sample["npu_load_pct"] = [int(x) for x in re.findall(
        r"Core\d+:\s*(\d+)%", Path("/sys/kernel/debug/rknpu/load").read_text())]
    samples.append(sample)
    deadline += a.interval
    time.sleep(max(0, deadline - time.monotonic()))
elapsed = time.monotonic() - start
after = ticks(base / "stat")
threads = []
for thread in (base / "task").iterdir():
    if thread.name in thread_before:
        threads.append({"tid": int(thread.name), "name": (thread / "comm").read_text().strip(),
                        "cpu_pct": 100 * (ticks(thread / "stat")["cpu"] - thread_before[thread.name]) / hz / elapsed})
valid = [x["npu_load_pct"] for x in samples if len(x["npu_load_pct"]) == 3]
summary = {"pid": a.pid, "elapsed_s": elapsed,
           "process_cpu_pct": 100 * (after["cpu"] - before["cpu"]) / hz / elapsed,
           "minor_faults": after["minflt"] - before["minflt"],
           "major_faults": after["majflt"] - before["majflt"],
           "threads": sorted(threads, key=lambda x: -x["cpu_pct"]),
           "npu_mean_pct": [statistics.mean(x[i] for x in valid) for i in range(3)] if valid else [],
           "npu_peak_pct": [max(x[i] for x in valid) for i in range(3)] if valid else [],
           "npu_all_zero_fraction": sum(not any(x) for x in valid) / len(valid) if valid else None,
           "frequencies_khz": {f"cpu{i}": Path(f"/sys/devices/system/cpu/cpufreq/policy{i}/scaling_cur_freq").read_text().strip() for i in (4, 6)},
           "npu_frequency_hz": Path("/sys/class/devfreq/fdab0000.npu/cur_freq").read_text().strip(),
           "cpu_temperature_millic": Path("/sys/class/thermal/thermal_zone0/temp").read_text().strip()}
a.output.parent.mkdir(parents=True, exist_ok=True)
a.output.write_text(json.dumps({"summary": summary, "samples": samples}, indent=2) + "\n")
print(json.dumps(summary, indent=2), flush=True)