"""Grades generations produced by gen.py against the frozen eval suite.

Only the final answer (text after </think>) is graded; a response whose thinking never closed counts as wrong.
Writes results/<tag>/graded.jsonl (one row per sample) and prints per-task accuracy.

usage: grade.py <tag>
"""
import glob
import json
import os
import re
import subprocess

import sandbox  # noqa: E402  (same folder)
import sys
import tempfile
import threading
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
TMP = os.path.join(tempfile.gettempdir(), "lfm3_codeexec")
os.makedirs(TMP, exist_ok=True)


def final_answer(completion):
    if "</think>" not in completion:
        return None
    return completion.split("</think>", 1)[1].strip()


def last_boxed(text):
    idx = text.rfind("\\boxed")
    if idx < 0:
        idx = text.rfind("\\fbox")
        if idx < 0:
            return None
    i = text.find("{", idx)
    if i < 0:
        return None
    depth = 0
    for j in range(i, len(text)):
        if text[j] == "{":
            depth += 1
        elif text[j] == "}":
            depth -= 1
            if depth == 0:
                return text[i + 1: j]
    return None


def grade_math(ans_text, gold):
    from math_verify import parse, verify
    if ans_text is None:
        return False
    pred = last_boxed(ans_text)
    if pred is None:
        return False
    try:
        g = parse("$" + gold + "$", parsing_timeout=None)
        p = parse("$" + pred + "$", parsing_timeout=None)
        if verify(g, p, timeout_seconds=None):
            return True
    except Exception:
        pass
    norm = lambda s: re.sub(r"\\text\{(.*?)\}|\\!|\s|\$|\\left|\\right|\.$", r"\1", s).replace("dfrac", "frac")
    return norm(pred) == norm(gold)


def grade_gsm8k(ans_text, gold):
    if ans_text is None:
        return False
    pred = last_boxed(ans_text)
    if pred is None:
        nums = re.findall(r"-?\d[\d,]*\.?\d*", ans_text)
        if not nums:
            return False
        pred = nums[-1]
    pred = re.sub(r"\\text\{.*?\}|\\\$|\$|,|\\%|%|\\!|\s", "", pred)
    m = re.search(r"-?\d+(\.\d+)?", pred.replace("{", "").replace("}", ""))
    if not m:
        return False
    try:
        return abs(float(m.group(0)) - float(gold)) < 1e-6
    except ValueError:
        return False


def extract_code(ans_text):
    if ans_text is None:
        return None
    blocks = re.findall(r"```(?:python|py|Python)?\s*\n(.*?)```", ans_text, re.S)
    if not blocks:
        return None
    return max(blocks, key=len) if len(blocks) > 1 and "def " not in blocks[-1] else blocks[-1]


