#!/usr/bin/env python3
"""Generate best-of-N candidates for the rows greedy decoding failed, with verified early stop.

Sampling every row would be waste: a row the greedy answer already certifies needs no candidates,
because certification is a PROOF and nothing is gained by a second opinion. So this samples only
the greedy failures, and stops as soon as one candidate certifies -- the certified@N curve depends
only on the INDEX of the first passing sample, so stopping after it is lossless for every N.

That early stop is why the published run averaged ~3.8 generations per failed input rather than 31.

    PYBYTECODE_ENDPOINT=... PYBYTECODE_MODEL=... ./bestofn_gen.py \\
        --bench ../benchmarks/csn-3.12-licensed/bench.jsonl \\
        --greedy gen.jsonl --out boN.jsonl --max-samples 31

Resumable: rows already present in --out are skipped.
"""
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, strip_fences  # noqa: E402
from config import ENDPOINT, MODEL  # noqa: E402
from generate import INSTRUCTION, complete  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bench", required=True)
    ap.add_argument("--greedy", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--max-samples", type=int, default=31)
    ap.add_argument("--temperature", type=float, default=0.8)
    ap.add_argument("--max-tokens", type=int, default=2048)
    a = ap.parse_args()

    bench, _ = load_bench(a.bench)
    greedy = {r["i"]: r for r in load_jsonl(a.greedy)}

    failed = [
        i for i in sorted(bench)
        if not (i in greedy and ours_ok(strip_fences(greedy[i]["got"]), bench[i]["expected"]))
    ]
    print(f"greedy certified {len(bench) - len(failed)}/{len(bench)}; "
          f"sampling {len(failed)} failures", file=sys.stderr, flush=True)

    out_path = Path(a.out)
    done: set[int] = set()
    if out_path.exists():
        for r in load_jsonl(out_path):
            done.add(r["i"])
        print(f"resuming: {len(done)} rows already sampled", file=sys.stderr)

    gen_count = 0
    recovered = 0
    with out_path.open("a") as f:
        for n, i in enumerate(failed):
            if i in done:
                continue
            prompt = f"{INSTRUCTION}\n\n{bench[i]['input']}"
            for s in range(a.max_samples):
                got = complete(prompt, a.temperature, a.max_tokens, s)
                gen_count += 1
                f.write(json.dumps({"i": i, "s": s, "got": got}) + "\n")
                f.flush()
                if ours_ok(strip_fences(got), bench[i]["expected"]):
                    recovered += 1
                    break  # lossless: certified@N depends only on this index
            if n % 10 == 0:
                print(f"  {n}/{len(failed)} failures processed, {gen_count} generations, "
                      f"{recovered} recovered", file=sys.stderr, flush=True)

    print(f"DONE: {gen_count} generations, {recovered}/{len(failed)} recovered "
          f"(mean {gen_count / max(1, len(failed)):.2f} gens/failed input)", file=sys.stderr)


if __name__ == "__main__":
    main()