"""Decode speed with experts streamed from the SSD, for several prefetch settings: per answer the decode tok/s and the
time the NPU waited for expert reads per token. Each setting in a fresh process (cold cache), settings interleaved.
usage: python bench_decode.py MODEL_DIR GB K1,K2,... [ROUNDS=2]"""
import json, os, subprocess, sys

CHILD = r'''
import json, sys, time
from onw.chat import ChatEngine
e = ChatEngine(sys.argv[1], "NPU", pld=False)
qs = ["NPUとGPUの違いを、身近なたとえを使って説明してください。",
      "ある商品を定価の2割引きで買うと960円でした。定価はいくらですか？途中の式も書いてください。",
      "Explain the difference between a process and a thread, with a short Python example."]
out = []
for q in qs:
    e.checkpoint = None
    b0, st = e.model.bank.stats(), None
    toks = []
    for x in e.stream_chat([{"role": "user", "content": q}], 192):
        if isinstance(x, dict):
            st = x
        else:
            toks.append(x)
    b1 = e.model.bank.stats()
    n = st["completion_tokens"]
    out.append({"tok_s": round(st["decode_tok_s"], 2), "wait_ms_per_tok": round((b1["read_s"] - b0["read_s"]) * 1000 / n, 1),
                "hit": round(b1["hit_rate"], 3), "text": "".join(toks)[:30]})
print("RESULT " + json.dumps(out, ensure_ascii=False), flush=True)
'''


def main():
    d, gb, ks = sys.argv[1], sys.argv[2], sys.argv[3].split(",")
    rounds = int(sys.argv[4]) if len(sys.argv) > 4 else 2
    for r in range(rounds):
        for k in ks:
            env = {**os.environ, "PYTHONUTF8": "1", "ONW_EXPERT_GB": gb, "ONW_PREFETCH_K": k}
            p = subprocess.run([sys.executable, "-c", CHILD, d], capture_output=True, text=True, encoding="utf8",
                               errors="replace", env=env)
            line = next((l for l in p.stdout.splitlines() if l.startswith("RESULT ")), None)
            if not line:
                print(f"round {r} K={k}: FAILED {p.stderr[-300:]}", flush=True)
                continue
            res = json.loads(line[7:])
            print(f"round {r} K={k:>2}: " + " | ".join(f"{x['tok_s']:5.2f} tok/s, wait {x['wait_ms_per_tok']:5.1f} ms/tok"
                                                    for x in res) + f" | {res[0]['text']!r}", flush=True)


if __name__ == "__main__":
    main()