"""Apply REAP's expert selection to an original checkpoint (any expert dtype: NVFP4, FP8, BF16).

REAP's prune.py picks, per layer, the n experts with the lowest saliency (torch.topk(saliency, n,
largest=False) on the observer data) and keeps the rest in index order. This recomputes that exact
selection from the saved observations, checks it against the router rows of REAP's own pruned
checkpoint (--reference), and writes the source checkpoint with only the kept experts: per-expert
tensors (`...layers.<l>.mlp.experts.<e>.*`; DeepSeek-V4.1 `layers.<l>.ffn.experts.<e>.*`) renumbered
0..K-1 in kept order, the router (`...mlp.gate.weight`; DeepSeek `ffn.gate.weight/bias/bias_vl`) sliced
to the kept rows, the MTP draft layers (`mtp.*`) dropped,
every other tensor and file unchanged, and num_experts (DeepSeek: n_routed_experts) = K in the
(text) config. The kept indices
go to <out>/reap_selection.json. --pad-to N appends N - K experts no token can reach (zero weights,
zero router rows, router bias -1e4; DeepSeek selects by score + bias with biases 9.6-14), for kernels
that take only some expert counts; the model computes exactly what the K-expert model does.

    external/reap/.venv/bin/python external/reap/experiments/apex/apply_selection.py \
        --observations <artifacts>/all/observations_1024_cosine-seed_42.pt --ratio 0.5 \
        --source <HF snapshot> --reference <REAP pruned dir> --out <dir>
"""

import argparse
import json
import re
import shutil
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file

# Qwen: ...layers.<l>.mlp.experts.<e>.* and ...mlp.gate.weight; DeepSeek-V4.1: layers.<l>.ffn.experts.<e>.*
# and layers.<l>.ffn.gate.{weight, bias, bias_vl} (all indexed by expert)
EXPERT = re.compile(r"^((?:.*\.)?layers\.(\d+)\.(?:mlp|ffn)\.experts\.)(\d+)(\..*)$")
ROUTER = re.compile(r"^(?:.*\.)?layers\.(\d+)\.(?:mlp|ffn)\.gate\.(?:weight|bias|bias_vl)$")
SHARD_BYTES = 5 * 2**30


def selection(observations: Path, ratio: float, metric: str) -> dict[int, list[int]]:
    """REAP prune.py's choice: per layer, prune the n lowest-saliency experts, keep the rest."""
    data = torch.load(observations, weights_only=False)
    kept = {}
    for layer, state in data.items():
        num_experts = state["expert_frequency"].shape[0]
        n = int(num_experts * ratio)
        _, pruned = torch.topk(state[metric], n, largest=False)
        pruned = set(pruned.tolist())
        kept[int(layer)] = [i for i in range(num_experts) if i not in pruned]
    return kept


