"""Failure anatomy of a LiveCodeBench run: for each problem, does the final program pass the public examples, and on
the hidden tests is it right, wrong, too slow, or crashing?  usage: lcb_failtypes.py <results tag>"""
import collections
import glob
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, final_answer  # noqa: E402
from lcb_tests import run_lcb  # noqa: E402

tag = sys.argv[1]
items = {json.loads(l)["id"]: json.loads(l) for l in open("data/lcb/lcb6.jsonl", encoding="utf-8")}
recs = []
for p in glob.glob(f"results/{tag}/*.jsonl"):
    if not p.endswith("graded.jsonl"):
        recs += [json.loads(l) for l in open(p, encoding="utf-8")]
recs = [r for r in recs if r["id"] in items and r.get("sample", 0) == 0]


def kind(msg):
    if msg is None:
        return "ok"
    m = msg.lower()
    if "time limit" in m or "timeout" in m:
        return "slow"
    if "error" in m or "traceback" in m or "exception" in m:
        return "crash"
    return "wrong"


def one(r):
    it = items[r["id"]]
    code = extract_code(final_answer(r["completion"]) or "")
    if not code:
        return r["id"], it["difficulty"], "no_code", "-"
    k, n, msg = run_lcb(code, it, it["public"])
    pub = "pub_ok" if k == n else "pub_fail"
    k, n, msg = run_lcb(code, it, it["hidden"])
    return r["id"], it["difficulty"], pub, kind(msg)


with ThreadPoolExecutor(8) as ex:
    rows = list(ex.map(one, recs))
out = collections.Counter((d, p, h) for _, d, p, h in rows)
for k in sorted(out):
    print(k, out[k])
json.dump(rows, open(f"results/{tag}/failtypes.json", "w"))