#!/usr/bin/env python3
"""Rescore harness samples with format-robust extraction.

Why: thinking models break stock filters two ways — (1) they echo the
answer prefix back (`The best answer is C` vs target `C`), failing exact
match on a correct answer; (2) long thought can leave `content` empty on a
tight token budget (counted separately as UNANSWERED, not wrong).
This script scores the raw `resps` text both sides can audit.

Letter tasks (ARC-chat, MMLU-generative): last bare A-D token vs target.
GSM8K: `#### N` if present, else last integer vs target number.
Usage: python3 rescore.py <samples.jsonl> [letter|gsm8k]
"""
import json
import re
import sys

path, kind = sys.argv[1], (sys.argv[2] if len(sys.argv) > 2 else "letter")
rows = [json.loads(l) for l in open(path)]

correct = wrong = unanswered = 0
for r in rows:
    raw = r.get("resps") or []
    text = raw[0] if raw and isinstance(raw[0], str) else (
        raw[0][0] if raw and raw[0] else "")
    if not (text or "").strip():
        unanswered += 1
        continue
    if kind == "gsm8k":
        m = re.search(r"####\s*(-?\d[\d,]*)", text)
        if not m:
            nums = re.findall(r"-?\d[\d,]*", text)
            m = [None, nums[-1]] if nums else None
        pred = (m[1] if isinstance(m, list) else m.group(1)).replace(",", "") \
            if m else None
        tm = re.search(r"####\s*(-?\d[\d,]*)", r.get("target", ""))
        tgt = tm.group(1).replace(",", "") if tm else None
        correct += pred is not None and tgt is not None and pred == tgt
        wrong += not (pred is not None and tgt is not None and pred == tgt)
    else:  # letter tasks: explicit "answer is X" wins, else last bare letter
        m = re.search(r"[Aa]nswer is\s*\**([A-D])", text)
        pred = m.group(1) if m else None
        if pred is None:
            letters = re.findall(r"\b([A-D])\b", text)
            if not letters:
                unanswered += 1
                continue
            pred = letters[-1]
        if pred == (r.get("target") or "").strip():
            correct += 1
        else:
            wrong += 1

n = len(rows)
print(f"{path.split('/')[-1]}: n={n} correct={correct} wrong={wrong} "
      f"unanswered={unanswered} acc={correct / n:.3f} "
      f"(answered-acc={correct / max(1, correct + wrong):.3f})")