#!/usr/bin/env python
"""Hand-written mutation check for the P2 modules: each mutant removes ONE fix (a guard, a mask
offset, a statistic's correction) 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.

    .venv/bin/python scripts/p2/mutation_check.py [--keep-going]

Writes `.tmp/p2/mutants_result.json`. A SURVIVED or BAD-MUTANT line fails the run: a surviving
mutant is a guard no test exercises, a bad mutant is a pattern that no longer matches the source.
"""
import json, os, shutil, subprocess, sys
from pathlib import Path
ROOT = str(Path(__file__).resolve().parents[2])
MR = f"{ROOT}/.tmp/p2/mutroot"
PY = sys.executable

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__"))
    os.makedirs(f"{MR}/tests")
    for f in os.listdir(f"{ROOT}/tests"):
        if f.startswith("test_p2_") or f == "conftest.py":
            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")  # p2_common loads it by path
    shutil.copy(f"{ROOT}/pyproject.toml", MR)
    os.makedirs(f"{MR}/docs")
    shutil.copy(f"{ROOT}/docs/P2.md", f"{MR}/docs/P2.md")  # the pre-registration test reads it

T = {
 "stats": ["tests/test_p2_proof_stats.py"],
 "prob": ["tests/test_p2_problems.py"],
 "ver": ["tests/test_p2_verify.py"],
 "sft": ["tests/test_p2_sft_data.py"],
 "rep": ["tests/test_p2_eval_report.py"],
}
PS = "align/eval/proof_stats.py"; CM = "scripts/p2/p2_common.py"; BP = "scripts/p2/build_problems.py"
VG = "scripts/p2/verify_generations.py"; SD = "scripts/p2/build_sft_data.py"; ER = "scripts/p2/eval_report.py"
M = [
 ("M01 plug-in pass@k estimator", PS, [("return float(1 - Fraction(math.comb(n - c, k), math.comb(n, k)))", "return 1 - (1 - c / n) ** k")], "stats"),
 ("M02 allow k>n", PS, [("    if k > n:\n        raise ValueError(\n            f\"pass@{k} needs", "    if False:\n        raise ValueError(\n            f\"pass@{k} needs")], "stats"),
 ("M03 mcnemar tail off by one", PS, [("tail = sum(math.comb(n, i) for i in range(m + 1))", "tail = sum(math.comb(n, i) for i in range(m))")], "stats"),
 ("M04 mcnemar one-sided", PS, [("Fraction(2 * tail, 2**n)", "Fraction(tail, 2**n)")], "stats"),
 ("M05 mcnemar p not capped at 1", PS, [("float(min(Fraction(1), Fraction(2 * tail, 2**n)))", "float(Fraction(2 * tail, 2**n))")], "stats"),
 ("M06 solved@k over all samples not the prefix", PS, [("return any(accepted_by_sample_idx[:k])", "return any(accepted_by_sample_idx)")], "stats"),
 ("M07 mcnemar pairs unchecked", PS, [("    if set(solved_a) != set(solved_b):", "    if False:")], "stats"),
 ("M08 wrong z", PS, [("Z95 = 1.959963984540054", "Z95 = 1.5")], "stats"),
 ("M09 unpaired bootstrap (independent resamples)", PS, [("means = _boot_means(vb - va, seed, n_boot)", "means = _boot_means(vb, seed, n_boot) - _boot_means(va, seed + 1, n_boot)")], "stats"),
 ("M10 statement_sha includes the theorem name", CM, [("    _, _, sig = statement_key(statement)\n    return sha256_text(\" \".join(sig))", "    nm, _, sig = statement_key(statement)\n    return sha256_text(nm + \" \".join(sig))")], "prob"),
 ("M11 alpha_sha does not rename", CM, [("out.append(mp.setdefault(t, f\"v{len(mp)}\"))", "out.append(t)")], "prob"),
 ("M12 name layer off", BP, [("if n in self.names or (self._name_re is not None and self._name_re.search(n)):", "if False:")], "prob"),
 ("M13 name substring layer off", BP, [("if n in self.names or (self._name_re is not None and self._name_re.search(n)):", "if n in self.names:")], "prob"),
 ("M14 informal 13-gram layer off", BP, [("        if not self.grams.isdisjoint(C.ngrams(toks, self.ngram)):\n            return \"informal_13gram\"", "        if False:\n            return \"informal_13gram\"")], "prob"),
 ("M14b informal 8-gram layer off", BP, [("if len(self.short_grams & C.ngrams(toks, self.short_ngram)) >= self.short_hits:", "if False:")], "prob"),
 ("M14c a single 8-gram triggers", BP, [("SHORT_HITS = 3", "SHORT_HITS = 1")], "prob"),
 ("M15 sampler keeps largest hashes", BP, [("elif item > heap[0]:", "elif item < heap[0]:")], "prob"),
 ("M16 sampler skips decontamination", BP, [("        if why is not None:\n            counts[f\"contaminated_{why}\"] += 1\n            continue", "        if False:\n            continue")], "prob"),
 ("M17 failures counted beyond the cutoff", BP, [("cutoff = candidates.index(selected[-1]) + 1", "cutoff = len(log_rows)")], "prob"),
 ("M18 elaboration instrument not ceiling-tested", BP, [("        if res[0] is not want:", "        if False:")], "prob"),
 ("M19 normalize accepts several ':= by'", CM, [("if len(hits) != 1:", "if len(hits) < 1:")], "prob"),
 ("M20 length-capped samples sent to Lean", VG, [("    if gen[\"finish_reason\"] == \"length\":\n        return \"finish_reason=length\"", "    if False:\n        return \"finish_reason=length\"")], "ver"),
 ("M21 stale verdict accepted", VG, [("if existing[key][\"code_sha256\"] != C.code_sha(gen[\"lean_code\"]):", "if False:")], "ver"),
 ("M22 resume ignores existing rows", VG, [("            if key in existing:\n                if existing", "            if False:\n                if existing")], "ver"),
 ("M23 dead-verifier chunk written", VG, [("if len(part) >= 8 and len(held) / len(part) > 0.5:", "if False:")], "ver"),
 ("M24 feedback computed on the full file instead of compile_view", VG, [("    view = C.compile_view(code)", "    view = code")], "ver"),
 ("M24b feedback drops endPos None handling", VG, [("return None if not p else {\"line\": int(p[\"line\"]), \"column\": int(p[\"column\"])}", "return {\"line\": int((p or {\"line\": 0})[\"line\"]), \"column\": 0}")], "ver"),
 ("M25b torn single line kept", CM, [("    cut = data.rfind(b\"\\n\") + 1  # 0 when the file has a single unterminated line", "    cut = max(data.rfind(b\"\\n\"), 0)")], "ver"),
 ("M25c atomic write leaves tmp", CM, [("        if tmp.exists():\n            tmp.unlink()\n    return n", "        pass\n    return n")], "ver"),
 ("M25d round>0 parent unchecked", CM, [("        if parent[\"problem_id\"] != pid or parent[\"round\"] != rec[\"round\"] - 1:", "        if False:")], "ver"),
 ("M25 torn tail kept", CM, [("        with open(p, \"r+b\") as f:\n            f.truncate(cut)", "        with open(p, \"r+b\") as f:\n            f.truncate(len(data))")], "ver"),
 ("M26 prompt_len off by one (+ readback disabled)", SD, [("prompt_len=len(rec[\"prompt_ids\"]),", "prompt_len=len(rec[\"prompt_ids\"]) + 1,"), ("n_checked = readback_check(out, idx, selected, held)", "n_checked = len(selected)")], "sft"),
 ("M27 completion truncated by one token (+ readback disabled)", SD, [("ids = rec[\"prompt_ids\"] + rec[\"completion_ids\"]", "ids = rec[\"prompt_ids\"] + rec[\"completion_ids\"][:-1]"), ("n_checked = readback_check(out, idx, selected, held)", "n_checked = len(selected)")], "sft"),
 ("M28 prompt_len off by one (readback ON)", SD, [("prompt_len=len(rec[\"prompt_ids\"]),", "prompt_len=len(rec[\"prompt_ids\"]) - 1,")], "sft"),
 ("M29 easy problems kept", SD, [("if n0 >= EASY_MIN_SAMPLES and c0 == n0 and not a.keep_easy:", "if False:")], "sft"),
 ("M30 n<8 counted as easy", SD, [("if n0 >= EASY_MIN_SAMPLES and c0 == n0", "if c0 == n0")], "sft"),
 ("M31 other-cap samples mixed in", SD, [("                if rec[\"cap_tokens\"] != cap_tokens:", "                if False:")], "sft"),
 ("M32 correction of an accepted parent kept", SD, [("    if ver[pkey][\"accepted\"]:", "    if False:")], "sft"),
 ("M33 correction prompt not checked for the failed attempt", SD, [("if not tail or tail not in prompt or \"rror\" not in prompt:", "if False:")], "sft"),
 ("M34 duplicate proofs not de-duplicated", SD, [("if room(kind, kind_limit) and h not in seen_hash:", "if room(kind, kind_limit):")], "sft"),
 ("M35 EOS requirement removed", SD, [("if not a.allow_no_eos and m[\"last_token\"] not in eos:", "if False:")], "sft"),
 ("M36 eval-statement leak check removed", SD, [("    if leaked:", "    if False:")], "sft"),
 ("M37 foreign problems allowed", SD, [("    if foreign:", "    if False:")], "sft"),
 ("M38 stale verification accepted by the SFT builder", SD, [("if rec[\"code_sha256\"] != m[\"code_sha\"]:", "if False:")], "sft"),
 ("M39 per-problem cap ignored", SD, [("return len(chosen) < a.max_per_problem and", "return True and")], "sft"),
 ("M48 statement-fidelity filter removed from the SFT selection", SD, [("if why is None and not C.displayed_statement_matches(", "if False and not C.displayed_statement_matches(")], "sft"),
 ("M49 displayed theorem compared by presence only (any statement passes)", CM, [("return shown is not None and shown == statement_key(formal_statement)[2]", "return shown is not None")], "sft"),
 ("M51 PRIMARY label without the registered problem count", ER, [("    if base_metrics[\"n_problems\"] != PRIMARY[\"n_problems\"]:", "    if False:")], "rep"),
 ("M52 PRIMARY label without the registered baseline revision", ER, [("    if revs != [PRIMARY[\"baseline_model_rev\"]]:", "    if False:")], "rep"),
 ("M53 PRIMARY label without the registered sampling settings", ER, [("        if (t, p) != (want[\"temperature\"], want[\"top_p\"]) or kk not in want[\"top_k\"]:", "        if False:")], "rep"),
 ("M54 runs with different sampling settings compared", ER, [("    if len(sb) != 1 or len(sc) != 1 or sb != sc:", "    if False:")], "rep"),
 ("M54b a run with mixed sampling settings compared", ER, [("    if len(sb) != 1 or len(sc) != 1 or sb != sc:", "    if sb != sc:")], "rep"),
 ("M55 differently prompted runs compared", ER, [("    bad = sorted({k[0] for k in common if base[k].prompt_sha != cand[k].prompt_sha})", "    bad = []")], "rep"),
 ("M56 design mismatch accepted without the override flag", ER, [("            if not allow_design_mismatch:", "            if False:")], "rep"),
 ("M57 decision rule ignores the direction (any significant difference is a success)", ER, [("    if sig and m[\"solved_b\"] > m[\"solved_a\"]:", "    if sig:")], "rep"),
 ("M58 decision rule: alpha 0.5", ER, [("PRIMARY_ALPHA = 0.05", "PRIMARY_ALPHA = 0.5")], "rep"),
 ("M40 caps pooled in the report", ER, [("if rd == 0 and c == cap and pid in problem_ids:", "if rd == 0 and pid in problem_ids:")], "rep"),
 ("M41 missing problems scored as failures", ER, [("    if missing:\n        raise C.ContractError(f\"incomplete run", "    if False:\n        raise C.ContractError(f\"incomplete run")], "rep"),
 ("M42 verifier mismatch allowed silently", ER, [("            if not allow_verifier_mismatch:", "            if False:")], "rep"),
 ("M43 primary endpoint not flagged", ER, [("\"endpoint\": \"PRIMARY\" if not blockers else \"secondary\"", "\"endpoint\": \"secondary\"")], "rep"),
 ("M44 correction of accepted parent allowed", ER, [("            if samples[par].accepted:", "            if False:")], "rep"),
 ("M45 stale verification accepted by the report", ER, [("if v[\"code_sha256\"] != C.code_sha(rec[\"lean_code\"]):", "if False:")], "rep"),
 ("M47 pre-registration hash never mismatches", ER, [("\"verified\" if rec.group(1) == h else \"MODIFIED_AFTER_REGISTRATION\"", "\"verified\"")], "rep"),
 ("M46 McNemar on pass@k means not solved@k (uses any of all samples)", ER, [("sa = {p: S.solved_at_k(fb[p], k) for p in fb}", "sa = {p: any(fb[p]) for p in fb}")], "rep"),
]
def run(tests):
    env = dict(os.environ, PYTHONPATH=MR)
    r = subprocess.run([PY, "-m", "pytest", *tests, "-q", "-x", "-m", "not slow", "-p", "no:cacheprovider"], cwd=MR, env=env, capture_output=True, text=True, timeout=900)
    return r.returncode, r.stdout.strip().splitlines()[-1] if r.stdout.strip() else r.stderr[-200:]
if __name__ == "__main__":
    setup()
    # baseline must be green
    allt = sum(T.values(), [])
    rc, last = run(allt)
    print("BASELINE", rc, last, flush=True)
    assert rc == 0
    out = []
    for name, path, reps, grp in M:
        full = f"{MR}/{path}"
        orig = open(full).read()
        new = orig
        for o, n in reps:
            if o not in new:
                print("[BAD-MUTANT]", name); new = None; break
            new = new.replace(o, n, 1)
        if new is None:
            out.append((name, "bad", "")); continue
        open(full, "w").write(new)
        try:
            rc, last = run(T[grp])
        finally:
            open(full, "w").write(orig)
        st = "KILLED" if rc != 0 else "SURVIVED"
        print(f"[{st}] {name}  :: {last}", flush=True)
        out.append((name, st, last))
    json.dump(out, open(f"{ROOT}/.tmp/p2/mutants_result.json", "w"), indent=1)
    bad = [n for n, s, _ in out if s != "KILLED"]
    print(f"{len(out) - len(bad)}/{len(out)} killed; survivors/bad:", bad)
    sys.exit(1 if bad else 0)