#!/usr/bin/env python3
"""
pull_one_expert.py
==================

Step 3 of the Kimi K3 Atlas pipeline: pull ONE real expert out of the 1.56 TB
checkpoint without downloading the checkpoint.

The idea in one paragraph
-------------------------
A .safetensors file is not a sealed blob. It starts with a small table of
contents in plain JSON that lists every tensor inside it, its shape, and the
exact byte positions where its data lives. Web servers can serve a slice of a
file if you ask politely (an HTTP "Range" request). So we read the table of
contents, look up where one expert sits, and ask the server for just those
bytes. We download about 17 MB instead of 16 GB.

Then we have to decode. Kimi K3's routed experts are stored in MXFP4: each
number takes 4 bits, and every group of 32 numbers shares one scale byte.
Decoding is a 16-entry lookup multiplied by that group's scale. That is the
whole trick, and this file implements it and proves it works.

Usage
-----
    python pull_one_expert.py --selftest          # no network, verifies the decoder
    python pull_one_expert.py --list              # show tensor names for one expert
    python pull_one_expert.py --layer 47 --expert 312

Requires: numpy, requests, huggingface_hub.  Pillow is optional (for the PNG).
    pip install numpy requests huggingface_hub pillow
"""

import argparse, json, os, struct, sys
import numpy as np

REPO      = "moonshotai/Kimi-K3"
BASE      = f"https://huggingface.co/{REPO}/resolve/main"
INDEX     = "model.safetensors.index.json"
CACHE     = os.path.expanduser("~/.cache/kimi-k3-atlas")
GROUP     = 32          # config.json: quantization_config.group_size
BITS      = 4           # config.json: num_bits


# ---------------------------------------------------------------------------
# 1. The MXFP4 decoder
# ---------------------------------------------------------------------------
# Each 4-bit code is E2M1: 1 sign bit, 2 exponent bits, 1 mantissa bit.
# With an exponent bias of 1 that yields exactly sixteen possible values.
# There is no rounding to guess at: this table IS the format.
E2M1 = np.array([
    0.0,  0.5,  1.0,  1.5,  2.0,  3.0,  4.0,  6.0,      # sign bit 0
   -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,      # sign bit 1
], dtype=np.float32)


def unpack_nibbles(packed: np.ndarray, low_first: bool = True) -> np.ndarray:
    """Split each byte into its two 4-bit codes.

    `low_first` says which of the two numbers in a byte comes first. Both
    conventions exist in the wild, so it is a switch rather than an assumption;
    `detect_nibble_order` below picks the right one from the data itself.
    """
    lo = packed & 0x0F
    hi = (packed >> 4) & 0x0F
    pair = (lo, hi) if low_first else (hi, lo)
    return np.stack(pair, axis=-1).reshape(*packed.shape[:-1], -1)


def decode_mxfp4(packed: np.ndarray, scales: np.ndarray,
                 low_first: bool = True) -> np.ndarray:
    """Turn packed 4-bit codes plus E8M0 group scales into real numbers.

    packed : uint8, two codes per byte, last axis is along the row
    scales : uint8, one per group of 32 numbers, E8M0 (value = 2**(byte-127))
    """
    codes = unpack_nibbles(packed, low_first)          # -> 0..15
    vals  = E2M1[codes]                                # -> the sixteen values

    n_groups = scales.shape[-1]
    if vals.shape[-1] != n_groups * GROUP:
        raise ValueError(
            f"{vals.shape[-1]} values do not divide into {n_groups} groups of {GROUP}. "
            "The packing convention or the group size is different than assumed."
        )

    # E8M0 is exponent-only: the stored byte is a power of two, offset by 127.
    # 255 is reserved for NaN. Using float64 for the exponent avoids overflow
    # at the extremes before we cast back down.
    exp = scales.astype(np.int16) - 127
    scale = np.where(scales == 255, np.nan, np.exp2(exp.astype(np.float64)))

    grouped = vals.reshape(*vals.shape[:-1], n_groups, GROUP)
    out = grouped * scale[..., :, None]
    return out.reshape(*vals.shape).astype(np.float32)


