import sys
import torch

if len(sys.argv) < 2:
    print("用法: python diagnose_codebook.py result/seperated/sep-XXX.pth")
    sys.exit(1)

ckpt_path = sys.argv[1]
ckpt = torch.load(ckpt_path, map_location="cpu")

names = ["quantizer1", "quantizer2", "quantizer3"]
labels = ["q1(stage1墙)", "q2(stage2门/怪/入口)", "q3(stage3资源)"]

print(f"Checkpoint: {ckpt_path}")
print(f"Epoch: {ckpt.get('epoch', '?')}")
print()
print(f"{'名称':<20s} {'存活/总量':>10s} {'利用率':>8s} {'perplexity':>12s} {'最近大小中位':>14s}")
print("-" * 68)

for name, label in zip(names, labels):
    sd = ckpt.get(name)
    if sd is None:
        print(f"{label:<20s} {'未找到':>10s}")
        continue
    cs = sd["ema_cluster_size"]
    K = cs.numel()
    alive = int((cs > 1.0).sum())
    p = cs / cs.sum()
    ppl = float(torch.exp(-(p * torch.log(p.clamp_min(1e-10))).sum()))
    usage = alive / K * 100

    # 最近（cluster_size 中位数，反映码字被使用频率）
    median_size = float(cs.median())

    print(f"{label:<20s} {alive:>3d}/{K:<3d}   {usage:>5.1f}%  {ppl:>8.2f}/{K:<8d} {median_size:>10.4f}")

# 检查是否有码字范数为零的死码
print()
print("--- 死码检查（范数=0 的码字） ---")
for name, label in zip(names, labels):
    sd = ckpt.get(name)
    if sd is None:
        continue
    w = sd["codebook.weight"]
    if w is None:
        w = sd.get("weight")  # 可能的备用 key
    if w is None:
        print(f"{label}: 无法读取 codebook.weight")
        continue
    norms = w.norm(dim=1)
    dead = int((norms < 1e-6).sum().item())
    if dead > 0:
        print(f"{label}: {dead} / {w.size(0)} 死码（范数=0）")
    else:
        print(f"{label}: 全部码字有非零范数")