"""Per-segment NPU time of one decode step (S=1), plus host overhead.
usage: python prof_segments.py ENGINE_DIR"""
import sys, time
import numpy as np
import openvino as ov
from onw.runtime import SegmentedModel


def main():
    core = ov.Core()
    m = SegmentedModel(core, sys.argv[1], "NPU", {"CACHE_DIR": "D:/moe-npu/npu_cache"})
    m.step(list(range(100, 116)))
    times = np.zeros(len(m.segs[1]))
    proxies = []
    for i, (req, meta, names) in enumerate(m.segs[1]):
        class P:
            pass
        p = P()
        p.set_tensor, p.model_inputs = req.set_tensor, req.model_inputs
        def g(*a, _f=req.infer, _i=i, **k):
            t0 = time.perf_counter()
            r = _f(*a, **k)
            times[_i] += time.perf_counter() - t0
            return r
        p.infer = g
        proxies.append((p, meta, names))
    m.segs[1] = proxies
    n = 30
    t0 = time.perf_counter()
    for i in range(n):
        m.step([200 + i])
    tot = (time.perf_counter() - t0) / n * 1000
    seg = times / n * 1000
    print(f"per token {tot:.1f} ms = NPU calls {seg.sum():.1f} ms + host {tot - seg.sum():.1f} ms")
    print("segment ms:", " ".join(f"{t:.1f}" for t in seg))


if __name__ == "__main__":
    main()