# MATH grader v2 (Oct 27 2026). Superset of math_grade.py: anything v1 accepts, v2 accepts.
# Adds conservative equivalences for the 16 E17 false negatives: base subscripts, option letters,
# units, decimal vs fraction (incl. a/b), \frac9{19}, leading 'x \in', ordinals, unordered lists and +-,
# and element-wise pmatrix comparison.
import re
from math_grade import last_boxed, norm as norm1, num as num1, correct as correct_v1

def norm(s):
    s = norm1(s)
    if s is None: return None
    s = re.sub(r'\\(?:mbox|textbf|mathrm)\{([^{}]*)\}', r'\1', s)
    s = re.sub(r'\\text\{([^{}]*)\}', r'\1', s)
    s = re.sub(r'\\frac(\d)\{', r'\\frac{\1}{', s)            # \frac9{19}
    s = re.sub(r'\\frac\{([^{}]+)\}(\d)(?!\d)', r'\\frac{\1}{\2}', s)
    s = re.sub(r'^[a-zA-Z]\\in', '', s)                         # x \in [..]
    s = re.sub(r'_\{(\d+)\}$', r'_\1', s)                       # 4210_{5} -> 4210_5
    m = re.fullmatch(r'\(?([A-E])\)?', s)                       # (C) / C
    if m: s = m.group(1)
    s = re.sub(r'^(-?\d+)(?:\^\{?(?:\\mathrm\{)?(?:st|nd|rd|th)\}?\}?|st|nd|rd|th)$', r'\1', s)  # ordinals
    m = re.fullmatch(r'(-?[\d.,]+|\\frac\{\d+\}\{\d+\})((?:[a-zA-Z]+)(?:\^\{?\d\}?)?)', s)   # trailing units
    if m and m.group(2) not in ('i', 'e', 'pi'): s = m.group(1)
    return s

def num(s):
    v = num1(s)
    if v is not None: return v
    m = re.fullmatch(r'(-?\d+)/(\d+)', s)
    if m and int(m.group(2)): return int(m.group(1)) / int(m.group(2))
    m = re.fullmatch(r'-\\frac\{(\d+)\}\{(\d+)\}', s)
    if m: return -int(m.group(1)) / int(m.group(2))
    return None

def eq(a, b):
    if a == b: return True
    x, y = num(a), num(b)
    return x is not None and y is not None and abs(x - y) < 1e-9

def split_top(s):
    out, d, cur = [], 0, ''
    for ch in s:
        if ch in '([{': d += 1
        elif ch in ')]}': d -= 1
        if ch == ',' and d == 0: out.append(cur); cur = ''
        else: cur += ch
    out.append(cur)
    return out

def expand_pm(parts):
    out = []
    for p in parts:
        m = re.fullmatch(r'(.+?)\\pm(.+)', p)
        out += [m.group(1) + '+' + m.group(2), m.group(1) + '-' + m.group(2)] if m else [p]
    return out

def pmat(s):
    m = re.fullmatch(r'\\begin\{[pb]matrix\}(.*)\\end\{[pb]matrix\}', s)
    return [e for e in re.split(r'\\\\|&', m.group(1))] if m else None

def correct(text, gold):
    if correct_v1(text, gold): return True
    p, g = norm(last_boxed(text)), norm(gold)
    if p is None: return False
    if eq(p, g): return True
    # base subscript present on only one side
    bp, bg = re.fullmatch(r'(\w+)_(\d+)', p), re.fullmatch(r'(\w+)_(\d+)', g)
    if bg and not bp and p == bg.group(1): return True
    if bp and not bg and g == bp.group(1): return True
    # matrices element-wise
    mp, mg = pmat(re.sub(r'\s+', '', last_boxed(text))), pmat(re.sub(r'\s+', '', gold))
    if mp is not None and mg is not None:  # raw (pre-norm) so \\\\ row breaks survive
        return len(mp) == len(mg) and all(eq(norm(a), norm(b)) for a, b in zip(mp, mg))
    # unordered lists of roots (gold not an interval/tuple), with +- expansion
    if not g.startswith(('(', '[')) and not p.startswith(('(', '[')):
        A, B = expand_pm(split_top(p)), expand_pm(split_top(g))
        if len(A) == len(B) and len(A) > 1:
            A, B = sorted(norm(a) for a in A), sorted(norm(b) for b in B)
            used = [False] * len(B)
            for a in A:
                k = next((j for j, b in enumerate(B) if not used[j] and eq(a, b)), None)
                if k is None: return False
                used[k] = True
            return True
    return False