"""Where does a prompt block (largest S) spend its time: NPU calls vs expert rebinding vs the rest?
usage: python prof_block.py ENGINE_DIR"""
import sys, time
import numpy as np
import openvino as ov
from onw.runtime import SegmentedModel


def main():
    core = ov.Core()
    d = sys.argv[1]
    m = SegmentedModel(core, d, "NPU", {"CACHE_DIR": d + "/npu_cache"})
    S = m.S
    acc = {"infer": 0.0, "set_tensor": 0.0, "n_set": 0}
    wrapped = []
    for req, meta, names in m.segs[S]:
        class Pxy:
            pass
        p = Pxy()
        p.model_inputs = req.model_inputs

        def st(*a, _f=req.set_tensor):
            t = time.perf_counter(); _f(*a); acc["set_tensor"] += time.perf_counter() - t; acc["n_set"] += 1

        def inf(*a, _f=req.infer, **k):
            t = time.perf_counter(); r = _f(*a, **k); acc["infer"] += time.perf_counter() - t; return r
        p.set_tensor, p.infer = st, inf
        wrapped.append((p, meta, names))
    m.segs[S] = wrapped
    rng = np.random.default_rng(0)
    m.step(list(rng.integers(1000, 50000, S)))
    for k in acc:
        acc[k] = 0
    n = 5
    t0 = time.perf_counter()
    for _ in range(n):
        m.reset()
        m.step(list(rng.integers(1000, 50000, S)))
    tot = (time.perf_counter() - t0) / n * 1000
    print(f"S={S} block: {tot:.0f} ms = NPU {acc['infer'] / n * 1000:.0f} ms + set_tensor {acc['set_tensor'] / n * 1000:.0f} ms "
          f"({acc['n_set'] // n} calls) + other {tot - (acc['infer'] + acc['set_tensor']) / n * 1000:.0f} ms")


if __name__ == "__main__":
    main()