def detect_nibble_order(packed: np.ndarray, scales: np.ndarray):
    """Decide low-nibble-first vs high-nibble-first from the data.

    Trained weights are dense near zero and roughly symmetric. The wrong
    nibble order shuffles which value pairs with which group scale, which
    inflates the spread. So we decode both ways and keep the tamer one.
    This is a heuristic; --list plus a known-good reference beats it if you
    have one.
    """
    a = decode_mxfp4(packed, scales, True)
    b = decode_mxfp4(packed, scales, False)
    ka, kb = np.nanstd(a), np.nanstd(b)
    return (True, ka, kb) if ka <= kb else (False, ka, kb)


# ---------------------------------------------------------------------------
# 2. Reading a safetensors table of contents over the network
# ---------------------------------------------------------------------------
def http_range(url: str, start: int, end_inclusive: int) -> bytes:
    import requests
    r = requests.get(url, headers={"Range": f"bytes={start}-{end_inclusive}"}, timeout=120)
    if r.status_code not in (200, 206):
        raise RuntimeError(f"server refused the range request: HTTP {r.status_code}")
    if r.status_code == 200 and len(r.content) > (end_inclusive - start + 1) * 4:
        raise RuntimeError("server ignored Range and sent the whole file; aborting")
    return r.content


def read_header(url: str):
    """A safetensors file opens with 8 bytes of length, then that much JSON."""
    n = struct.unpack("<Q", http_range(url, 0, 7))[0]
    if n > 400_000_000:
        raise RuntimeError(f"header claims {n} bytes, which is not plausible")
    raw = http_range(url, 8, 8 + n - 1)
    return json.loads(raw), 8 + n     # header dict, and where the data starts


def load_index():
    """The index maps every tensor name to the shard file holding it.

    It is about 60 MB, so it is fetched once and cached. This is the only
    large download in the whole step.
    """
    os.makedirs(CACHE, exist_ok=True)
    local = os.path.join(CACHE, INDEX)
    if not os.path.exists(local):
        print(f"  fetching the tensor index once (~60 MB) into {CACHE} ...")
        from huggingface_hub import hf_hub_download
        p = hf_hub_download(repo_id=REPO, filename=INDEX)
        import shutil; shutil.copy(p, local)
    with open(local) as f:
        return json.load(f)["weight_map"]


def find_expert_tensors(weight_map, layer: int, expert: int):
    """Locate the tensors for one expert without hardcoding a naming scheme.

    Different releases name things differently, so we search for anything
    mentioning this layer and this expert rather than guessing the exact path.
    """
    want = (f".{layer}.", f".{expert}.")
    hits = {
        name: shard for name, shard in weight_map.items()
        if "expert" in name and want[0] in name and want[1] in name
    }
    return hits


