#!/usr/bin/env python
"""Decontaminate ingested statement sets and the existing train pool against the evaluation sets
(miniF2F test+valid, MOBench, ProofNet# valid+test, optionally FormalMATH). Uses align.data.decontam.

    .venv/bin/python scripts/p2/supply/decontaminate.py --supply-dir .tmp/p2/supply

Writes, per input set `<name>`: `<name>.clean.jsonl` (kept rows), `<name>.flagged.jsonl` (id, reason, hits),
and `decontam_report.json` (hit counts per stage per source, independent and first-hit, plus each stage's
unique contribution). Nothing is deleted from the inputs.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import sys
from collections import Counter
from pathlib import Path

REPO = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO))

from align.data import hub  # noqa: E402
from align.data.decontam import STAGES, ContaminationLedger, Decontaminator, EvalItem  # noqa: E402

spec = importlib.util.spec_from_file_location("p2_common", REPO / "scripts" / "p2" / "p2_common.py")
C = sys.modules.get("p2_common") or importlib.util.module_from_spec(spec)
if "p2_common" not in sys.modules:
    sys.modules["p2_common"] = C
    spec.loader.exec_module(C)

import os  # noqa: E402

#: .tmp/ is gitignored and lives in the main checkout; worktrees share it (override with ALIGN_P2_TMP).
P2_TMP = Path(os.environ.get("ALIGN_P2_TMP", "/apps/development/rlhf/align/.tmp/p2"))
PROBLEMS = P2_TMP / "problems"
GOEDEL_MOBENCH = Path("/apps/development/rlhf/reference-repos/Goedel-Prover-V2/dataset/MOBench.jsonl")
PROOFNET = "PAug/ProofNetSharp"


def fetch_proofnet(out_dir: Path) -> list[EvalItem]:
    """ProofNet# (Lean 4 port, valid 185 + test 186) through the datasets-server JSON API; cached as JSONL."""
    items = []
    with urllib_open(f"https://huggingface.co/api/datasets/{PROOFNET}") as r:
        rev = json.loads(r.read())["sha"]
    for split in ("valid", "test"):
        cache = out_dir / f"proofnet_{split}.jsonl"
        if not cache.exists():
            rows = list(hub.iter_rows(PROOFNET, "default", split))
            with open(cache, "w", encoding="utf-8") as f:
                for i, row in enumerate(rows):
                    f.write(json.dumps({"row_index": i, "revision": rev, **row}, ensure_ascii=False) + "\n")
        for row in C.iter_jsonl(cache):
            stmt = (row.get("lean4_src_header") or "").rstrip() + "\n\n" + (row.get("lean4_formalization") or "")
            items.append(EvalItem(f"proofnet-{split}:{row.get('id')}", f"proofnet_{split}", stmt, row.get("nl_statement")))
    return items


def urllib_open(url):
    import urllib.request
    return urllib.request.urlopen(url, timeout=60)


def eval_items(supply_dir: Path, formalmath_as_eval: bool) -> tuple[list[EvalItem], dict]:
    items: list[EvalItem] = []
    for name, path in (("minif2f_test", PROBLEMS / "minif2f.jsonl"), ("minif2f_valid", PROBLEMS / "minif2f_valid.jsonl")):
        for r in C.iter_jsonl(path):
            items.append(EvalItem(f"{name}:{r['problem_id']}", name, r["formal_statement"], r.get("informal_prefix")))
    # all 360 MOBench problems (the verifier supports only 135 of them, contamination does not care)
    for r in C.iter_jsonl(GOEDEL_MOBENCH):
        items.append(EvalItem(f"mobench:{r['problem_id']}", "mobench", r.get("lean4_code") or r["formal_statement"], r.get("informal_prefix")))
    items += fetch_proofnet(supply_dir / "evalsets")
    if formalmath_as_eval:
        for r in C.iter_jsonl(supply_dir / "formalmath-all.jsonl"):
            items.append(EvalItem(f"formalmath:{r['problem_id']}", "formalmath", r["formal_statement"], r.get("informal_prefix")))
    return items, dict(Counter(i.eval_set for i in items))


