"""Example: test a TRL-style batch reward function with rewardprobe.

TRL reward functions take (completions: list[str], **kwargs) -> list[float].
rewardprobe auto-detects this pattern — no special flags needed.
"""

import re
from rewardprobe import Probe
from rewardprobe.report import print_report


# A TRL-style reward function (batch signature)
def accuracy_reward(completions, solution, **kwargs):
    """Check if the boxed answer matches ground truth."""
    rewards = []
    for completion, sol in zip(completions, solution):
        match = re.search(r"\\boxed\{([^}]*)\}", completion)
        if match and match.group(1).strip() == str(sol).strip():
            rewards.append(1.0)
        else:
            rewards.append(0.0)
    return rewards


# A format-only reward function (no answer check)
def format_reward(completions, **kwargs):
    """Check if completion uses <think>...</think> format."""
    pattern = r"<think>.*?</think>"
    return [
        1.0 if re.search(pattern, c, re.DOTALL) else 0.0
        for c in completions
    ]


tasks = [
    {"prompt": "What is 2+2?", "answer": "4", "solution": "4"},
    {"prompt": "What is 3+3?", "answer": "6", "solution": "6"},
    {"prompt": "What is 5*5?", "answer": "25", "solution": "25"},
    {"prompt": "What is 7+3?", "answer": "10", "solution": "10"},
    {"prompt": "What is 6*7?", "answer": "42", "solution": "42"},
]

probe = Probe()

print("=== Testing accuracy_reward (TRL batch) ===\n")
report = probe.test_fn(accuracy_reward, tasks)
print_report(report)

print("\n=== Testing format_reward (format-only, no answer check) ===\n")
report = probe.test_fn(format_reward, tasks)
print_report(report)