# ---------------------------------------------------------------------------
# 3. Reduce a matrix to something a browser can actually draw
# ---------------------------------------------------------------------------
def thumbnail(matrix: np.ndarray, size: int = 64) -> np.ndarray:
    """Average magnitude over a size x size grid of blocks.

    This is the honest reduction: 11 million numbers become 4,096. We keep
    magnitude because that is what "how much does this weight matter" means.
    """
    h, w = matrix.shape
    bh, bw = max(1, h // size), max(1, w // size)
    m = np.abs(matrix[: bh * size, : bw * size])
    return m.reshape(size, bh, size, bw).mean(axis=(1, 3))


def summarise(name: str, m: np.ndarray) -> dict:
    finite = m[np.isfinite(m)]
    return {
        "tensor": name,
        "shape": list(m.shape),
        "count": int(m.size),
        "mean": float(finite.mean()),
        "std": float(finite.std()),
        "min": float(finite.min()),
        "max": float(finite.max()),
        "mean_abs": float(np.abs(finite).mean()),
        "pct_zero": float((finite == 0).mean() * 100),
        "pct_positive": float((finite > 0).mean() * 100),
    }


# ---------------------------------------------------------------------------
# 4. Self-test: prove the decoder before spending any bandwidth
# ---------------------------------------------------------------------------
def selftest() -> bool:
    print("MXFP4 decoder self-test")
    print("-" * 58)
    ok = True

    # (a) every one of the sixteen codes, at scale 2^0, must come back exactly.
    codes = np.arange(16, dtype=np.uint8)
    byte8 = (codes[1::2] << 4 | codes[0::2]).astype(np.uint8)       # 8 bytes = 16 values
    packed = np.tile(byte8, 2)                                      # 16 bytes = 32 values = 1 group
    scales = np.array([127], dtype=np.uint8)                        # 2**(127-127) = 1
    got = decode_mxfp4(packed[None, :], scales[None, :])[0][:16]
    exp = E2M1
    if np.allclose(got, exp, equal_nan=True):
        print("  [ok] all 16 E2M1 codes decode exactly at scale 1")
    else:
        ok = False
        print("  [FAIL] code table mismatch")
        print("     expected", exp); print("     got     ", got)

    # (b) the group scale must multiply cleanly.
    for byte, factor in [(127, 1.0), (128, 2.0), (126, 0.5), (120, 2**-7), (134, 2**7)]:
        s = np.array([byte], dtype=np.uint8)
        v = decode_mxfp4(packed[None, :], s[None, :])[0][:8]
        want = E2M1[:8] * factor
        if not np.allclose(v, want):
            ok = False
            print(f"  [FAIL] scale byte {byte} should multiply by {factor}")
            break
    else:
        print("  [ok] E8M0 group scales multiply correctly across 7 orders of magnitude")

    # (c) two groups with different scales must not bleed into each other.
    p2 = np.tile(packed, 2)                                   # 32 bytes = 64 values = 2 groups
    s2 = np.array([127, 130], dtype=np.uint8)                 # x1 then x8
    v2 = decode_mxfp4(p2[None, :], s2[None, :])[0]
    if np.allclose(v2[:16], E2M1) and np.allclose(v2[32:48], E2M1 * 8):
        print("  [ok] group boundaries are respected (no scale bleeding)")
    else:
        ok = False
        print("  [FAIL] scales are being applied to the wrong group")

    # (d) NaN handling for the reserved byte.
    sn = np.array([255], dtype=np.uint8)
    if np.all(np.isnan(decode_mxfp4(packed[None, :], sn[None, :])[0])):
        print("  [ok] reserved scale byte 255 yields NaN")
    else:
        ok = False; print("  [FAIL] byte 255 should be NaN")

    # (e) the reduction keeps the size promise.
    rng = np.random.default_rng(0)
    fake = rng.normal(0, 0.02, (3072, 3584)).astype(np.float32)
    t = thumbnail(fake)
    ratio = fake.size / t.size
    print(f"  [ok] thumbnail reduces {fake.size:,} numbers to {t.size:,} ({ratio:,.0f} to 1)")

    # (f) a round trip on realistic data: quantise, decode, compare.
    vals = rng.normal(0, 1.0, 32).astype(np.float32)
    gmax = np.abs(vals).max()
    # The scale must be rounded UP, so the largest value in the group still
    # fits under 6.0, which is the biggest number E2M1 can represent. Rounding
    # down lets the peaks clip, and the error explodes. This is the single
    # easiest mistake to make when writing an MXFP4 quantiser.
    e = int(np.ceil(np.log2(gmax / 6.0))) + 127
    sc = np.exp2(e - 127)
    idx = np.abs(E2M1[:8][None, :] - np.abs(vals)[:, None] / sc).argmin(axis=1)
    sign = (vals < 0).astype(np.uint8) * 8
    q = (idx + sign).astype(np.uint8)
    pk = (q[1::2] << 4 | q[0::2]).astype(np.uint8)
    back = decode_mxfp4(pk[None, :], np.array([[e]], dtype=np.uint8))[0]
    err = np.abs(back - vals).max() / gmax
    print(f"  [ok] quantise -> decode round trip, worst error {err:.1%} of group max")
    if err > 0.30:
        ok = False; print("  [FAIL] round-trip error is implausibly large")

    print("-" * 58)
    print("PASS: the decoder is correct." if ok else "FAIL: do not run this on real data yet.")
    return ok


# ---------------------------------------------------------------------------
# 5. Main
# ---------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser(description="Pull and decode one Kimi K3 expert.")
    ap.add_argument("--selftest", action="store_true", help="verify the decoder, no network")
    ap.add_argument("--list", action="store_true", help="list the tensors found for one expert")
    ap.add_argument("--layer", type=int, default=47)
    ap.add_argument("--expert", type=int, default=312)
    ap.add_argument("--out", default="expert_out")
    ap.add_argument("--nibble", choices=["auto", "low", "high"], default="auto")
    a = ap.parse_args()

    if a.selftest:
        sys.exit(0 if selftest() else 1)

    print(f"Kimi K3 | layer {a.layer} | expert {a.expert}")
    print("Reading the tensor index ...")
    wm = load_index()
    hits = find_expert_tensors(wm, a.layer, a.expert)
    if not hits:
        print("No tensors matched. Names may differ from what was assumed.")
        sample = [n for n in list(wm)[:400] if "expert" in n][:10]
        print("A few expert tensor names from the index, to adapt the matcher:")
        for s in sample: print("   ", s)
        sys.exit(1)

    print(f"Found {len(hits)} tensors:")
    for n, shard in sorted(hits.items()):
        print(f"   {n}   ->   {shard}")
    if a.list:
        return

    os.makedirs(a.out, exist_ok=True)
    headers, report = {}, []

    # group tensors by their base name, so packed data and scales stay together
    bases = sorted({n.rsplit(".", 1)[0] for n in hits})
    for base in bases:
        pk_name = next((n for n in hits if n.startswith(base) and "packed" in n), None)
        sc_name = next((n for n in hits if n.startswith(base) and "scale" in n), None)
        if not pk_name or not sc_name:
            print(f"  skipping {base}: needs both a packed tensor and a scale tensor")
            continue

        shard = hits[pk_name]
        url = f"{BASE}/{shard}"
        if shard not in headers:
            print(f"\nReading the table of contents of {shard} ...")
            headers[shard] = read_header(url)
        hdr, data0 = headers[shard]

        def fetch(nm):
            info = hdr[nm]
            s, e = info["data_offsets"]
            raw = http_range(url, data0 + s, data0 + e - 1)
            arr = np.frombuffer(raw, dtype=np.uint8)
            return arr.reshape(info["shape"]), info

        print(f"  {base}")
        pk, pki = fetch(pk_name)
        sc, sci = fetch(sc_name)
        mb = (pk.nbytes + sc.nbytes) / 1e6
        print(f"    downloaded {mb:.1f} MB  (packed {pki['shape']}, scales {sci['shape']})")

        if a.nibble == "auto":
            low_first, ka, kb = detect_nibble_order(pk, sc)
            print(f"    nibble order: {'low' if low_first else 'high'} first "
                  f"(spread low={ka:.4g} high={kb:.4g})")
        else:
            low_first = (a.nibble == "low")

        m = decode_mxfp4(pk, sc, low_first)
        st = summarise(base, m)
        report.append(st)
        print(f"    decoded {st['count']:,} real weights | "
              f"mean_abs {st['mean_abs']:.5f} | std {st['std']:.5f} | "
              f"{st['pct_positive']:.1f}% positive")

        th = thumbnail(m)
        np.save(os.path.join(a.out, f"{base.replace('.', '_')}_thumb.npy"), th)
        try:
            from PIL import Image
            img = (255 * th / (th.max() or 1)).astype(np.uint8)
            Image.fromarray(img).resize((256, 256), Image.NEAREST).save(
                os.path.join(a.out, f"{base.replace('.', '_')}_thumb.png"))
        except ImportError:
            pass

    with open(os.path.join(a.out, "stats.json"), "w") as f:
        json.dump({"layer": a.layer, "expert": a.expert, "tensors": report}, f, indent=2)
    print(f"\nWrote {len(report)} tensor summaries to {a.out}/stats.json")
    print("These are real numbers out of the real checkpoint.")


if __name__ == "__main__":
    main()