def run_set(name: str, path: Path, dec: Decontaminator, out_dir: Path, elab: dict[str, str]) -> dict:
    ledger = ContaminationLedger()
    rows_meta: list[tuple[str, dict]] = []
    n = 0
    per_stage: Counter = Counter()
    first: Counter = Counter()
    by_eval: Counter = Counter()
    unique: Counter = Counter()
    flagged_f = open(out_dir / f"{name}.flagged.jsonl", "w", encoding="utf-8")
    clean_f = open(out_dir / f"{name}.clean.jsonl", "w", encoding="utf-8")
    inherited = 0
    # single pass: the ledger needs ancestry, and mutation children may precede their parent in a file, so
    # rows are held only as (id, line) -- the statement text is dropped once checked
    pending: list[tuple[str, str]] = []
    for r in C.iter_jsonl(path):
        n += 1
        hits = dec.check(r["formal_statement"], r.get("informal_prefix"), elab.get(r["problem_id"]))
        ledger.add(r["problem_id"], hits, r.get("ancestry"))
        pending.append((r["problem_id"], json.dumps(r, ensure_ascii=False)))
        if hits:
            stages = {h.stage for h in hits}
            for s in stages:
                per_stage[s] += 1
            first[min(stages, key=STAGES.index)] += 1
            if len(stages) == 1:
                unique[next(iter(stages))] += 1
            for es in {dec.items[h.eval_id].eval_set for h in hits}:
                by_eval[es] += 1
    flagged = 0
    for pid, line in pending:
        bad, why = ledger.resolve(pid)
        if bad:
            flagged += 1
            inherited += why.startswith("inherited")
            h = ledger.direct.get(pid, [])
            flagged_f.write(json.dumps({"problem_id": pid, "reason": why, "hits": [
                {"stage": x.stage, "eval_id": x.eval_id, "eval_set": dec.items[x.eval_id].eval_set, "score": round(x.score, 3)}
                for x in h[:8]]}, ensure_ascii=False) + "\n")
        else:
            clean_f.write(line + "\n")
    flagged_f.close(); clean_f.close()
    return {"rows": n, "flagged": flagged, "clean": n - flagged, "inherited": inherited,
            "flagged_by_stage_independent": {s: per_stage[s] for s in STAGES if per_stage[s]},
            "first_hit_stage": {s: first[s] for s in STAGES if first[s]},
            "only_this_stage_fired": {s: unique[s] for s in STAGES if unique[s]},
            "flagged_rows_by_eval_set": dict(by_eval)}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--supply-dir", default=str(P2_TMP / "supply"))
    ap.add_argument("--pool", default=str(PROBLEMS / "train_pool.jsonl"))
    ap.add_argument("--formalmath-as-eval", action="store_true")
    ap.add_argument("--elab-hashes", help="JSONL {problem_id, elab_sha} from elab_hash.py (eval + candidate sides)")
    a = ap.parse_args(argv)
    sd = Path(a.supply_dir)
    (sd / "evalsets").mkdir(parents=True, exist_ok=True)
    out = sd / "decontam"
    out.mkdir(exist_ok=True)
    items, eval_counts = eval_items(sd, a.formalmath_as_eval)
    elab: dict[str, str] = {}
    if a.elab_hashes:
        for r in C.iter_jsonl(a.elab_hashes):
            elab[r["problem_id"]] = r["elab_sha"]
        for it in items:
            if it.eval_id in elab:
                it.elab_sha = elab[it.eval_id]
    dec = Decontaminator(items)
    report = {"eval_sets": eval_counts, "eval_index_skipped": dict(dec.skipped), "sets": {}}
    inputs = {"train_pool": Path(a.pool)}
    for p in sorted(sd.glob("*.jsonl")):
        if p.name.endswith((".clean.jsonl", ".flagged.jsonl")):
            continue
        if a.formalmath_as_eval and p.name == "formalmath-all.jsonl":
            continue
        inputs[p.stem] = p
    for name, path in inputs.items():
        report["sets"][name] = run_set(name, path, dec, out, elab)
        print(name, json.dumps(report["sets"][name]), flush=True)
    (out / "decontam_report.json").write_text(json.dumps(report, indent=1))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())