#!/usr/bin/env python3
"""Dry-run the quant target regexes against the real module names from the
safetensors index. Verifies: full coverage of Linear weights, zero overlap
between groups, and that everything intentionally excluded is excluded."""
import json, re, sys, collections

BASE = sys.argv[1]
G0 = [r".*self_attn\.(q|k|v|o)_proj$",
      r".*linear_attn\.(in_proj_qkv|in_proj_z|out_proj)$",
      r".*lm_head",
      r".*layers\.(56|57|58|59|60|61|62|63)\.mlp\.(gate|up|down)_proj$"]
G1 = [r".*layers\.([0-9]|[1-4][0-9]|5[0-5])\.mlp\.(gate|up|down)_proj$"]
IG = [r".*visual.*", r".*linear_attn\.(norm|in_proj_a|in_proj_b)$", r"^mtp.*"]

idx = json.load(open(BASE + "/model.safetensors.index.json"))
mods = sorted({k.rsplit(".", 1)[0] for k in idx["weight_map"] if k.endswith(".weight")})

def hit(pats, name):
    return any(re.fullmatch(p, name) for p in pats)

g0 = [m for m in mods if hit(G0, m)]
g1 = [m for m in mods if hit(G1, m)]
ig = [m for m in mods if hit(IG, m)]
overlap = sorted(set(g0) & set(g1))
covered = set(g0) | set(g1) | set(ig)
uncov = [m for m in mods if m not in covered]

print(f"total modules with .weight : {len(mods)}")
print(f"group_0 (FP8 W8A8)         : {len(g0)}")
print(f"group_1 (NVFP4 W4A4)       : {len(g1)}")
print(f"ignored                    : {len(ig)}")
print(f"OVERLAP g0&g1              : {len(overlap)}  {'<-- BUG' if overlap else 'OK'}")
if overlap:
    print("   ", overlap[:10])

# sanity: which layers landed in which MLP group
def layers_of(lst, kind):
    out = set()
    for m in lst:
        mm = re.search(r"layers\.(\d+)\.mlp\.", m)
        if mm:
            out.add(int(mm.group(1)))
    return sorted(out)

l0, l1 = layers_of(g0, "g0"), layers_of(g1, "g1")
print(f"\nMLP layers -> FP8   : {l0[:3]}..{l0[-3:] if l0 else []}  (n={len(l0)})")
print(f"MLP layers -> NVFP4 : {l1[:3]}..{l1[-3:] if l1 else []}  (n={len(l1)})")
print(f"MLP layer union covers 0-63: {sorted(set(l0)|set(l1)) == list(range(64))}")

print("\nuncovered modules (neither quantized nor explicitly ignored):", len(uncov))
buck = collections.Counter()
for m in uncov:
    if "visual" in m: buck["visual"] += 1
    elif "mtp" in m: buck["mtp"] += 1
    elif "norm" in m: buck["norm"] += 1
    elif "embed" in m: buck["embed"] += 1
    elif "linear_attn" in m: buck["linear_attn"] += 1
    else: buck["OTHER:" + m.split(".")[-1]] += 1
for k, v in buck.most_common():
    print(f"   {k}: {v}")
oth = [m for m in uncov if not any(s in m for s in ("visual", "mtp", "norm", "embed", "linear_attn"))]
if oth:
    print("   uncovered non-norm/embed sample:", oth[:12])