#!/usr/bin/env python3
"""P19.1 — inspect gpt-oss-20b safetensors + validate MXFP4 dequant (CPU only).

Run: python3 tools/mxfp4_probe.py /path/to/gpt-oss-20b

Does three things:
  1. List the layer-0 tensor names, shapes, dtypes (esp. expert _blocks/_scales).
  2. Dequantize one expert's gate_up_proj from MXFP4 -> fp32 with our own LUT/decode.
  3. Print stats so we can eyeball sanity (range, mean, a few values).

MXFP4 (OCP microscaling) as used by gpt-oss:
  - elements are FP4 E2M1 (4-bit), packed 2-per-byte.
  - every block of 32 consecutive elements shares one E8M0 scale (8-bit exponent).
  - value = fp4_e2m1_lut[code] * 2**(scale_e8m0 - 127)
The `_blocks` tensor holds the packed 4-bit codes; `_scales` holds the per-block
E8M0 exponents (uint8). We confirm shapes line up (last dim of blocks * 2 / ... )
and that decoded values are in a sane range.
"""
import sys, json, os
import numpy as np

# FP4 E2M1 code -> value lookup (OCP MX spec). 16 codes: sign(1) exp(2) mant(1).
# Values: 0, 0.5, 1, 1.5, 2, 3, 4, 6 and their negatives.
FP4_E2M1 = np.array(
    [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
     -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
    dtype=np.float32,
)

def dequant_mxfp4(blocks: np.ndarray, scales: np.ndarray) -> np.ndarray:
    """gpt-oss layout:
      blocks: uint8 [..., nblk, 16]  -> each 16-byte row = 32 FP4 codes (2/byte)
      scales: uint8 [..., nblk]      -> one E8M0 exponent per block
    Returns fp32 [..., nblk*32] (the input/contraction dim)."""
    lo = blocks & 0x0F
    hi = (blocks >> 4) & 0x0F
    # within a byte, low nibble is element 2i, high nibble is element 2i+1
    codes = np.empty(blocks.shape[:-1] + (blocks.shape[-1] * 2,), dtype=np.uint8)
    codes[..., 0::2] = lo
    codes[..., 1::2] = hi                          # [..., nblk, 32]
    vals = FP4_E2M1[codes]                         # [..., nblk, 32]
    scale_f = np.power(2.0, scales.astype(np.float32) - 127.0)  # [..., nblk]
    out = vals * scale_f[..., None]                # [..., nblk, 32]
    return out.reshape(out.shape[:-2] + (out.shape[-2] * out.shape[-1],))  # [..., nblk*32]

def main():
    d = sys.argv[1] if len(sys.argv) > 1 else "/home/gahow/models/gpt-oss-20b"
    from safetensors import safe_open
    idx = json.load(open(os.path.join(d, "model.safetensors.index.json")))
    wm = idx["weight_map"]
    l0 = {k: v for k, v in wm.items() if "layers.0." in k}
    print("=== layer 0 tensors ===")
    # open each shard lazily
    handles = {}
    def get(name):
        shard = wm[name]
        if shard not in handles:
            handles[shard] = safe_open(os.path.join(d, shard), framework="numpy")
        return handles[shard]
    for k in sorted(l0):
        h = get(k)
        t = h.get_slice(k)
        print(f"  {k.replace('model.layers.0.','L0.')}  shape={t.get_shape()} dtype={t.get_dtype()}")

    # find expert gate_up blocks/scales
    gu_b = next((k for k in l0 if "gate_up_proj_blocks" in k), None)
    gu_s = next((k for k in l0 if "gate_up_proj_scales" in k), None)
    if gu_b and gu_s:
        print(f"\n=== dequant {gu_b.split('.')[-1]} (expert 0) ===")
        blocks = get(gu_b).get_tensor(gu_b)
        scales = get(gu_s).get_tensor(gu_s)
        print("  blocks", blocks.shape, blocks.dtype, " scales", scales.shape, scales.dtype)
        b0 = blocks[0]; s0 = scales[0]   # expert 0: blocks [5760,90,16], scales [5760,90]
        deq = dequant_mxfp4(b0.astype(np.uint8), s0.astype(np.uint8))  # -> [5760, 2880]
        print("  expert0 dequant shape", deq.shape, "(expect [5760, 2880] = [2*inter, hidden])",
              "\n  min %.4f max %.4f mean %.4f std %.4f" % (deq.min(), deq.max(), deq.mean(), deq.std()))
        print("  row0 first 8 vals:", np.round(deq[0, :8], 4))
    else:
        print("\n(no gate_up_proj_blocks/_scales found — check tensor names above)")

if __name__ == "__main__":
    main()