"""Model config for a REAP-pruned checkpoint (apply_selection.py output): the base model's config
with pretrained = the checkpoint directory, revision = SHA-1 of its reap_selection.json (a local
directory ignores the revision; the hash ties every run to the selection) and run_name = --name.
Everything else (serving, sampling, chat template) stays the base config's.

    cd apex-moe && .venv/bin/python external/reap/experiments/apex/make_config.py \
        --base configs/models/qwen3.6-35b-a3b-nvfp4.yaml --checkpoint data/reap-official/ckpt/<dir> \
        --name qwen3.6-35b-a3b-nvfp4-reapofficial-keep50
"""

import argparse
import hashlib
import json
import re
from pathlib import Path

parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
parser.add_argument("--base", required=True)
parser.add_argument("--checkpoint", required=True)
parser.add_argument("--name", required=True)
args = parser.parse_args()

ckpt = Path(args.checkpoint)
selection = (ckpt / "reap_selection.json").read_bytes()
digest = hashlib.sha1(selection).hexdigest()
meta = json.loads(selection)
text = Path(args.base).read_text()
header = (
    f"# REAP-pruned (official REAP, apex port: external/reap) {Path(args.base).name}: "
    f"{meta['kept_per_layer']} experts per layer kept\n"
    f"# ({100 * meta['ratio_pruned']:.0f}% pruned by {meta['metric']} saliency on apex-moe's calibration prompts\n"
    "# with their datasets' answers, REAP style); revision = SHA-1 of the checkpoint's reap_selection.json\n"
)
if meta.get("padding_experts"):
    header += (
        f"# plus {meta['padding_experts']} padding experts no token can reach (zero weights, router bias -1e4),\n"
        "# for vLLM's sqrtsoftplus top-k kernel (64/128/192/256/384 experts)\n"
    )
text, n1 = re.subn(r"(?m)^(  pretrained:\s*)\S+", rf"\g<1>{ckpt}", text)
text, n2 = re.subn(r"(?m)^(run_name:\s*)\S+", rf"\g<1>{args.name}", text)
assert n1 == 1 and n2 == 1, (n1, n2)
if re.search(r"(?m)^  revision:", text):
    text = re.sub(r"(?m)^(  revision:\s*)\S+", rf"\g<1>{digest}", text)
else:
    text = re.sub(r"(?m)^(  pretrained:.*\n)", rf"\g<1>  revision: {digest}\n", text, count=1)
out = Path("configs/models") / f"{args.name}.yaml"
out.write_text(header + text)
print(f"{out}: pretrained {ckpt}, revision {digest}")