"""Answer extraction, grading and loop signals.

The extraction and the math-verify check are copied from the paper's harness
(external/lm_understand-logs/exp8_math/code/mathqa.py, MIT licence) so our numbers
are graded exactly the way the paper's were:

- the answer is the content of the LAST \\boxed{...} (or \\fbox{...}) in the text
  after the last </think> (or in the whole text when there is no </think>);
- correctness is math-verify 0.9.0: verify(parse('\\boxed{gold}'), parse('\\boxed{pred}')).

`repetitive_end` re-implements the rule described in mathqa.loop_signals, because
the module it imports (mcqa) is not in the public repo.
"""
from __future__ import annotations

import functools
import re
import threading
from collections import Counter


def answer_text(text: str) -> str:
    if "</think>" in text:
        text = text.split("</think>")[-1].strip()
    return text


_BOX_FORMS = re.compile(r"\\boxed|\\fbox")


def _box_at(text: str, idx: int):
    key = "\\boxed" if text.startswith("\\boxed", idx) else "\\fbox"
    rest = text[idx + len(key):]
    r = rest.lstrip()
    if r.startswith("{"):
        depth = 0
        for k, ch in enumerate(r):
            if ch == "{":
                depth += 1
            elif ch == "}":
                depth -= 1
                if depth == 0:
                    out = r[1:k].strip()
                    return out or None
        return None
    if key == "\\boxed" and rest.startswith(" "):
        out = rest.strip().split("$")[0].split("\n")[0].strip()
        return out or None
    return None


def last_boxed(text: str):
    idx = text.rfind("\\boxed")
    if idx < 0:
        idx = text.rfind("\\fbox")
        if idx < 0:
            return None
    return _box_at(text, idx)


def extract(text: str):
    """(answer or None, number of boxes in the answer text)."""
    t = answer_text(text)
    return last_boxed(t), len(_BOX_FORMS.findall(t))


def _timeout():
    return 5 if threading.current_thread() is threading.main_thread() else None


@functools.lru_cache(maxsize=1 << 16)
def _parsed(s: str, to):
    import math_verify
    return math_verify.parse("\\boxed{" + s + "}", parsing_timeout=to)


def is_correct(pred, gold) -> bool:
    if pred is None or gold is None:
        return False
    import math_verify
    to = _timeout()
    g, p = _parsed(gold, to), _parsed(pred, to)
    if not g or not p:
        return False
    try:
        return bool(math_verify.verify(g, p, timeout_seconds=to))
    except Exception:  # noqa: BLE001
        return False


# ---------------------------------------------------------------- loop signals
LOOP_MAX_PERIOD, LOOP_MIN_REPEATS, LOOP_MIN_SPAN = 200, 3, 40


def verbatim_loop_end(text: str) -> bool:
    """The text (minus its last, possibly cut, word) ends with one block of 1..200 words
    repeated back to back at least 3 times and over at least 40 words (mathqa.loop_signals)."""
    w = text.split()[:-1]
    n = len(w)
    for p in range(1, LOOP_MAX_PERIOD + 1):
        reps = max(LOOP_MIN_REPEATS, -(-LOOP_MIN_SPAN // p))
        if reps * p > n:
            break
        blk = w[n - p:]
        if all(w[n - (j + 1) * p:n - j * p] == blk for j in range(1, reps)):
            return True
    return False


def repetitive_end(text: str) -> bool:
    """In the last 200 words, the 4-word positions whose phrase occurs 3+ times there make up
    at least half of all positions (circling with small variations)."""
    w = text.split()[-200:]
    grams = [tuple(w[i:i + 4]) for i in range(len(w) - 3)]
    if not grams:
        return False
    c = Counter(grams)
    return sum(1 for g in grams if c[g] >= 3) >= len(grams) / 2


def tail_dup8(text: str, window: int = 512) -> float:
    """Share of word 8-gram positions in the last `window` words whose 8-gram already
    appeared earlier in that window. About 0 for fresh text, close to 1 for a loop."""
    w = text.split()[-window:]
    if len(w) < 9:
        return 0.0
    seen, dup, total = set(), 0, 0
    for i in range(len(w) - 7):
        g = tuple(w[i:i + 8])
        total += 1
        if g in seen:
            dup += 1
        seen.add(g)
    return dup / total


def ifeval_follow(row: dict, text: str, thinking: bool):
    """IFEval strict prompt-level check on the answer after </think> (as the paper did):
    with thinking on and no </think>, the answer is empty and fails."""
    from lm_eval.tasks.ifeval import instructions_registry as IR

    resp = text.split("</think>")[-1].strip() if "</think>" in text else ("" if thinking else text.strip())
    follow = []
    for iid, kw in zip(row["instruction_id_list"], row["kwargs"]):
        ins = IR.INSTRUCTION_DICT[iid](iid)
        ins.build_description(**{k: v for k, v in (kw or {}).items() if v is not None})
        args = ins.get_instruction_args()
        if args and "prompt" in args:
            ins.build_description(prompt=row["prompt"])
        follow.append(bool(resp) and bool(ins.check_following(resp)))
    return all(follow), follow