"""NPU vs CPU, segment by segment, for one block of S tokens (same inputs fed to both at every segment).
usage: python probe_segments.py MODEL_DIR [S=1]"""
import sys
import numpy as np
import openvino as ov
from onw.runtime import SegmentedModel


def main():
    d = sys.argv[1]
    S = int(sys.argv[2]) if len(sys.argv) > 2 else 1
    core = ov.Core()
    cpu = SegmentedModel(core, d, "CPU", {"INFERENCE_PRECISION_HINT": "f32"}, sizes=[1, S] if S > 1 else [1])
    npu = SegmentedModel(core, d, "NPU", {"CACHE_DIR": d + "/npu_cache"}, sizes=[1, S] if S > 1 else [1])
    ids = list(range(1000, 1000 + S))
    # run CPU normally, capturing every segment's feed; replay each feed on the NPU segment
    feeds = []
    for req, m, names in cpu.segs[S]:
        orig = req.infer

        def cap(feed, *a, _o=orig, **k):
            feeds.append(dict(feed))
            return _o(feed, *a, **k)
        req.infer = cap
    cpu.step(ids)
    cpu.reset()                                          # replay from the same (zero) state on both devices
    npu.reset()
    for i, ((creq, m, names), (nreq, _, _)) in enumerate(zip(cpu.segs[S], npu.segs[S])):
        f = feeds[i]
        r_cpu = {o.get_any_name(): v.astype(np.float32) for o, v in creq.infer(f).items()}
        r_npu = {o.get_any_name(): v.astype(np.float32) for o, v in nreq.infer(f).items()}
        errs = {k: float(np.linalg.norm(r_npu[k] - r_cpu[k]) / (np.linalg.norm(r_cpu[k]) + 1e-12)) for k in r_cpu
                if r_cpu[k].dtype.kind == "f"}
        worst = sorted(errs.items(), key=lambda kv: -kv[1])[:12]
        print(f"seg {i}: " + "  ".join(f"{k} {v:.3f}" for k, v in worst), flush=True)


if __name__ == "__main__":
    main()