#!/usr/bin/env python3
"""Reward functions for GRPO / RLVR fine-tuning — design patterns and examples.

Demonstrates: composing format + accuracy rewards, shaping continuous signals,
and sanity-checking against reward hacking. Run standalone to test reward logic
against sample completions.
"""

from __future__ import annotations

import re


# ── Format reward: check for required answer envelope ───────────────────

def format_reward(completions: list[str], **kwargs) -> list[float]:
    """Reward completions that wrap final answer in \\boxed{...}."""
    rewards = []
    for completion in completions:
        has_boxed = bool(re.search(r"\\boxed\{.*?\}", completion))
        rewards.append(1.0 if has_boxed else 0.0)
    return rewards


# ── Accuracy reward: compare answer to ground truth ────────────────────

def accuracy_reward(completions: list[str], ground_truth: list[str], **kwargs) -> list[float]:
    """Extract boxed answer and compare to ground truth."""
    rewards = []
    for completion, truth in zip(completions, ground_truth):
        match = re.search(r"\\boxed\{(.*?)\}", completion)
        if not match:
            rewards.append(0.0)
            continue
        answer = match.group(1).strip()
        # Normalize: strip whitespace, standardize fractions if needed
        truth = truth.strip()
        if answer == truth:
            rewards.append(1.0)
        else:
            # Partial credit: exact match only for demo
            rewards.append(0.0)
    return rewards


# ── Shaped reward: continuous signal (fractional) ──────────────────────

def fraction_passing_reward(completions: list[str], test_cases: list[str], **kwargs) -> list[float]:
    """Reward by fraction of test cases that pass (for code generation tasks)."""
    rewards = []
    for completion in completions:
        passed = 0
        for test in test_cases:
            # In real RLVR, you'd execute the code in a sandbox
            # Here we simulate with a heuristic check
            if test.lower() in completion.lower():
                passed += 1
        rewards.append(passed / max(len(test_cases), 1))
    return rewards


# ── Composite reward ────────────────────────────────────────────────────

def composite_reward(
    completions: list[str],
    ground_truth: list[str] | None = None,
    format_weight: float = 0.3,
    accuracy_weight: float = 0.7,
    **kwargs,
) -> list[float]:
    """Combine format and accuracy rewards."""
    fmt = format_reward(completions)
    acc = accuracy_reward(completions, ground_truth=ground_truth or [""] * len(completions))
    return [format_weight * f + accuracy_weight * a for f, a in zip(fmt, acc)]


# ── Sanity checks for reward hacking ────────────────────────────────────

def sanity_check():
    """Verify rewards don't reward degenerate outputs."""
    completions = [
        r"The answer is \boxed{42}.",     # Correct format + answer
        r"\boxed{42}",                     # Minimal correct
        r"The answer is 42.",             # Missing boxed format
        r"\boxed{}\boxed{}\boxed{}",      # Multiple boxes, no content
        r"\boxed{42} is the answer.",     # Correct with extra text
        r"",                              # Empty
    ]
    ground_truth = ["42"] * len(completions)

    fmt = format_reward(completions)
    acc = accuracy_reward(completions, ground_truth=ground_truth)

    print("Format rewards:", fmt)
    print("Accuracy rewards:", acc)
    print("Composite:", composite_reward(completions, ground_truth=ground_truth))

    # Degenerate check
    degenerate = [r"\boxed{}"]  # Well-formed but empty — format passes, accuracy fails
    print("\nDegenerate check:")
    print("  Format:", format_reward(degenerate))
    print("  Accuracy:", accuracy_reward(degenerate, ground_truth=["42"]))
    print("  → Format reward passed but accuracy failed. Watch for this divergence.")


if __name__ == "__main__":
    sanity_check()
    print("\n✓ Reward function demo complete.")