"""Custom lm-eval task utils for MATH-500 (HuggingFaceH4/MATH-500).

Scoring (locked 2026-07-29, pre-registered -- not changed after seeing 3-arm numbers):
  1. math_verify.verify on parse(gold) vs parse(full_generation) (bidirectional).
     math_verify robustly extracts the answer (\\boxed, last expr) and checks latex
     equivalence (handles 1/2 == 0.5, etc.).
  2. Fallback: clean-string equality of the boxed/Final-Answer extraction vs gold,
     stripping formatting (\\left \\right \\text ^\\circ braces spaces). This recovers
     tuples/points ('(3,pi/2)' == '\\left(3,\\frac{\\pi}{2}\\right)') and text answers
     ('Evelyn' == '\\text{Evelyn}') that math_verify's parser misses.
  NB: minerva's is_equiv is NOT reflexive on non-trivial latex; normalize_final_answer
  mangles '\\left'->'\\le', so neither is used for the decision.
Few-shot = minerva's fixed 4 examples. MATH-500 ships gold `answer` directly.
"""
from math_verify import parse, verify
from math_verify.parser import LatexExtractionConfig, ExprExtractionConfig
from lm_eval.tasks.minerva_math.utils import (
    list_fewshot_samples,
    last_boxed_only_string,
    remove_boxed,
    get_unnormalized_answer,
)

_CFG = [LatexExtractionConfig(), ExprExtractionConfig()]
_STRIP = ["\\left", "\\right", "\\!", "\\,", "\\;", "\\ ", "\\text",
          "\\mathrm", "\\displaystyle", "$", "^\\circ", "^{\\circ}"]


def doc_to_text(doc: dict) -> str:
    return "Problem:\n" + doc["problem"] + "\n\nSolution:"


def _clean(s) -> str:
    s = str(s)
    for a in _STRIP:
        s = s.replace(a, "")
    return s.replace("{", "").replace("}", "").replace(" ", "").replace("degrees", "").strip().lower()


def _boxed_or_final(gen: str) -> str:
    try:
        b = last_boxed_only_string(gen)
        if b:
            return remove_boxed(b)
    except Exception:
        pass
    return get_unnormalized_answer(gen)


def process_results(doc: dict, results):
    gen = results[0]
    gold = doc["answer"]
    ok = 0
    try:
        gp = parse(gold, extraction_config=_CFG)
        pp = parse(gen, extraction_config=_CFG)
        if verify(gp, pp) or verify(pp, gp):
            ok = 1
    except Exception:
        pass
    if not ok:
        pred = _boxed_or_final(gen)
        if _clean(pred) != "" and _clean(pred) == _clean(gold):
            ok = 1
    return {"exact_match": ok}