#!/usr/bin/env python3
"""Emit staged next-experiment commands for ASPA/STES candidates."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


ENV_PREFIX = (
    "HF_HOME=${HF_HOME:-cache/huggingface} "
    "TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE:-${HF_HOME:-cache/huggingface}} "
    "PYTORCH_ENABLE_MPS_FALLBACK=1 TOKENIZERS_PARALLELISM=false"
)


def load_json(path: Path) -> dict:
    return json.loads(path.read_text())


def gamma_label(gamma: float) -> str:
    return f"srcgamma{int(round(gamma * 1000)):03d}"


def emit_plan(source: str, edited_candidate: str, gammas: list[float]) -> str:
    lines = [
        "# ASPA/STES Next Experiments",
        "",
        "Run only one MPS model command at a time.",
        "",
        "## Build Missing Gamma Candidates",
        "",
    ]
    for gamma in gammas:
        label = gamma_label(gamma)
        out = f"outputs/qwen3.6-27b-aspa-n2-reg05-{label}"
        if Path(out).exists():
            lines.append(f"- `{out}` already exists.")
            continue
        lines.extend(
            [
                f"### Build `{label}`",
                "",
                "```bash",
                ".venv/bin/python scripts/interpolate_hf_shards.py \\",
                f"  --source {source} \\",
                f"  --candidate {edited_candidate} \\",
                f"  --alpha {gamma:.3f} \\",
                f"  --output {out}",
                "```",
                "",
            ]
        )

    lines.extend(["## Triage Commands", ""])
    for gamma in gammas:
        label = gamma_label(gamma)
        model = f"outputs/qwen3.6-27b-aspa-n2-reg05-{label}"
        lines.extend(
            [
                f"### `{label}` capability",
                "",
                "```bash",
                f"{ENV_PREFIX} .venv/bin/python scripts/qwen36_capability_probe.py \\",
                f"  --label aspa_n2_reg05_{label} \\",
                f"  --model {model} \\",
                f"  --out runs/qwen36-capability/aspa_n2_reg05_{label}.json",
                "```",
                "",
                f"### `{label}` n30 ship gate",
                "",
                "```bash",
                f"{ENV_PREFIX} .venv/bin/python scripts/qwen36_ship_gate.py \\",
                f"  --label aspa_n2_reg05_{label}_n30 \\",
                f"  --model {model} \\",
                "  --harmful-n 30 \\",
                f"  --out runs/qwen36-ship-gate/aspa_n2_reg05_{label}_n30.json",
                "```",
                "",
                f"### `{label}` first-token KL",
                "",
                "```bash",
                f"{ENV_PREFIX} .venv/bin/python scripts/qwen36_kl_probe.py \\",
                f"  --label aspa_n2_reg05_{label}_source_kl \\",
                f"  --source {source} \\",
                f"  --candidate {model} \\",
                f"  --out runs/qwen36-kl/aspa_n2_reg05_{label}_source_kl.json \\",
                "  --device mps --dtype bfloat16",
                "```",
                "",
            ]
        )

    lines.extend(
        [
            "## Promotion Rule",
            "",
            "Promote a candidate to full n120 only if it beats the current leader on at least two of:",
            "",
            "- n30 ship score",
            "- capability score",
            "- mean KL",
            "- allowed/adversarial boundary score",
            "",
            "and loses none of the hard gates.",
            "",
            "After adding new JSON artifacts, refresh:",
            "",
            "```bash",
            ".venv/bin/python scripts/aspa_pareto_controller.py --out runs/qwen36-pareto/frontier.json",
            "```",
            "",
        ]
    )
    return "\n".join(lines)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--leader-metadata",
        default="outputs/qwen3.6-27b-aspa-n2-reg05-srcgamma090/source_interpolation_metadata.json",
    )
    ap.add_argument("--gamma", action="append", type=float, default=[0.875, 0.925])
    ap.add_argument("--out", default="runs/qwen36-pareto/next_experiments.md")
    args = ap.parse_args()

    meta = load_json(Path(args.leader_metadata))
    text = emit_plan(meta["source"], meta["candidate"], sorted(set(args.gamma)))
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(text)
    print(json.dumps({"event": "wrote_plan", "out": str(out), "gammas": sorted(set(args.gamma))}))


if __name__ == "__main__":
    main()