#!/usr/bin/env python3
"""Verified best-of-N, graded from the cached samples. No model, no GPU, no network.

certified@N = the greedy (temperature 0) answer verifies, OR any of the first N-1 sampled
candidates verifies -- all under the strict oracle (byte-identical code object, docstrings and
exception tables included).

This is the number that matters for a decompiler product rather than for a leaderboard: because
the oracle is SOUND, a certified answer is proven correct, so best-of-N buys real accuracy rather
than a better guess. The uncertified remainder is reported as unknown, never as wrong.

    ./bestofn_grade.py --bench <bench.jsonl> --greedy <gen.jsonl> --samples <boN.jsonl> \\
                       --name CSN-3.12 --out ../results/bestofn.json
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import load_bench, load_jsonl, ours_ok, self_test, require_sound, strip_fences  # noqa: E402

DEFAULT_NS = (1, 4, 8, 16, 32)


def grade(name: str, bench_p: str, greedy_p: str, samples_p: str, Ns=DEFAULT_NS) -> dict:
    bench, _ = load_bench(bench_p)
    greedy = {r["i"]: r for r in load_jsonl(greedy_p)}
    n = len(bench)

    greedy_pass = {
        i for i in bench
        if i in greedy and ours_ok(strip_fences(greedy[i]["got"]), bench[i]["expected"])
    }

    samples: dict[int, dict[int, str]] = {}
    for r in load_jsonl(samples_p):
        samples.setdefault(r["i"], {})[r["s"]] = r["got"]

    first_pass: dict[int, int | None] = {}
    diversity: dict[int, tuple[int, int]] = {}
    for i in bench:
        if i in greedy_pass:
            continue
        exp = bench[i]["expected"]
        fp, seen = None, set()
        smap = samples.get(i, {})
        for s in sorted(smap):
            src = strip_fences(smap[s])
            seen.add(src)
            if fp is None and ours_ok(src, exp):
                fp = s
        first_pass[i] = fp
        diversity[i] = (len(seen), len(smap))

    curve = {}
    for N in Ns:
        cert = len(greedy_pass) + sum(1 for fp in first_pass.values() if fp is not None and fp <= N - 1)
        curve[str(N)] = {
            "certified": cert,
            "pct": round(100 * cert / n, 2),
            "recovered_from_sampling": sum(
                1 for fp in first_pass.values() if fp is not None and fp <= N - 1),
        }

    return {
        "benchmark": name,
        "n": n,
        "oracle": "strict: byte-identical code object incl. docstrings and co_exceptiontable",
        "greedy_pass": len(greedy_pass),
        "greedy_pass_pct": round(100 * len(greedy_pass) / n, 2),
        "greedy_fail": len(first_pass),
        "ever_recovered_by_sampling": sum(1 for fp in first_pass.values() if fp is not None),
        "curve": curve,
        "avg_unique_candidates_per_failed_input":
            round(sum(d[0] for d in diversity.values()) / max(1, len(diversity)), 2),
        "avg_samples_generated_per_failed_input":
            round(sum(d[1] for d in diversity.values()) / max(1, len(diversity)), 2),
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bench", required=True)
    ap.add_argument("--greedy", required=True, help="temperature-0 generations")
    ap.add_argument("--samples", required=True, help="sampled generations, each row carrying `s`")
    ap.add_argument("--name", default="benchmark")
    ap.add_argument("--out", required=True)
    ap.add_argument("--skip-self-test", action="store_true",
                    help="only for iterating; a reported number must never use it")
    a = ap.parse_args()

    if not a.skip_self_test:
        bench, _ = load_bench(a.bench)
        st = self_test(bench, "ours")
        print(f"self-test: preflight {st['preflight_pct']}%  mutation kill "
              f"{st['mutation_kill_rate_pct']}%", file=sys.stderr, flush=True)
        require_sound(st)

    res = grade(a.name, a.bench, a.greedy, a.samples)
    Path(a.out).parent.mkdir(parents=True, exist_ok=True)
    Path(a.out).write_text(json.dumps(res, indent=2))
    print(json.dumps(res, indent=2))


if __name__ == "__main__":
    main()