"""Answer extraction for BBH when the model is a chat model.

The Base protocol relies on the model CONTINUING the 3-shot "A:" pattern, so the
whole continuation is the answer and `strip` (remove_whitespace + take_first) is
enough. LLaDA-8B-Instruct under a chat template does not continue -- it answers,
usually with a short chain of thought and a concluding sentence. Measured
2026-08-18, dense, 108 samples: 33% of responses contain "answer is".

The cost of not extracting is not a small bias, it is a wrong verdict. snarks doc 0:
    target       "(B)"
    response     "...contains verbal irony, which is a hallmark of sarcasm.
                  So the answer is (B)."
    exact_match  0.0
-- a correct answer scored wrong purely on format.

This filter is ADDITIVE: the stock `strip` filter stays, so a Base run keeps the
exact protocol the finished Base table used, and analysis/lmeval_to_items.py picks
the right one per family.
"""
import re

# "So the answer is (B)." / "The answer is: valid" / "answer is 5"
# "is", ":", and "is:" are all seen in practice, hence the optional colon after `is`.
_ANSWER_IS = re.compile(r"(?:the\s+)?answer\s*(?:is\s*:?|:)\s*(.+?)\s*[.。]?\s*$",
                        re.I | re.S)
# last standalone option letter in parentheses
_PAREN = re.compile(r"\(([A-Za-z])\)")


def _one(text):
    if not text:
        return ""
    s = text.strip()
    # Prefer an explicit concluding statement, taking the LAST one (the model
    # sometimes restates the question's phrasing before concluding).
    tail = s.split("\n")[-1] if "\n" in s else s
    for cand in (tail, s):
        m = None
        for m in _ANSWER_IS.finditer(cand):
            pass                      # keep the last match
        if m:
            ans = m.group(1).strip()
            # collapse a multi-line tail down to its first line
            ans = ans.split("\n")[0].strip()
            if ans:
                return ans
    # No concluding phrase: fall back to the last "(X)" if the task is multiple
    # choice, else the stock behaviour (the stripped string itself).
    paren = _PAREN.findall(s)
    if paren:
        return f"({paren[-1]})"
    return s


def extract_bbh_answer(resps, docs):
    return [[_one(r) for r in resp] for resp in resps]