"""Example: test a simple reward function with rewardprobe."""

from rewardprobe import Probe
from rewardprobe.report import print_report


# A reward function with a subtle bug: substring match
def buggy_reward(completion, answer):
    return 1.0 if answer in completion else 0.0


# A better reward function: exact match with normalization
def better_reward(completion, answer):
    c = completion.strip().lower()
    a = str(answer).strip().lower()
    if not c or not a:
        return 0.0
    return 1.0 if c == a else 0.0


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

probe = Probe()

print("=== Testing buggy_reward (substring match) ===\n")
report = probe.test_fn(buggy_reward, tasks)
print_report(report)

print("\n=== Testing better_reward (exact match) ===\n")
report = probe.test_fn(better_reward, tasks)
print_report(report)
