#!/usr/bin/env python
"""Hand-written mutation check for the expert-iteration driver and weight averaging.

Each mutant removes ONE rule in a private COPY of the repo subset and re-runs the targeted tests; a mutant
must make them fail ("KILLED"). The live tree is never touched. A SURVIVED or BAD-MUTANT line fails the run.

    CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/p2/mutation_check_ei.py
"""
import json, os, shutil, subprocess, sys
from pathlib import Path

ROOT = str(Path(__file__).resolve().parents[2])
MR = f"{ROOT}/.tmp/p2/mutroot_ei"
PY = sys.executable
EI = "scripts/p2/expert_iteration.py"
WA = "align/model/weight_average.py"
SD = "scripts/p2/build_sft_data.py"
TESTS = ["tests/test_expert_iteration.py", "tests/test_weight_average.py"]

M = [
    ("M01 band lower bound inclusive", EI, [('elif m["p"] <= lo:', 'elif m["p"] < lo:')]),
    ("M02 band upper bound exclusive", EI, [('elif m["p"] > hi:', 'elif m["p"] >= hi:')]),
    ("M03 unsolved problems always resampled", EI, [("    if not include_zero:\n        for pid in groups[\"unsolved\"]:", "    if False:\n        for pid in groups[\"unsolved\"]:")]),
    ("M04 shortest rank ignored by the SFT builder", SD, [('if getattr(a, "rank", "order") == "shortest":', 'if False:')]),
    ("M05 shortest picks the LONGEST", SD, [('first.sort(key=lambda k: (idx.meta[k]["n_completion"], k))', 'first.sort(key=lambda k: (-idx.meta[k]["n_completion"], k))')]),
    ("M06 stale marker accepted (config change not detected)", EI, [('    if mk.get("chain") != s.chain:\n        return "stale"', '    if False:\n        return "stale"')]),
    ("M07 non-zero exit still writes the marker", EI, [("        if rc != 0:\n            raise DriverError", "        if False:\n            raise DriverError")]),
    ("M08 partial outputs not cleaned before a re-run", EI, [("    for p in s.pre_clean:\n        if Path(p).exists():", "    for p in []:\n        if Path(p).exists():")]),
    ("M09 resume re-runs finished stages", EI, [('    if status == "done":\n        say(f"[round {s.r}] {s.name}: done (skipped)")\n        return "skipped"', '    if False:\n        return "skipped"')]),
    ("M10 export picks the OLDEST checkpoint", EI, [("    return cands[-1] if cands else None", "    return cands[0] if cands else None")]),
    ("M11 export merges onto the hard-coded snapshot", EI, [('{**env_py, "P2_EXPORT_BASE_SNAPSHOT": gen_model}', "dict(env_py)")]),
    ("M12 pass-rate table not carried across rounds", EI, [("    table = dict(prev)\n    for pid, m in new.items():", "    table = {}\n    for pid, m in new.items():")]),
    ("M13 regression boundary off by one", EI, [("last[\"solved\"] <= prior_best - max_regress", "last[\"solved\"] < prior_best - max_regress")]),
    ("M14 plateau needs only one flat round", EI, [("        if flat == patience:", "        if flat >= 1:")]),
    ("M15 stop rule compares across different n/cap", EI, [('comp = [h for h in history if (h["n_eval"], h["cap"]) == (last["n_eval"], last["cap"])]', "comp = list(history)")]),
    ("M16 average: storage-dtype arithmetic", WA, [("fa[lo:hi].to(torch.float64)", "fa[lo:hi]"), ("fb[lo:hi].to(torch.float64)", "fb[lo:hi]")]),
    ("M17 average: weights swapped", WA, [("((1.0 - w) * fa[lo:hi]", "((w) * fa[lo:hi]"), ("+ w * fb[lo:hi]", "+ (1.0 - w) * fb[lo:hi]")]),
    ("M18 average: last element of each chunk dropped", WA, [("hi = min(lo + chunk_elems, fa.numel())", "hi = min(lo + chunk_elems, fa.numel()) - (1 if fa.numel() > chunk_elems and lo + chunk_elems < fa.numel() else 0)")]),
    ("M19 average: shape mismatch tolerated", WA, [("        if ha[n][1] != hb[n][1] or ha[n][2] != hb[n][2]:", "        if False:")]),
    ("M20 average: partial output left behind on a crash", WA, [("        handles.clear()\n        shutil.rmtree(part, ignore_errors=True)\n        raise", "        handles.clear()\n        raise")]),
]


def setup():
    shutil.rmtree(MR, ignore_errors=True)
    os.makedirs(MR)
    shutil.copytree(f"{ROOT}/align", f"{MR}/align", ignore=shutil.ignore_patterns("__pycache__"))
    shutil.copytree(f"{ROOT}/scripts/p2", f"{MR}/scripts/p2", ignore=shutil.ignore_patterns("__pycache__"))
    shutil.copy(f"{ROOT}/scripts/qwen3_export_check.py", f"{MR}/scripts/qwen3_export_check.py")
    os.makedirs(f"{MR}/tests")
    for f in ("conftest.py", "test_p2_support.py", *[Path(t).name for t in TESTS]):
        shutil.copy(f"{ROOT}/tests/{f}", f"{MR}/tests/{f}")
    os.makedirs(f"{MR}/extras/vllm_rollout")
    shutil.copy(f"{ROOT}/extras/vllm_rollout/lean_prompts.py", f"{MR}/extras/vllm_rollout/lean_prompts.py")
    shutil.copy(f"{ROOT}/pyproject.toml", MR)
    os.makedirs(f"{MR}/docs")
    shutil.copy(f"{ROOT}/docs/P2.md", f"{MR}/docs/P2.md")


def run_tests():
    env = dict(os.environ, CUDA_VISIBLE_DEVICES="", PYTHONDONTWRITEBYTECODE="1")
    r = subprocess.run([PY, "-m", "pytest", *TESTS, "-q", "-x", "-p", "no:cacheprovider"], cwd=MR, env=env, capture_output=True, text=True)
    return r.returncode, r.stdout[-600:]


def main():
    setup()
    rc, out = run_tests()
    if rc != 0:
        print("BASELINE FAILS in the copy:\n", out)
        return 2
    results, bad = [], 0
    for name, rel, subs in M:
        path = f"{MR}/{rel}"
        orig = open(path).read()
        mut = orig
        ok = True
        for old, new in subs:
            if old not in mut:
                ok = False
                break
            mut = mut.replace(old, new, 1)
        if not ok:
            print(f"BAD-MUTANT {name}: pattern not found in {rel}")
            results.append({"mutant": name, "result": "BAD-MUTANT"})
            bad += 1
            continue
        open(path, "w").write(mut)
        try:
            rc, out = run_tests()
        finally:
            open(path, "w").write(orig)
        res = "KILLED" if rc != 0 else "SURVIVED"
        bad += res == "SURVIVED"
        failing = [l for l in out.splitlines() if l.startswith("FAILED")][:1]
        print(f"{res:9s} {name}  {failing[0][:110] if failing else ''}")
        results.append({"mutant": name, "result": res})
    os.makedirs(f"{ROOT}/.tmp/p2", exist_ok=True)
    json.dump(results, open(f"{ROOT}/.tmp/p2/mutants_result_ei.json", "w"), indent=1)
    print(f"{sum(r['result'] == 'KILLED' for r in results)}/{len(M)} killed")
    return 1 if bad else 0


if __name__ == "__main__":
    sys.exit(main())