"""Decode bitplanes_k2_c0.6.npz -> full bf16 weights (bit-exact vs model.safetensors).

The true quantized artifact of circus-0.4-t9. Each target linear layer is stored as
  W = alpha ⊙ (T1 + c·T2),  T1,T2 ∈ {-1,0,+1},  c = 0.6
with per-block-32 fp32 scales alpha (GPTQ column-permuted domain; `inv` restores
the original column order). Index encoding: idx = (T1+1)*3 + (T2+1) ∈ [0,8],
two 4-bit indices per byte.

Usage:
  python decode_bitplanes.py                    # verify all layers vs model.safetensors
  python decode_bitplanes.py --tensor NAME      # decode one tensor, print stats
"""
import argparse

import numpy as np


def decode(z, name):
    packed = z[f"{name}.idx"]          # (M, N//2) uint8
    alpha = z[f"{name}.alpha"]         # (M, nB) fp32
    inv = z[f"{name}.inv"]             # (N,) int32
    M, nB = alpha.shape
    N = nB * 32
    idx = np.empty((M, N), np.uint8)
    idx[:, 0::2] = packed >> 4
    idx[:, 1::2] = packed & 0x0F
    t1 = (idx.astype(np.float32) // 3) - 1.0
    t2 = (idx % 3).astype(np.float32) - 1.0
    v = (t1 + float(z["meta.c"]) * t2).reshape(M, nB, 32)
    w = (alpha[:, :, None] * v).reshape(M, N)
    return w[:, inv]                   # un-permute columns (fp32)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--planes", default="bitplanes_k2_c0.6.npz")
    ap.add_argument("--safetensors", default="model.safetensors")
    ap.add_argument("--tensor", default=None)
    args = ap.parse_args()
    z = np.load(args.planes)
    names = sorted({k.rsplit(".", 1)[0] for k in z.files if k.endswith(".idx")})
    if args.tensor:
        w = decode(z, args.tensor)
        print(args.tensor, w.shape, "std", w.std())
        return
    import torch
    from safetensors import safe_open
    bad = 0
    with safe_open(args.safetensors, framework="pt") as f:
        for i, n in enumerate(names):
            ref = f.get_tensor(n + ".weight")
            w = torch.from_numpy(decode(z, n)).to(torch.bfloat16)
            ok = torch.equal(w, ref)
            bad += not ok
            if not ok or i % 32 == 0:
                print(f"[{i+1}/{len(names)}] {n} bit-exact={ok}")
    print(f"verified {len(names)} tensors, mismatches={bad}")


if __name__ == "__main__":
    main()