#!/usr/bin/env python3
"""derive-protected-experts.py — rewrite the recipe's protected-expert section
from a calibration recording. RUNS ON DANKSERVER (venv: numpy).

The law (docs/MOE-W4.md "Protected experts", restated in the recipe):
an expert whose share of its layer's down-input energy is at least 3% AND
at least three times its share of the routing is kept bf16; the count must
be EVEN on each rank's half of the layer (the four-bit GEMM takes the rest
as two equal tables by parity) — an odd half gains its next most energetic
non-qualifying expert.

Data: the calibration dir's per-expert down Grams,
gram.r<rank>.blk.L.ffn_down_exps.<e>.bin — 32-byte RADGRAM2 header
(magic, k, rank, rows, domain) then k*k f32, named with the checkpoint's
GLOBAL expert numbering (arch/common/rad_block_moe_fp8.h: ge =
expert_base + e). rows (header) is the expert's routing count; the trace
of the Gram is its down-input energy. Rank halves: 0-255, 256-511
(contiguous expert_base split at TP=2).

Output: the cast-rule lines (recipe format) on stdout, a report on stderr.

Usage: derive-protected-experts.py <calib dir> [--min-share 0.03]
       [--min-route-mult 3.0]
"""
import argparse
import glob
import os
import re
import struct
import sys

import numpy as np

NAME_RE = re.compile(r"^gram\.r(\d)\.blk\.(\d+)\.ffn_down_exps\.(\d+)\.bin$")


def read_gram(path):
    """(k, rank, rows, trace) from a RADGRAM2 file (full matrix on disk)."""
    with open(path, "rb") as f:
        hdr = f.read(32)
        magic, k, rank, rows, domain = struct.unpack("<8sIIQI", hdr[:28])
        if magic != b"RADGRAM2":
            raise SystemExit(f"not a RADGRAM2 file: {path}")
        n = k * k
        a = np.frombuffer(f.read(n * 4), dtype=np.float32)
        if a.size != n:
            raise SystemExit(f"short gram file: {path}")
        tr = float(a.reshape(k, k).diagonal().sum())
    return k, rank, int(rows), tr


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("calib_dir")
    ap.add_argument("--min-share", type=float, default=0.03)
    ap.add_argument("--min-route-mult", type=float, default=3.0)
    a = ap.parse_args()

    # layer -> global expert -> (rank, rows, energy)
    data = {}
    for path in sorted(glob.glob(os.path.join(
            a.calib_dir, "gram.r*.ffn_down_exps.*.bin"))):
        m = NAME_RE.match(os.path.basename(path))
        if not m:
            continue
        rank, layer, expert = int(m.group(1)), int(m.group(2)), int(m.group(3))
        k, _, rows, tr = read_gram(path)
        data.setdefault(layer, {})[expert] = (rank, rows, tr)

    if not data:
        raise SystemExit("no per-expert down grams found")

    lines, report = [], []
    total = 0
    for layer in sorted(data):
        experts = data[layer]
        tot_e = sum(v[2] for v in experts.values())
        tot_r = sum(v[1] for v in experts.values())
        if tot_e <= 0 or tot_r <= 0:
            report.append(f"blk.{layer}: empty recording, skipped")
            continue
        qual, pool = [], []
        for e, (rank, rows, tr) in experts.items():
            es, rs = tr / tot_e, rows / max(tot_r, 1)
            rec = (e, rank, es, rs)
            if es >= a.min_share and es >= a.min_route_mult * rs:
                qual.append(rec)
            elif es > 0:
                pool.append(rec)
        chosen = {r[0]: (r[2], r[3]) for r in qual}
        for half in (0, 1):  # rank halves by the contiguous expert_base split
            n_half = sum(1 for e in chosen if e // 256 == half)
            if n_half % 2:
                pool.sort(key=lambda r: -r[2])
                for cand in pool:
                    if cand[1] == half and cand[0] not in chosen:
                        chosen[cand[0]] = (cand[2], cand[3])
                        report.append(
                            f"blk.{layer}: half {half} odd -> added e{cand[0]} "
                            f"(energy {cand[2]*100:.2f}%)")
                        break
        for e in sorted(chosen):
            es, rs = chosen[e]
            lines.append(f"blk.{layer}.ffn_*_exps.{e}.weight"
                         f"         cast dtype=bf16")
            report.append(f"blk.{layer} e{e:3d} (rank {e // 256}): "
                          f"energy {es*100:5.2f}%, routing {rs*100:5.3f}%, "
                          f"ratio {es/max(rs,1e-12):6.1f}")
        total += len(chosen)

    print("\n".join(lines))
    print("\n".join(report), file=sys.stderr)
    print(f"total protected: {total}", file=sys.stderr)


if __name__ == "__main__":
    main()