#!/usr/bin/env python3
"""SC-1: checkpoint schema + NaN audit (no GPU). Usage: sc1_schema.py CKPT"""
import json, struct, os, sys
import numpy as np

CKPT = sys.argv[1]
idx = json.load(open(f"{CKPT}/model.safetensors.index.json"))
wm = idx["weight_map"]
cfg = json.load(open(f"{CKPT}/config.json"))
books = {int(k): v for k, v in
         cfg["quantization_config"]["aqlm_layer_books"].items()}
assert cfg["quantization_config"]["quant_method"] == "nvfp4_aqlm_hybrid"

# 1. byte accounting
total = 0
headers = {}
for shard in sorted(set(wm.values())):
    with open(f"{CKPT}/{shard}", "rb") as fh:
        n = struct.unpack("<Q", fh.read(8))[0]
        hdr = json.loads(fh.read(n))
    headers[shard] = hdr
    for name, info in hdr.items():
        if name == "__metadata__": continue
        b, e = info["data_offsets"]
        total += e - b
assert total == idx["metadata"]["total_size"], \
    f"byte mismatch {total} vs {idx['metadata']['total_size']}"

# 2. per-layer schema
DT_OK = {"hyb_kind": "I8", "w13_codes": "I16", "w13_codebooks": "F16",
         "w13_scales": "F16", "w2m_codes": "I16", "w2m_codebooks": "F16",
         "w2m_scales": "F16", "w2c_codes": "I16", "w2c_codebooks": "F16",
         "w2c_scales": "F16", "nvfp4_w13_packed": "U8",
         "nvfp4_w13_bscale": "F8_E4M3", "nvfp4_w13_scale2": "F32",
         "nvfp4_w2_packed": "U8", "nvfp4_w2_bscale": "F8_E4M3",
         "nvfp4_w2_scale2": "F32"}
H, I = 6144, 2048
problems = []
for li in range(3, 78):
    b = books.get(li)
    if b is None:
        problems.append(f"L{li}: missing from aqlm_layer_books"); continue
    na, nm, nc = b["n_nvfp4"], b["n_base"], b["n_cold"]
    if na + nm + nc != 256:
        problems.append(f"L{li}: counts {na}+{nm}+{nc} != 256")
    p = f"model.layers.{li}.mlp.experts"
    want = {
        "hyb_kind": [256], "w13_codes": [nm + nc, 1, 2*I, H//8],
        "w13_codebooks": [1, 65536, 8], "w13_scales": [nm + nc, 2*I],
        "w2m_codes": [nm, 2, H, I//8], "w2m_codebooks": [2, 65536, 8],
        "w2m_scales": [nm, H], "w2c_codes": [nc, 1, H, I//8],
        "w2c_codebooks": [1, 65536, 8], "w2c_scales": [nc, H],
        "nvfp4_w13_packed": [na, 2*I, H//2],
        "nvfp4_w13_bscale": [na, 2*I, H//16], "nvfp4_w13_scale2": [na, 2],
        "nvfp4_w2_packed": [na, H, I//2], "nvfp4_w2_bscale": [na, H, I//16],
        "nvfp4_w2_scale2": [na, 1],
    }
    for tname, shape in want.items():
        full = f"{p}.{tname}"
        if full not in wm:
            problems.append(f"L{li}: missing {tname}"); continue
        info = headers[wm[full]][full]
        if list(info["shape"]) != shape:
            problems.append(f"L{li}.{tname}: shape {info['shape']} != {shape}")
        okdt = DT_OK[tname]
        if info["dtype"] not in (okdt, "U8" if okdt == "F8_E4M3" else okdt):
            problems.append(f"L{li}.{tname}: dtype {info['dtype']} != {okdt}")

# 3. NaN/Inf spot check on fp tensors (sampled)
import random
rng = random.Random(0)
fp_names = [n for n in wm if n.endswith(("codebooks", "scales", "scale2"))]
for name in rng.sample(fp_names, min(120, len(fp_names))):
    shard = wm[name]; info = headers[shard][name]
    b, e = info["data_offsets"]
    # header sizes are per shard; data starts after 8+hlen
    with open(f"{CKPT}/{shard}", "rb") as fh:
        n = struct.unpack("<Q", fh.read(8))[0]
        fh.seek(8 + n + b)
        buf = fh.read(min(e - b, 1 << 20))
    dt = {"F16": np.float16, "F32": np.float32, "BF16": None}[info["dtype"]]
    if dt is None: continue
    arr = np.frombuffer(buf, dtype=dt)
    if not np.isfinite(arr.astype(np.float32)).all():
        problems.append(f"{name}: NaN/Inf detected")

if problems:
    print("SC1 FAIL"); [print(" ", p) for p in problems[:30]]; sys.exit(1)
print(f"SC1 PASS ({total/1e9:.1f} GB, {len(wm)} tensors, "
      f"{len(books)} hybrid layers)")