#!/usr/bin/env python
"""Agreement + throughput: our strict verifier vs the Kimina Lean server pre-filter.

    hostguard.py run --cores 4 --mem-gb 24 --label agree --wait -- \\
        .venv/bin/python scripts/p2/backend_agreement.py --out .tmp/kimina_agreement/run1

Item sets (all verified by BOTH backends on the same (formal statement, extracted Lean code)):
  kg      known-good proofs of the ceiling test (150 Lean-Workbook-proofs rows saved by the registered
          ceiling run + 30 restated), statements derived exactly as `lean_ceiling_test.py` does
  adv     the 61 frozen adversarial cases of `lean_ceiling_cases.py` (3 bombs run on a 1 GiB / 30 s
          pool for ours and on a second, small Kimina server for the candidate)
  sample  N random code-bearing records (finish_reason != length) of a P2 generation file

Phases (resumable, JSONL per phase): `ours` (StrictVerifier, the authority), `control` (OUR REPL pool,
compile-only: separates "Kimina server overhead" from "strict-tier overhead"), `kimina` (compile-only
through the server). They run strictly one after another: never two heavy jobs at once. The server(s)
must already be running (toolchains/work/kimina-run/start.sh); this script never starts or stops them.

Semantics of the comparison. Kimina is a PRE-FILTER: its verdict is `reject` / `maybe` / `error`, never
an accept. The "gated" outcome (accept iff strict accepts AND the pre-filter did not reject) is computed
from the two per-item records; the only accept that can exist is the strict verifier's.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import random
import statistics
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

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

MAIN_REPO = Path("/apps/development/rlhf/align")  # data lives in the main checkout's .tmp, not in worktrees
DEFAULT_GENS = MAIN_REPO / ".tmp/p2/r1/baseline/gen_minif2f_c40960.jsonl"
DEFAULT_PROBLEMS = MAIN_REPO / ".tmp/p2/problems/minif2f.jsonl"
DEFAULT_KG_RAW = MAIN_REPO / ".tmp/lean_ceiling/registered/known_good_raw.jsonl"
CORES = 4


def _load(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod


LC = _load("lean_ceiling_test", REPO / "scripts" / "lean_ceiling_test.py")
CASES = LC.C


def pctl(xs, q):
    s = sorted(xs)
    return s[min(len(s) - 1, int(q * len(s)))] if s else float("nan")


# --------------------------------------------------------------------------- items


def build_items(args) -> list[dict]:
    from align.verify.lean.reward import extract_lean_code

    items: list[dict] = []
    # known-good
    cands = []
    for line in open(args.kg_raw):
        r = json.loads(line)
        cands.append((r["problem_id"], "lean-workbook-proofs", r["full_proof"]))
    for rs in CASES.RESTATED:
        cands.append((rs["id"], "restated", CASES.HEADER + rs["code"]))
    for pid, src, code in cands:
        ds = LC.derive_statement(code)
        if ds is None:
            continue
        items.append(dict(id=f"kg:{pid}", set="kg", sub=src, statement=ds[0], code=code, pool="main"))
    # adversarial
    for c in CASES.ADVERSARIAL:
        stmt = CASES.HEADER + CASES.STATEMENTS[c["statement"]]
        comp = c["completion"] if c["completion"] is not None else LC.as_completion(
            (CASES.HEADER if c["raw_header"] is None else c["raw_header"]) + c["code"])
        items.append(dict(id=f"adv:{c['id']}", set="adv", sub=c["category"], statement=stmt,
                          code=extract_lean_code(comp), pool=c["pool"]))
    # sampled generations
    probs = {}
    for line in open(args.problems):
        r = json.loads(line)
        probs[r["problem_id"]] = r["formal_statement"]
    pop = []
    for line in open(args.generations):
        r = json.loads(line)
        if r["finish_reason"] == "length" or not (r["lean_code"] or "").strip() or r["problem_id"] not in probs:
            continue
        pop.append(r)
    rng = random.Random(args.seed)
    for r in rng.sample(pop, min(args.sample, len(pop))):
        items.append(dict(id=f"gen:{r['problem_id']}#{r['sample_idx']}", set="sample", sub="minif2f",
                          statement=probs[r["problem_id"]], code=r["lean_code"], pool="main"))
    items.append(dict(id="__meta__", set="meta", population=len(pop)))
    return items


def done_ids(path: Path) -> dict:
    out = {}
    if path.exists():
        for line in open(path):
            try:
                r = json.loads(line)
                out[r["id"]] = r
            except json.JSONDecodeError:
                pass
    return out


class Sink:
    def __init__(self, path: Path):
        self.f = open(path, "a")
        self.lock = threading.Lock()

    def write(self, rec: dict):
        with self.lock:
            self.f.write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
            self.f.flush()


# --------------------------------------------------------------------------- phases


def phase_ours(items, out: Path, args) -> dict:
    from align.verify.lean import ReplLimits, ReplPool, StrictConfig, StrictVerifier, find_toolchain

    tc = find_toolchain()
    sink = Sink(out / "ours.jsonl")
    have = done_ids(out / "ours.jsonl")
    walls = {}
    for pool_name, lim, cfg in (("main", ReplLimits(), StrictConfig(proof_timeout_s=args.timeout)),
                                ("bomb", ReplLimits(anon_cap_gb=1.0), StrictConfig(proof_timeout_s=30))):
        todo = [it for it in items if it["set"] != "meta" and it["pool"] == pool_name and it["id"] not in have]
        if not todo:
            continue
        n = args.workers if pool_name == "main" else 1
        pool = ReplPool.from_toolchain(tc, n_workers=n, limits=lim, preflight=True,
                                       scratch_dir=str(tc.root / "tmp" / f"agree_{pool_name}"))
        try:
            ver = StrictVerifier(pool, cfg)
            t0 = time.monotonic()

            def one(w, it):
                r = ver._verify_on(w, (it["statement"], it["code"]))
                sink.write(dict(id=it["id"], set=it["set"], accepted=bool(r.accepted), reasons=list(r.reasons),
                                stage=r.stage, elapsed_s=r.elapsed_s, undecided=bool(r.undecided)))

            pool.map(one, todo)
            walls[pool_name] = time.monotonic() - t0
        finally:
            pool.close()
    return {"walls": walls, "workers": args.workers}


def phase_control(items, out: Path, args) -> dict:
    """OUR REPL pool, compile-only (no strictness): the engine-level control for the Kimina server."""
    from align.verify.lean import ReplLimits, ReplPool, find_toolchain

    tc = find_toolchain()
    sink = Sink(out / "control.jsonl")
    have = done_ids(out / "control.jsonl")
    todo = [it for it in items if it["set"] in ("kg", "sample") and it["id"] not in have]
    if not todo:
        return {}
    pool = ReplPool.from_toolchain(tc, n_workers=args.workers, limits=ReplLimits(), preflight=True,
                                   scratch_dir=str(tc.root / "tmp" / "agree_control"))
    try:
        t0 = time.monotonic()

        def one(w, it):
            t = time.monotonic()
            r = LC.compile_only(w, it["statement"], it["code"], args.timeout) if it["code"] else {"ok": False, "why": "no_code"}
            sink.write(dict(id=it["id"], set=it["set"], ok=bool(r.get("ok")), clean=bool(r.get("clean")), why=r.get("why", ""),
                            elapsed_s=time.monotonic() - t))

        pool.map(one, todo)
        return {"wall": time.monotonic() - t0, "workers": args.workers}
    finally:
        pool.close()


def phase_kimina(items, out: Path, args) -> dict:
    from align.verify_backends.kimina_server import KiminaClient, KiminaPrefilter

    sink = Sink(out / "kimina.jsonl")
    have = done_ids(out / "kimina.jsonl")
    walls = {}
    for pool_name, url, tmo, workers in (("main", args.kimina_url, args.timeout, args.workers),
                                         ("bomb", args.kimina_bomb_url, 30, 1)):
        todo = [it for it in items if it["set"] != "meta" and it["pool"] == pool_name and it["id"] not in have]
        if not todo:
            continue
        cl = KiminaClient(url)
        if not cl.healthy():
            raise SystemExit(f"Kimina server at {url} is not healthy; start it first")
        pf = KiminaPrefilter(cl, tmo)

        def one(it):
            t = time.monotonic()
            v = pf.check_code(it["code"])
            sink.write(dict(id=it["id"], set=it["set"], verdict=v.verdict, reasons=list(v.reasons), detail=v.detail,
                            elapsed_s=time.monotonic() - t, server_s=v.elapsed_s))

        t0 = time.monotonic()
        with ThreadPoolExecutor(workers) as ex:
            list(ex.map(one, todo))
        walls[pool_name] = time.monotonic() - t0
    return {"walls": walls, "workers": args.workers}


# --------------------------------------------------------------------------- report


def stats(xs):
    return {"n": len(xs), "sum_s": round(sum(xs), 1), "mean": round(statistics.mean(xs), 3) if xs else None,
            "p50": round(pctl(xs, .5), 3), "p90": round(pctl(xs, .9), 3), "p99": round(pctl(xs, .99), 3),
            "max": round(max(xs), 3) if xs else None}


def report(items, out: Path, phase_meta: dict) -> dict:
    ours, ctl, kim = done_ids(out / "ours.jsonl"), done_ids(out / "control.jsonl"), done_ids(out / "kimina.jsonl")
    by_id = {it["id"]: it for it in items if it["set"] != "meta"}
    rep: dict = {"phase_meta": phase_meta, "matrix": {}, "disagreements": [], "n_items": {}}

    def okind(r):  # ours: accept / reject / undecided
        return "accept" if r["accepted"] else ("undecided" if r["undecided"] else "reject")

    for s in ("kg", "adv", "sample", "ALL"):
        ids = [i for i in by_id if (s == "ALL" or by_id[i]["set"] == s) and i in ours and i in kim]
        m: dict = {}
        for i in ids:
            m.setdefault(okind(ours[i]), {}).setdefault(kim[i]["verdict"], 0)
            m[okind(ours[i])][kim[i]["verdict"]] += 1
        rep["matrix"][s] = {"n": len(ids), "rows(ours)xcols(kimina)": m}
    # ceiling truth
    kg_ids = [i for i in by_id if by_id[i]["set"] == "kg" and i in ours and i in kim]
    adv_ids = [i for i in by_id if by_id[i]["set"] == "adv" and i in ours and i in kim]
    rep["ceiling"] = {
        "kg_ours_accept": sum(ours[i]["accepted"] for i in kg_ids), "kg_n": len(kg_ids),
        "kg_kimina_maybe": sum(kim[i]["verdict"] == "maybe" for i in kg_ids),
        "kg_kimina_reject": sum(kim[i]["verdict"] == "reject" for i in kg_ids),
        "adv_ours_accept": sum(ours[i]["accepted"] for i in adv_ids), "adv_n": len(adv_ids),
        "adv_kimina_maybe(compile-only would pass these on)": sum(kim[i]["verdict"] == "maybe" for i in adv_ids),
        "adv_kimina_reject": sum(kim[i]["verdict"] == "reject" for i in adv_ids),
        "gated_adv_accepts": sum(ours[i]["accepted"] and kim[i]["verdict"] != "reject" for i in adv_ids),
    }
    # disagreements: everything that is not (ours accept & kimina maybe) or (ours reject & kimina reject)
    for i in by_id:
        if i not in ours or i not in kim:
            continue
        o, k = ours[i], kim[i]
        ok = okind(o)
        kv = k["verdict"]
        if (ok == "accept" and kv == "maybe") or (ok == "reject" and kv == "reject"):
            continue
        if ok == "accept" and kv == "reject":
            klass = "PREFILTER_FALSE_REJECT (recall loss)"
        elif ok == "reject" and kv == "maybe":
            klass = "STRICT_ONLY_REJECT (expected: policy/audit layer)"
        elif ok == "accept" and kv == "error":
            klass = "PREFILTER_UNDECIDED_STRICT_ACCEPT"
        elif ok == "undecided" or kv == "error":
            klass = "UNDECIDED_ON_ONE_SIDE"
        else:
            klass = "OTHER"
        rep["disagreements"].append(dict(id=i, set=by_id[i]["set"], klass=klass, ours=ok, ours_reasons=o["reasons"],
                                         ours_s=round(o["elapsed_s"], 2), kimina=kv, kimina_reasons=k["reasons"],
                                         kimina_detail=k.get("detail"), kimina_s=round(k["elapsed_s"], 2),
                                         control=({"ok": ctl[i]["ok"], "clean": ctl[i]["clean"], "why": ctl[i]["why"]} if i in ctl else None)))
    from collections import Counter

    rep["disagreement_classes"] = dict(Counter(d["klass"] for d in rep["disagreements"]))
    # throughput over the sample set (+ kg) : same items, three engines
    for name, rows, key in (("ours_strict", ours, "elapsed_s"), ("ours_compile_only", ctl, "elapsed_s"), ("kimina_compile_only", kim, "elapsed_s")):
        for s in ("sample", "kg"):
            xs = [rows[i][key] for i in by_id if by_id[i]["set"] == s and i in rows]
            rep.setdefault("latency_s", {}).setdefault(name, {})[s] = stats(xs)
    # busy-time throughput on the SAME items (kg + sample): n / sum(per-item latency) = proofs per second per
    # worker-core, independent of the heavy tail's wall-clock packing and of which extra items a phase ran
    same = [i for i in by_id if by_id[i]["set"] in ("sample", "kg") and i in ours and i in kim and i in ctl]
    rep["busy_throughput_proofs_per_s_per_core"] = {
        "n": len(same),
        **{nm: round(len(same) / sum(rows[i]["elapsed_s"] for i in same), 4)
           for nm, rows in (("ours_strict", ours), ("ours_compile_only", ctl), ("kimina_compile_only", kim))}}
    # gated pipeline, from the per-item records (never re-run)
    samp = [i for i in by_id if by_id[i]["set"] in ("sample", "kg") and i in ours and i in kim]
    if samp:
        rep["gated"] = {
            "n": len(samp),
            "strict_accepts": sum(ours[i]["accepted"] for i in samp),
            "gated_accepts": sum(ours[i]["accepted"] and kim[i]["verdict"] != "reject" for i in samp),
            "prefilter_rejects": sum(kim[i]["verdict"] == "reject" for i in samp),
            "strict_calls_saved": sum(kim[i]["verdict"] == "reject" for i in samp),
            "sum_strict_s_saved": round(sum(ours[i]["elapsed_s"] for i in samp if kim[i]["verdict"] == "reject"), 1),
            "sum_kimina_s_spent": round(sum(kim[i]["elapsed_s"] for i in samp), 1),
            "sum_strict_s_total": round(sum(ours[i]["elapsed_s"] for i in samp), 1),
            "sum_strict_s_on_survivors": round(sum(ours[i]["elapsed_s"] for i in samp if kim[i]["verdict"] != "reject"), 1),
        }
    # agreement on the sample set between the two COMPILE-ONLY engines (our REPL vs Kimina)
    both = [i for i in by_id if i in ctl and i in kim and by_id[i]["set"] in ("sample", "kg")]
    if both:
        agree = sum((ctl[i]["clean"]) == (kim[i]["verdict"] == "maybe") for i in both if kim[i]["verdict"] != "error" and ctl[i]["why"] != "timeout")
        n2 = sum(1 for i in both if kim[i]["verdict"] != "error" and ctl[i]["why"] != "timeout")
        rep["engine_control"] = {"n_decided_by_both": n2, "clean_agree": agree,
                                 "diffs": [dict(id=i, ours_compile_clean=ctl[i]["clean"], ours_why=ctl[i]["why"], kimina=kim[i]["verdict"], kimina_reasons=kim[i]["reasons"])
                                           for i in both if kim[i]["verdict"] != "error" and ctl[i]["why"] != "timeout" and ctl[i]["clean"] != (kim[i]["verdict"] == "maybe")]}
    return rep


def throughput(rep, items, out: Path, phase_meta: dict) -> None:
    by_set = {}
    for it in items:
        if it["set"] != "meta":
            by_set.setdefault(it["pool"], []).append(it["id"])
    tp = {}
    for name, f, wkey in (("ours_strict", "ours.jsonl", "ours"), ("ours_compile_only", "control.jsonl", "control"), ("kimina_compile_only", "kimina.jsonl", "kimina")):
        m = phase_meta.get(wkey) or {}
        wall = m.get("wall") if "wall" in m else (m.get("walls") or {}).get("main")
        n_main = sum(1 for r in done_ids(out / f).values() if by_set.get("main") and r["id"] in set(by_set["main"]))
        if wall:
            tp[name] = {"n": n_main, "wall_s": round(wall, 1), "proofs_per_s": round(n_main / wall, 4),
                        "proofs_per_s_per_core": round(n_main / wall / CORES, 4), "cores": CORES}
    rep["throughput"] = tp


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--out", required=True)
    ap.add_argument("--phases", default="ours,control,kimina")
    ap.add_argument("--workers", type=int, default=CORES)
    ap.add_argument("--timeout", type=float, default=180.0)
    ap.add_argument("--sample", type=int, default=500)
    ap.add_argument("--seed", type=int, default=20261019)
    ap.add_argument("--generations", default=str(DEFAULT_GENS))
    ap.add_argument("--problems", default=str(DEFAULT_PROBLEMS))
    ap.add_argument("--kg-raw", default=str(DEFAULT_KG_RAW))
    ap.add_argument("--kimina-url", default="http://127.0.0.1:18000")
    ap.add_argument("--kimina-bomb-url", default="http://127.0.0.1:18001")
    ap.add_argument("--sets", default="kg,adv,sample", help="comma list; smoke runs use e.g. --sets sample --sample 6")
    ap.add_argument("--report-only", action="store_true")
    args = ap.parse_args()
    if args.workers > CORES:
        print(f"host rule: at most {CORES} cores/workers", file=sys.stderr)
        return 2
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    items = [i for i in build_items(args) if i["set"] == "meta" or i["set"] in args.sets.split(",")]
    meta_it = [i for i in items if i["set"] == "meta"][0]
    (out / "items_manifest.json").write_text(json.dumps(
        {"n": len(items) - 1, "population_code_bearing": meta_it["population"], "seed": args.seed,
         "sets": {s: sum(1 for i in items if i["set"] == s) for s in ("kg", "adv", "sample")},
         "ids": [i["id"] for i in items if i["set"] != "meta"]}, indent=1))
    meta_path = out / "phase_meta.json"
    phase_meta = json.loads(meta_path.read_text()) if meta_path.exists() else {}
    if not args.report_only:
        for ph, fn in (("ours", phase_ours), ("control", phase_control), ("kimina", phase_kimina)):
            if ph in args.phases.split(","):
                print(f"== phase {ph} ==", flush=True)
                t = time.monotonic()
                m = fn(items, out, args)
                if m:
                    phase_meta[ph] = m
                    meta_path.write_text(json.dumps(phase_meta, indent=1))
                print(f"   done in {time.monotonic() - t:.0f}s", flush=True)
    rep = report(items, out, phase_meta)
    throughput(rep, items, out, phase_meta)
    (out / "report.json").write_text(json.dumps(rep, indent=1, default=str))
    print(json.dumps({k: rep[k] for k in ("matrix", "ceiling", "disagreement_classes", "throughput", "gated") if k in rep}, indent=1))
    return 0


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