#!/usr/bin/env python3
"""
量化"Hexagon NPU 推理时 CPU 占用"到底是多少。

方法：用 resource.getrusage(RUSAGE_CHILDREN) 取子进程 utime+stime，
      除以墙钟时间得到「占用了几个核」，再除以推理张数得到「每帧 CPU 时间」。
"""
import os
import resource
import subprocess
import sys
import time

SDK = os.path.expanduser("~/QAIRT/qairt/2.42.0.251225")
NPROC = os.cpu_count()


def run_measured(name, cmd, env, n_infer, note=""):
    before = resource.getrusage(resource.RUSAGE_CHILDREN)
    t0 = time.monotonic()
    r = subprocess.run(cmd, env=env, capture_output=True, text=True)
    wall = time.monotonic() - t0
    after = resource.getrusage(resource.RUSAGE_CHILDREN)

    cpu = (after.ru_utime - before.ru_utime) + (after.ru_stime - before.ru_stime)
    cores = cpu / wall if wall else 0
    per_inf_ms = cpu / n_infer * 1000 if n_infer else 0

    print(f"\n### {name}")
    if note:
        print(f"    {note}")
    print(f"    退出码      : {r.returncode}")
    print(f"    墙钟时间    : {wall:.2f} s   ({n_infer} 张 → {wall / n_infer * 1000:.2f} ms/张)")
    print(f"    子进程 CPU  : {cpu:.2f} s (user {after.ru_utime - before.ru_utime:.2f} + sys {after.ru_stime - before.ru_stime:.2f})")
    print(f"    平均占用    : {cores:.2f} 核 / {NPROC} 核  = {cores / NPROC * 100:.1f}% 总 CPU")
    print(f"    每帧 CPU    : {per_inf_ms:.2f} ms")
    if r.returncode != 0:
        print("    stderr:", (r.stderr or r.stdout)[-300:])
    return dict(name=name, wall=wall, cpu=cpu, cores=cores, per_inf=per_inf_ms, rc=r.returncode)


def make_env():
    env = os.environ.copy()
    env["LD_LIBRARY_PATH"] = f"{SDK}/lib/aarch64-oe-linux-gcc11.2"
    env["ADSP_LIBRARY_PATH"] = f"{SDK}/lib/hexagon-v68/unsigned;/dsp"
    return env


def main():
    env = make_env()
    work = os.path.expanduser("~/QAIRT/aihub_work")
    res_dir = os.path.expanduser("~/resnet50_qairt/6490")
    results = []

    # ---------- 1. 空闲基线 ----------
    b = resource.getrusage(resource.RUSAGE_CHILDREN)
    t0 = time.monotonic()
    time.sleep(5)
    wall = time.monotonic() - t0
    a = resource.getrusage(resource.RUSAGE_CHILDREN)
    cpu = (a.ru_utime - b.ru_utime) + (a.ru_stime - b.ru_stime)
    print(f"### 空闲基线\n    5 s 内子进程 CPU: {cpu:.2f} s")

    # ---------- 2. ResNet50 @ NPU (context binary) ----------
    big_list = "/tmp/bench_resnet500.txt"
    with open(os.path.join(res_dir, "test_list.txt")) as f:
        lines = [l.strip() for l in f if l.strip()]
    # test_list.txt 里是相对路径（../data/...），这里转成绝对路径
    lines = [os.path.normpath(os.path.join(res_dir, l)) for l in lines]
    with open(big_list, "w") as f:
        for _ in range(125):
            f.write("\n".join(lines) + "\n")

    results.append(run_measured(
        "ResNet50 INT8 @ NPU (qnn-net-run + context binary)",
        [f"{SDK}/bin/aarch64-oe-linux-gcc11.2/qnn-net-run",
         "--backend", f"{SDK}/lib/aarch64-oe-linux-gcc11.2/libQnnHtp.so",
         "--retrieve_context", os.path.join(res_dir, "resnet50_aimet_quantized_6490.bin"),
         "--input_list", big_list,
         "--output_dir", "/tmp/bench_out_resnet"],
        env, 500, "HTP v68，纯 NPU 推理"))

    # ---------- 3. MobileNetV2 INT8 @ NPU (TFLite delegate) ----------
    mb_list = os.path.join(work, "list_mbv2.txt")
    one = open(mb_list).read().strip().splitlines()[0]
    multi = "/tmp/bench_mbv2_100.txt"
    with open(multi, "w") as f:
        f.write("\n".join([one] * 100) + "\n")

    results.append(run_measured(
        "MobileNetV2 INT8 @ NPU (qtld-net-run，HTP)",
        ["qtld-net-run",
         "--model", os.path.join(work, "mbv2_int8.tflite"),
         "--input", multi,
         "--output", "/tmp/bench_out_npu",
         "--backend", "htp",
         "--library_path", f"{SDK}/lib/aarch64-oe-linux-gcc11.2/libQnnHtp.so",
         "--skel_library_dir", f"{SDK}/lib/hexagon-v68/unsigned"],
        env, 100, "INT8 量化，0 个 FP16 报错 → 全程 NPU"))

    # ---------- 4. MobileNetV2 float @ CPU（对照组） ----------
    os.makedirs("/tmp/bench_out_cpu", exist_ok=True)
    results.append(run_measured(
        "MobileNetV2 float @ CPU (qtld-net-run，HTP 拒绝 FP16 后回退)",
        ["qtld-net-run",
         "--model", os.path.join(work, "mobilenet_v2.tflite"),
         "--input", multi,
         "--output", "/tmp/bench_out_cpu",
         "--backend", "htp",
         "--library_path", f"{SDK}/lib/aarch64-oe-linux-gcc11.2/libQnnHtp.so",
         "--skel_library_dir", f"{SDK}/lib/hexagon-v68/unsigned"],
        env, 100, "同一模型 float 版，算子全被 HTP 拒绝 → CPU 执行"))

    # ---------- 汇总 ----------
    print("\n" + "=" * 74)
    print(f"{'场景':<52}{'每帧CPU(ms)':>12}{'占用核数':>10}")
    print("-" * 74)
    for r in results:
        print(f"{r['name'][:50]:<52}{r['per_inf']:>12.2f}{r['cores']:>10.2f}")
    print("=" * 74)


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