"""Check a hand-off bundle with nothing but numpy: python verify.py [dir]"""
import json
import sys

import numpy as np

root = sys.argv[1] if len(sys.argv) > 1 else "."
manifest = json.load(open(f"{root}/manifest.json"))
npz = np.load(f"{root}/weights.npz")

print(f"{manifest['model']}  pattern={manifest['pattern']}  "
      f"axis={manifest['group_axis']}  N_inference={manifest['N_inference']}")
print(f"{'matrix':<20}{'M':>6}{'K':>6}{'sparsity':>10}{'pattern':>10}{'ref':>10}")
print("-" * 62)

bad = 0
for m in manifest["matrices"]:
    name, M, K = m["name"], m["M"], m["K"]
    w = npz[f"{name}.weight"]
    mask = npz[f"{name}.mask"]

    assert w.shape == (M, K), f"{name}: shape {w.shape} != {(M, K)}"
    assert np.all(w[mask == 0] == 0), f"{name}: nonzero weight under a zero mask"
    assert int((mask != 0).sum()) == m["nonzero"], f"{name}: nonzero count"

    # Structural check: for an N:M pattern every full group of M columns along
    # K must hold at most N nonzeros. This is the contract a generated kernel
    # relies on, so verify it against the shipped array, not the manifest.
    ok = "n/a"
    pat = m["pattern"]
    if pat.startswith("1x"):
        # Block pattern: every full block of B columns along K must be kept or
        # dropped whole. Partial blocks would break a block-packed kernel.
        block = int(pat[2:].split(":")[0])
        nb = K // block
        if nb:
            counts = (w[:, :nb * block] != 0).reshape(M, nb, block).sum(-1)
            viol = int(((counts != 0) & (counts != block)).sum())
            bad += viol
            ok = "OK" if viol == 0 else f"{viol} BAD"
    elif ":" in pat and not pat.startswith("unstructured"):
        n, g = (int(v) for v in pat.split(":"))
        ng = K // g
        if ng:
            counts = (w[:, :ng * g] != 0).reshape(M, ng, g).sum(-1)
            viol = int((counts > n).sum())
            bad += viol
            ok = "OK" if viol == 0 else f"{viol} BAD"

    ref = "-"
    if f"{name}.ref_x" in npz.files:
        y = w @ npz[f"{name}.ref_x"]
        if f"{name}.bias" in npz.files:
            y = y + npz[f"{name}.bias"]
        err = np.max(np.abs(y - npz[f"{name}.ref_y"]))
        ref = f"{err:.2e}"
        if err > 1e-3:
            bad += 1
            ref += " BAD"

    print(f"{name:<20}{M:>6}{K:>6}{m['sparsity']:>10.4f}{ok:>10}{ref:>10}")

print("-" * 62)
print("FAILED" if bad else "all checks passed")
sys.exit(1 if bad else 0)