"""Candidate selection for code: K independent answers (each already through the deliberation fallback) are run on
the statement's public tests; the kept answer is the first candidate that passes all of them, else the candidate
passing the most (ties: the earliest sample). Only public tests are used, never the hidden grading tests.
The output keeps one record per (problem, sample 0) so it is graded like a single answer; n_tokens = sum over the
candidates actually generated (all K are counted, the cost of selection).

usage: code_select.py --src results/lcb_dm/lcb.jsonl results/lcb_dm2/lcb.jsonl ... --items data/lcb/lcb6.jsonl
       --out results/lcb_sel/lcb.jsonl [--k 4]
"""
import argparse
import collections
import json
import os
import sys
from concurrent.futures import ThreadPoolExecutor

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from grade import extract_code  # noqa: E402
from public_tests import run_public  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--src", nargs="+", required=True)
ap.add_argument("--items", nargs="+", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--k", type=int, default=4)
a = ap.parse_args()
items = {}
for p in a.items:
    for l in open(p, encoding="utf-8"):
        r = json.loads(l)
        items[r["id"]] = r
cands = collections.defaultdict(list)
for p in a.src:
    for l in open(p, encoding="utf-8"):
        r = json.loads(l)
        cands[r["id"]].append(r)


def kind(it):
    return "lcb" if it["task"] == "lcb" else "humaneval" if it["task"] == "humaneval" else \
        "mbpp" if it["task"] in ("mbpp", "code_fn") else "code_io"


def pick(pid):
    rs = sorted(cands[pid], key=lambda r: r["sample"])[: a.k]
    it = items[pid]
    best, best_key = rs[0], (-1, 0)
    for i, r in enumerate(rs):
        fin = r["completion"].split("</think>", 1)[1] if "</think>" in r["completion"] else ""
        code = extract_code(fin)
        ok, tot = (run_public(kind(it), code, it["prompt"], it)[:2] if code else (0, 1))
        if (ok, -i) > best_key:
            best, best_key = r, (ok, -i)
        if tot and ok == tot:
            break
    return {**best, "sample": 0, "n_tokens": sum(r["n_tokens"] for r in rs), "picked": best["sample"],
            "public_passed": best_key[0]}


with ThreadPoolExecutor(12) as ex:
    out = list(ex.map(pick, list(cands)))
os.makedirs(os.path.dirname(a.out), exist_ok=True)
with open(a.out, "w", encoding="utf-8") as f:
    for r in out:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")
print(len(out), "problems; picked sample distribution:", collections.Counter(r["picked"] for r in out))