import re

# Keep integer conversion bounded; 12 digits is ample for this toy task.
FINAL_INTEGER = re.compile(
    r"<answer>\s*([+-]?[0-9]{1,12})\s*</answer>\s*\Z"
)


def extract_final_count(response: str) -> int | None:
    # Don't let a response submit several answers and win on the last one.
    if response.count("<answer>") != 1 or response.count("</answer>") != 1:
        return None

    match = FINAL_INTEGER.search(response)
    if match is None:
        return None

    return int(match.group(1))


def correctness_reward(response: str, expected_count: int) -> float:
    count = extract_final_count(response)
    if count is None:
        return 0.0

    return float(count == expected_count)


assert correctness_reward("48 - 6 = 42. <answer>42</answer>", 42) == 1.0
assert correctness_reward("I considered 42. <answer>40</answer>", 42) == 0.0
assert correctness_reward("<answer>40</answer><answer>42</answer>", 42) == 0.0
assert correctness_reward("<answer>0</answer>", 0) == 1.0