import pytest

from inspect_evals.mathematics.utils import extract_answer

test_cases = [
    ("Answer: $\\frac{1}{2}$", "$\\frac{1}{2}$"),
    ("ANSWER: 42.", "42."),  # note trailing period is included in extracted answer
    (
        "**Final answer:**\n\nANSWER: 42",
        "42",
    ),  # gpt 4o/4.1 families tend to produce this format
    ("I think the result equals $x^2 + 2x + 1$.", ""),
    ("The solution is $\\sqrt{2}$.", ""),
    ("Therefore, $\\boxed{42}$ is our answer.", ""),
    ("This is a multi-line answer: below\n\nMy answer is ANSWER: 42", "42"),
]


@pytest.mark.parametrize("completion, expected_answer", test_cases)
def test_extract_answer(completion, expected_answer):
    """
    Test extract_answer with different completions.

    The prompt specifies
    "put your answer on its own line at the end in the form "ANSWER: $ANSWER" (without quotes)"
    """
    assert extract_answer(completion) == expected_answer
