"""Layer-major prefill (SegmentedModel.prefill) == block-by-block prefill (step): same logits and same state.
usage: python test_layer_major.py MODEL_DIR   (set ONW_EXPERT_GB to test it with experts streamed from the SSD)"""
import sys, time
import numpy as np
import openvino as ov
from onw.runtime import SegmentedModel


def main():
    m = SegmentedModel(ov.Core(), sys.argv[1], "NPU", {"CACHE_DIR": sys.argv[1] + "/npu_cache"})
    ids = list(np.random.default_rng(0).integers(1000, 30000, 83))       # 5 full blocks + a 3-token tail
    m.reset()
    t = time.time()
    a = m.prefill(ids)
    ta, sa = time.time() - t, {k: v.copy() for k, v in m.state.items()}
    ca = m.cur
    m.reset()
    t, b, s = time.time(), None, 0
    while s < len(ids):
        k = min(m.S, len(ids) - s)
        b = m.step(ids[s:s + k], want_logits=s + k >= len(ids))
        s += k
    tb = time.time() - t
    err = float(np.abs(a.astype(np.float32) - b.astype(np.float32)).max())
    serr = max(float(np.abs(sa[k].astype(np.float32)[..., :ca, :] if k.startswith("past_") else sa[k].astype(np.float32)
                            - (m.state[k].astype(np.float32)[..., :ca, :] if k.startswith("past_") else m.state[k].astype(np.float32))).max())
               for k in sa if not k.startswith("past_"))
    kv = max(float(np.abs(sa[k][:, :, :ca].astype(np.float32) - m.state[k][:, :, :ca].astype(np.float32)).max())
             for k in sa if k.startswith("past_"))
    print(f"{len(ids)} tokens: layer-major {ta:.2f}s, block-major {tb:.2f}s | cur {ca} vs {m.cur} | "
          f"max |logit diff| {err:.4f}, argmax equal {np.array_equal(a.argmax(-1), b.argmax(-1))}, "
          f"max state diff {serr:.4f}, max KV diff {kv:.4f}")
    if hasattr(m.bank, "stats"):
        print("cache:", m.bank.stats())


if __name__ == "__main__":
    main()