def run_python(src, timeout):
    fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        f.write(src)
    try:
        r = sandbox.run([sys.executable, "-I", path], capture_output=True, timeout=timeout, cwd=TMP,
                           creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        return r.returncode == 0
    except subprocess.TimeoutExpired:
        return False
    finally:
        try:
            os.remove(path)
        except OSError:
            pass


def grade_humaneval(ans_text, item):
    code = extract_code(ans_text)
    if code is None:
        return False
    if f"def {item['entry_point']}" not in code:
        code = item["code_prompt"] + code
    src = code + "\n\n" + item["test"] + f"\n\ncheck({item['entry_point']})\n"
    return run_python(src, 60)


def grade_mbpp(ans_text, item):
    code = extract_code(ans_text)
    if code is None:
        return False
    src = "\n".join(item.get("test_imports") or []) + "\n" + code + "\n\n" + item["test"] + "\n"
    return run_python(src, 60)


def run_io(code, inp, out, timeout=6):
    fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        f.write(code)
    try:
        r = sandbox.run([sys.executable, "-I", path], input=inp.encode(), capture_output=True, timeout=timeout,
                           cwd=TMP, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        if r.returncode != 0:
            return False
        return r.stdout.decode(errors="replace").split() == out.split()
    except subprocess.TimeoutExpired:
        return False
    finally:
        try:
            os.remove(path)
        except OSError:
            pass


def grade_code_io(ans_text, item):
    code = extract_code(ans_text)
    if code is None:
        return False
    return all(run_io(code, i, o) for i, o in zip(item["inputs"], item["outputs"]))


def grade_lcb(ans_text, item):
    from lcb_tests import run_lcb
    code = extract_code(ans_text)
    if code is None:
        return False
    k, n, _ = run_lcb(code, item, item["hidden"])
    return k == n


_IFEVAL_LOCK = threading.Lock()


def grade_ifeval(ans_text, item):
    # lm_eval fills empty instruction arguments with random values and langdetect is random too: seed both under a
    # lock so that a response always gets the same grade
    import random
    from langdetect import DetectorFactory
    from lm_eval.tasks.ifeval.utils import process_results
    doc = {"key": 0, "prompt": item["prompt"], "instruction_id_list": item["instruction_id_list"],
           "kwargs": item["kwargs"]}
    with _IFEVAL_LOCK:
        random.seed(0)
        DetectorFactory.seed = 0
        res = process_results(doc, [ans_text or ""])
    return {"strict": res["prompt_level_strict_acc"], "loose": res["prompt_level_loose_acc"],
            "inst_strict": res["inst_level_strict_acc"]}


def grade_mmlupro(ans_text, gold):
    if ans_text is None:
        return False
    m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", ans_text)
    if not m:
        m = re.findall(r"\\boxed\{\(?([A-J])\)?\}", ans_text)
    return bool(m) and m[-1] == gold


def grade_bbh(ans_text, gold, prompt=None):
    """Exact match after extraction. With the prompt (flexible mode) it also accepts answers that state the same
    thing in an unambiguous other form: the text of the chosen option instead of its letter, a word list separated by
    commas, the whole Dyck sequence (input + completion) instead of the completion, plausible/not plausible for
    yes/no. Each rule maps the answer to exactly one gold-comparable form, so it cannot accept a different answer."""
    if ans_text is None:
        return False
    m = re.findall(r"answer is:?\s*(.+)", ans_text, re.I)
    pred = m[-1] if m else (last_boxed(ans_text) or "")
    pred = re.sub(r"\\boxed\{(.*)\}", r"\1", pred).strip().strip("*").strip().rstrip(".").strip()
    gold = gold.strip()
    norm = lambda s: re.sub(r"\s+", " ", s.lower().replace("`", "").replace("*", "")).strip().rstrip(".").strip()
    if re.fullmatch(r"\([A-R]\)", gold):
        mm = re.search(r"\(?([A-R])\)", pred) or re.fullmatch(r"([A-R])", pred)
        if mm or prompt is None:
            return bool(mm) and mm.group(1) == gold[1]
        opts = {norm(t): k for k, t in re.findall(r"^\(([A-R])\)\s*(.+?)\s*$", prompt, re.M)}
        return opts.get(norm(pred)) == gold[1]
    if norm(pred) == norm(gold):
        return True
    if prompt is None:
        return False
    if "Sort the following words alphabetically" in prompt:
        return re.findall(r"[^\s,]+", norm(pred)) == norm(gold).split()
    if "parentheses are closed properly" in prompt:
        inp = re.search(r"Input:\s*(.+)", prompt).group(1)
        sp = lambda s: re.sub(r"\s+", "", norm(s))
        return sp(pred) == sp(inp) + sp(gold)
    yn = {"plausible": "yes", "not plausible": "no", "implausible": "no"}
    return norm(gold) in ("yes", "no") and yn.get(norm(pred)) == norm(gold)


def load_suite(pattern=os.path.join(ROOT, "data", "eval", "*.jsonl")):
    items = {}
    for path in glob.glob(pattern):
        with open(path, encoding="utf-8") as f:
            for l in f:
                r = json.loads(l)
                items[r["id"]] = r
    return items


def grade_one(rec, item):
    a = final_answer(rec["completion"])
    t = item["task"]
    extra = {}
    if t in ("math500", "aime"):
        ok = grade_math(a, item["answer"])
    elif t == "gsm8k":
        ok = grade_gsm8k(a, item["answer"])
    elif t == "humaneval":
        ok = grade_humaneval(a, item)
    elif t in ("mbpp", "code_fn"):
        ok = grade_mbpp(a, item)
    elif t == "code_io":
        ok = grade_code_io(a, item)
    elif t == "lcb":
        ok = grade_lcb(a, item)
    elif t == "math":
        ok = grade_math(a, item["answer"])
    elif t == "ifeval":
        extra = grade_ifeval(a, item)
        ok = extra["strict"]
    elif t == "mmlupro":
        ok = grade_mmlupro(a, item["answer"])
    elif t == "bbh":
        ok = grade_bbh(a, item["answer"], item["prompt"])
        extra = {"strict": grade_bbh(a, item["answer"])}
    else:
        raise ValueError(t)
    return {"id": rec["id"], "task": t, "sample": rec["sample"], "correct": bool(ok), "n_tokens": rec["n_tokens"],
            "truncated": rec["truncated"], "closed_think": a is not None,
            "subtask": item.get("subtask") or item.get("category") or item.get("level"), **extra}


def main(tag, pattern=None):
    items = load_suite(pattern) if pattern else load_suite()
    recs = []
    for path in glob.glob(os.path.join(ROOT, "results", tag, "*.jsonl")):
        if path.endswith("graded.jsonl"):
            continue
        with open(path, encoding="utf-8") as f:
            recs += [json.loads(l) for l in f]
    with ThreadPoolExecutor(12) as ex:
        graded = list(ex.map(lambda r: grade_one(r, items[r["id"]]), recs))
    with open(os.path.join(ROOT, "results", tag, "graded.jsonl"), "w", encoding="utf-8") as f:
        for g in graded:
            f.write(json.dumps(g) + "\n")
    agg = defaultdict(list)
    for g in graded:
        agg[g["task"]].append(g)
    print(f"{'task':10s} {'n':>5s} {'acc':>7s} {'trunc':>6s} {'no</think>':>10s} {'avg_tok':>8s}")
    for t in sorted(agg):
        rows = agg[t]
        n = len(rows)
        print(f"{t:10s} {n:5d} {100 * sum(r['correct'] for r in rows) / n:7.2f} "
              f"{100 * sum(r['truncated'] for r in rows) / n:6.1f} {100 * sum(not r['closed_think'] for r in rows) / n:10.1f} "
              f"{sum(r['n_tokens'] for r in rows) / n:8.0f}")
        if t == "ifeval":
            print(f"{'  loose':10s} {'':5s} {100 * sum(r['loose'] for r in rows) / n:7.2f}   inst_strict "
                  f"{100 * sum(sum(r['inst_strict']) for r in rows) / sum(len(r['inst_strict']) for r in rows):.2f}")


if __name__ == "__main__":
    main(*sys.argv[1:])