#!/usr/bin/env python3
"""Probe indexer and compressor weight shapes from the checkpoint.
This tells us the ACTUAL dimensions, not what we assume.
Run via: fire_b200_test probe_indexer_shapes.py
"""
import json, sys
from pathlib import Path
from safetensors.torch import load_file

CHECKPOINT = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"

def main():
    cdir = Path(CHECKPOINT)
    with open(cdir / "config.json") as f:
        cfg = json.load(f)
    
    n_layers = cfg["num_hidden_layers"]
    n_ih = cfg.get("index_n_heads", 64)
    ihd = cfg.get("index_head_dim", 128)
    hd = cfg["head_dim"]
    cr = cfg.get("compress_ratios", [128] * n_layers)
    
    print(f"Config: n_ih={n_ih}, ihd={ihd}, hd={hd}")
    print(f"n_ih * ihd = {n_ih * ihd}")
    print(f"2 * ihd = {2 * ihd}")
    print(f"2 * hd = {2 * hd}")
    print(f"Compress ratios: first5={cr[:5]}")
    print()
    
    # Load weight map to find indexer weights
    idx_file = cdir / "model.safetensors.index.json"
    if idx_file.exists():
        with open(idx_file) as f:
            wmap = json.load(f).get("weight_map", {})
        
        # Find indexer/compressor weights for layer 2 (first CSA layer)
        for li in [0, 1, 2, 3]:
            pfx = f"model.layers.{li}.self_attn"
            print(f"\n=== Layer {li} (ratio={cr[li] if li < len(cr) else '?'}) ===")
            for k in sorted(wmap.keys()):
                if k.startswith(pfx) and ('compressor' in k or 'indexer' in k or 'q_b_proj' in k or 'kv_proj' in k or 'gate_proj' in k):
                    shard = cdir / wmap[k]
                    print(f"  {k} -> shard {wmap[k]}")
    else:
        print("No index file, loading all weights...")
    
    # Actually load some weights and print shapes
    # Just load the first shard to get shapes
    print("\n=== Loading weight shapes ===")
    all_w = {}
    if idx_file.exists():
        shards = set(wmap.values())
        for sn in sorted(shards):
            sf = cdir / sn
            if sf.exists():
                w = load_file(str(sf))
                # Only print relevant keys
                for k, v in w.items():
                    if ('compressor' in k or 'indexer' in k) and 'layers.2' in k:
                        print(f"  {k}: shape={list(v.shape)} dtype={v.dtype}")
                del w
    
    # Also check q_b_proj for layer 2
    print("\n=== Layer 2 attention projection shapes ===")
    for sn in sorted(shards):
        sf = cdir / sn
        if sf.exists():
            w = load_file(str(sf))
            for k, v in w.items():
                if 'layers.2.self_attn' in k and ('q_b' in k or 'kv_proj' in k or 'gate_proj' in k):
                    print(f"  {k}: shape={list(v.shape)} dtype={v.dtype}")
            del w

if __name__ == "__main__":
    main()