def reference_routers(reference: Path) -> dict[int, torch.Tensor]:
    out = {}
    for f in sorted(reference.glob("*.safetensors")):
        with safe_open(f, "pt") as fh:
            for k in fh.keys():
                if (m := ROUTER.match(k)) and k.endswith(".weight"):
                    out[int(m.group(1))] = fh.get_tensor(k)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--observations", required=True)
    parser.add_argument("--ratio", type=float, required=True, help="share of experts pruned per layer")
    parser.add_argument("--metric", default="reap")
    parser.add_argument("--source", required=True, help="original checkpoint directory")
    parser.add_argument("--reference", default=None, help="REAP's pruned checkpoint (router check)")
    parser.add_argument("--out", required=True)
    parser.add_argument(
        "--pad-to", type=int, default=None,
        help="append never-routed experts up to this count (DeepSeek: vLLM's sqrtsoftplus top-k kernel "
        "takes 64/128/192/256/384 experts): zero weights, zero router rows, router bias -1e4",
    )
    args = parser.parse_args()

    source, out = Path(args.source), Path(args.out)
    kept = selection(Path(args.observations), args.ratio, args.metric)
    k = {len(v) for v in kept.values()}
    assert len(k) == 1, k
    k = k.pop()
    index = json.loads((source / "model.safetensors.index.json").read_text())["weight_map"]

    routers = {}
    for name, f in index.items():
        if (m := ROUTER.match(name)) and not name.startswith("mtp.") and name.endswith(".weight"):
            with safe_open(source / f, "pt") as fh:
                routers[int(m.group(1))] = fh.get_tensor(name)
    assert sorted(routers) == sorted(kept), (sorted(routers)[:5], sorted(kept)[:5])
    if args.reference:
        ref = reference_routers(Path(args.reference))
        assert sorted(ref) == sorted(kept), "reference layers differ"
        for layer, rows in kept.items():
            same = torch.equal(routers[layer][rows].to(ref[layer].dtype), ref[layer])
            assert same, f"layer {layer}: kept set differs"
        print(f"selection matches REAP's pruned checkpoint in all {len(kept)} layers")

    out.mkdir(parents=True, exist_ok=False)
    new_index, shard, size, n_shard = {}, {}, 0, 0

    def flush() -> None:
        nonlocal shard, size, n_shard
        if shard:
            n_shard += 1
            name = f"model-{n_shard:05d}.safetensors"
            save_file(shard, out / name, metadata={"format": "pt"})
            new_index.update({key: name for key in shard})
            shard, size = {}, 0

    position = {layer: {e: j for j, e in enumerate(rows)} for layer, rows in kept.items()}
    pad = 0 if args.pad_to is None else args.pad_to - k
    assert pad >= 0, (args.pad_to, k)
    templates: dict[int, dict[str, tuple[str, torch.Tensor]]] = {}  # layer -> suffix -> (prefix, tensor)
    by_file: dict[str, list[str]] = {}
    for name, f in index.items():
        by_file.setdefault(f, []).append(name)
    dropped = 0
    for f, names in sorted(by_file.items()):
        with safe_open(source / f, "pt") as fh:
            for name in sorted(names):
                if name.startswith("mtp."):
                    dropped += 1
                    continue
                if m := EXPERT.match(name):
                    layer, expert = int(m.group(2)), int(m.group(3))
                    if expert not in position[layer]:
                        dropped += 1
                        continue
                    key, tensor = f"{m.group(1)}{position[layer][expert]}{m.group(4)}", fh.get_tensor(name)
                    if pad and m.group(4) not in templates.setdefault(layer, {}):
                        templates[layer][m.group(4)] = (m.group(1), torch.zeros_like(tensor))
                elif (m := ROUTER.match(name)) and int(m.group(1)) in kept:
                    key, tensor = name, fh.get_tensor(name)[kept[int(m.group(1))]].contiguous()
                    if pad:
                        fill = -1e4 if not name.endswith(".weight") else 0.0
                        extra = torch.full((pad, *tensor.shape[1:]), fill, dtype=tensor.dtype)
                        tensor = torch.cat([tensor, extra])
                else:
                    key, tensor = name, fh.get_tensor(name)
                shard[key] = tensor
                size += tensor.numel() * tensor.element_size()
                if size >= SHARD_BYTES:
                    flush()
    for layer, suffixes in sorted(templates.items()):
        for suffix, (prefix, zero) in suffixes.items():
            for j in range(k, k + pad):
                shard[f"{prefix}{j}{suffix}"] = zero.clone()  # safetensors refuses shared storage
                size += zero.numel() * zero.element_size()
        if size >= SHARD_BYTES:
            flush()
    flush()
    total = sum((out / f).stat().st_size for f in set(new_index.values()))
    (out / "model.safetensors.index.json").write_text(
        json.dumps({"metadata": {"total_size": total}, "weight_map": dict(sorted(new_index.items()))}, indent=2)
    )
    config = json.loads((source / "config.json").read_text())
    text = config.get("text_config", config)
    text["n_routed_experts" if "n_routed_experts" in text else "num_experts"] = k + pad
    (out / "config.json").write_text(json.dumps(config, indent=2) + "\n")
    for f in source.iterdir():
        if f.is_file() and not f.name.endswith(".safetensors") and f.name not in (
            "model.safetensors.index.json",
            "config.json",
        ):
            shutil.copy(f, out / f.name)
    (out / "reap_selection.json").write_text(
        json.dumps(
            {
                "observations": str(Path(args.observations).resolve()),
                "ratio_pruned": args.ratio,
                "metric": args.metric,
                "source": str(source.resolve()),
                "kept_per_layer": k,
                "padding_experts": pad,  # ids k..k+pad-1: zero weights, router bias -1e4, never routed
                "kept": {str(layer): rows for layer, rows in sorted(kept.items())},
            }
        )
        + "\n"
    )
    print(f"{out}: {k} experts per layer kept, {dropped} tensors dropped, {total / 1e9:.1f} GB")


if __name__ == "__main__":
    main()