"""Reward functions a program can compute, in the shape TRL's GRPOTrainer expects.

Purpose: the course's reward library. Every function here takes the keyword arguments
    TRL passes to a reward function - `completions` plus any extra columns of the
    dataset - and returns one float per completion. Nothing here asks a model for an
    opinion: each reward is arithmetic, a regular expression or a subprocess that
    either passes or fails, so the same completion always scores the same.
Platform: all (pure Python; the unit-test reward needs a POSIX system for the
    resource limits, so on Windows run it inside WSL2)
Minimum memory: 8 GB
Assumes: Python 3.10 or newer. No third-party packages. `test-rewards.py` sits beside
    this file and is the proof that the functions do what the lesson says they do.

Usage: imported by train-grpo.py and eval-pass-at-1.py:
           from rewards import build_reward_functions, extract_final_answer
       run directly to score the built-in worked examples, including the
       reward-hacking ones, and print a table:
           python3 rewards.py --demo
           python3 rewards.py --demo --show-sandbox
"""

from __future__ import annotations

import argparse
import json
import math
import os
import re
import subprocess
import sys
import tempfile
from typing import Any, Callable, Iterable, Optional, Sequence

# --------------------------------------------------------------------------------------
# Reading an answer out of a completion
# --------------------------------------------------------------------------------------

# Three ways a model marks its final answer, in the order this course prefers them.
# GSM8K's own solutions end with "#### 18", the maths literature uses \boxed{18}, and
# the course's own format reward asks for "Answer: 18".
FINAL_PATTERNS = (
    re.compile(r"####\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),
    re.compile(r"\\boxed\{\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)\s*\}"),
    re.compile(r"(?i)\banswer\s*[:=]\s*\$?(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),
)
ANY_NUMBER = re.compile(r"-?[0-9][0-9,]*(?:\.[0-9]+)?")

# The two-block shape the format reward asks for. The working comes first inside
# <think> tags, then one line beginning "Answer:" with nothing after the number.
FORMAT_PATTERN = re.compile(
    r"^\s*<think>.*?</think>\s*Answer:\s*-?[0-9][0-9,]*(?:\.[0-9]+)?\s*$",
    re.DOTALL,
)


def text_of(completion: Any) -> str:
    """TRL hands completions as strings or as message lists; this flattens both."""
    if isinstance(completion, str):
        return completion
    if isinstance(completion, list):
        return "".join(str(m.get("content", "")) for m in completion if isinstance(m, dict))
    if isinstance(completion, dict):
        return str(completion.get("content", ""))
    return str(completion)


def extract_final_answer(text: str, fallback_to_last_number: bool = True) -> Optional[str]:
    """The number the completion is claiming as its answer, or None.

    Marked answers win over unmarked ones. The fallback to "the last number in the
    text" is convenient and slightly dangerous: a model that ends with a page of
    working can score by accident. The lesson says to turn it off once the format
    reward has taught the model to mark its answer.
    """
    for pattern in FINAL_PATTERNS:
        matches = pattern.findall(text)
        if matches:
            return matches[-1].replace(",", "")
    if fallback_to_last_number:
        numbers = ANY_NUMBER.findall(text)
        if numbers:
            return numbers[-1].replace(",", "")
    return None


def as_number(value: Any) -> Optional[float]:
    """Parse a string or number into a float, or None if it is not one."""
    if value is None:
        return None
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return float(value)
    text = str(value).strip().replace(",", "").rstrip(".")
    if text.startswith("$"):
        text = text[1:]
    try:
        return float(text)
    except ValueError:
        return None


# --------------------------------------------------------------------------------------
# Exact match and numeric rewards
# --------------------------------------------------------------------------------------


def exact_match_reward(completions, answer, **kwargs) -> list[float]:
    """1.0 when the extracted answer string equals the reference string, else 0.0.

    Strict, cheap and unforgiving: "18.0" does not equal "18", and neither does
    "the answer is 18" if the extractor picked up something else. Use it when the
    reference answers are canonical strings and you want no tolerance at all.
    """
    out = []
    for completion, reference in zip(completions, answer):
        found = extract_final_answer(text_of(completion))
        out.append(1.0 if found is not None and found == str(reference).strip() else 0.0)
    return out


def numeric_reward(tolerance: float = 1e-6, fallback_to_last_number: bool = True):
    """1.0 when the extracted number equals the reference within a tolerance.

    This is the reward the maths lab trains against. It fixes the "18.0" problem
    above by comparing numbers rather than strings, and it is still a program: no
    judgement, no model, no ambiguity about what it will return.
    """

    def reward(completions, answer, **kwargs) -> list[float]:
        out = []
        for completion, reference in zip(completions, answer):
            want = as_number(reference)
            got = as_number(extract_final_answer(text_of(completion), fallback_to_last_number))
            if want is None or got is None:
                out.append(0.0)
            else:
                out.append(1.0 if math.isclose(got, want, rel_tol=0.0, abs_tol=tolerance) else 0.0)
        return out

    reward.__name__ = "numeric_reward"
    return reward


# --------------------------------------------------------------------------------------
# Format and length rewards
# --------------------------------------------------------------------------------------


def format_reward(pattern: re.Pattern[str] = FORMAT_PATTERN, value: float = 0.2):
    """A small reward for answering in the shape the task asked for.

    Keep it small. A format reward is a hint about where to put the answer, not a
    goal in itself; if it is worth as much as being right, the model will learn to
    produce beautifully formatted wrong answers, which is the first reward-hacking
    example in the lesson.
    """

    def reward(completions, **kwargs) -> list[float]:
        return [value if pattern.match(text_of(c)) else 0.0 for c in completions]

    reward.__name__ = "format_reward"
    return reward


def length_reward(target_tokens: int = 256, penalty: float = 0.2, tokens_per_word: float = 1.35):
    """A penalty that grows once a completion runs past a soft length budget.

    Zero up to the budget, then linear to -`penalty` at twice the budget, then flat.
    Two reasons to have it: GRPO's own objective has a documented length bias (the
    Dr. GRPO paper), and a completion that runs to the generation cap is truncated,
    so its answer never appears and it scores zero for a reason that has nothing to
    do with reasoning. The word count is an approximation of the token count; the
    ratio is a rough constant for English prose, not a measurement.
    """

    def reward(completions, **kwargs) -> list[float]:
        out = []
        for completion in completions:
            words = len(text_of(completion).split())
            approx_tokens = words * tokens_per_word
            over = max(0.0, approx_tokens - target_tokens) / max(1.0, float(target_tokens))
            out.append(-penalty * min(1.0, over))
        return out

    reward.__name__ = "length_reward"
    return reward


# --------------------------------------------------------------------------------------
# Unit tests as a reward, in a subprocess with limits
# --------------------------------------------------------------------------------------

CODE_BLOCK = re.compile(r"```(?:python)?\s*\n(.*?)```", re.DOTALL)

# The child process is started with this preamble so that a submission cannot pass by
# printing the expected output: the tests run after the candidate code, in the same
# namespace, and the harness prints one line the parent parses.
RUNNER_TEMPLATE = """\
import resource, sys
resource.setrlimit(resource.RLIMIT_CPU, ({cpu}, {cpu}))
resource.setrlimit(resource.RLIMIT_AS, ({mem}, {mem}))
resource.setrlimit(resource.RLIMIT_FSIZE, ({fsize}, {fsize}))
resource.setrlimit(resource.RLIMIT_NPROC, ({nproc}, {nproc}))
sys.stdout = open(__import__("os").devnull, "w")
namespace = {{}}
exec(compile(open({candidate!r}).read(), "candidate.py", "exec"), namespace)
exec(compile(open({tests!r}).read(), "tests.py", "exec"), namespace)
sys.stderr.write("COURSE-TESTS-PASSED\\n")
"""


def extract_code(text: str) -> str:
    """The last fenced Python block, or the whole completion if there is no fence."""
    blocks = CODE_BLOCK.findall(text)
    return blocks[-1] if blocks else text


def run_tests_sandboxed(
    code: str,
    tests: str,
    timeout_s: int = 5,
    cpu_seconds: int = 5,
    memory_mb: int = 512,
    max_output_bytes: int = 1 << 20,
    max_processes: int = 64,
) -> tuple[bool, str]:
    """Run `tests` against `code` in a separate interpreter with resource limits.

    Returns (passed, reason). The child runs with -I (isolated: no user site
    directory, no PYTHON* environment variables) and -S, in a fresh temporary
    directory, with a wall-clock timeout on top of the CPU limit, so an infinite
    loop, a runaway allocation, a fork bomb and a program that writes a large file
    are all bounded.

    What this does NOT do, and the lesson says so on the page: it is not a security
    boundary. The child runs as your user, on your filesystem, with your network.
    It can read any file you can read, open a socket, and delete your work. It stops
    accidents and cheap denial of service; it does not stop a program written to do
    harm. Part 24 builds the container-and-dedicated-user sandbox that an agent
    running untrusted code needs.
    """
    if os.name != "posix":
        return False, "resource limits need a POSIX system; run this inside WSL2 on Windows"

    with tempfile.TemporaryDirectory(prefix="course-reward-") as workdir:
        candidate_path = os.path.join(workdir, "candidate.py")
        tests_path = os.path.join(workdir, "tests.py")
        runner_path = os.path.join(workdir, "runner.py")
        with open(candidate_path, "w", encoding="utf-8") as handle:
            handle.write(code)
        with open(tests_path, "w", encoding="utf-8") as handle:
            handle.write(tests)
        with open(runner_path, "w", encoding="utf-8") as handle:
            handle.write(
                RUNNER_TEMPLATE.format(
                    cpu=cpu_seconds,
                    mem=memory_mb * 1024 * 1024,
                    fsize=max_output_bytes,
                    nproc=max_processes,
                    candidate=candidate_path,
                    tests=tests_path,
                )
            )
        try:
            done = subprocess.run(
                [sys.executable, "-I", "-S", runner_path],
                cwd=workdir,
                env={"PATH": os.environ.get("PATH", ""), "HOME": workdir},
                capture_output=True,
                text=True,
                timeout=timeout_s,
                check=False,
            )
        except subprocess.TimeoutExpired:
            return False, f"timed out after {timeout_s} s"
        except OSError as exc:
            return False, f"could not start the sandbox: {exc}"

    if "COURSE-TESTS-PASSED" in (done.stderr or ""):
        return True, "passed"
    tail = (done.stderr or "").strip().splitlines()
    return False, (tail[-1][:200] if tail else f"exit code {done.returncode}")


def unit_test_reward(timeout_s: int = 5, memory_mb: int = 512, partial_credit: bool = False):
    """1.0 when the completion's code passes the task's tests, else 0.0.

    The tests come from the dataset: each row carries a `tests` column holding the
    assertions to run against the candidate. `partial_credit` splits the test source
    on blank lines and scores the fraction of blocks that pass, which gives GRPO a
    gradient on tasks where all-or-nothing rewards make every rollout in a group
    score zero and the group's advantages collapse to nothing.
    """

    def reward(completions, tests, **kwargs) -> list[float]:
        out = []
        for completion, test_source in zip(completions, tests):
            code = extract_code(text_of(completion))
            if partial_credit:
                blocks = [b for b in re.split(r"\n\s*\n", str(test_source)) if b.strip()]
                if not blocks:
                    out.append(0.0)
                    continue
                passed = sum(1 for b in blocks if run_tests_sandboxed(code, b, timeout_s, memory_mb=memory_mb)[0])
                out.append(passed / len(blocks))
            else:
                ok, _ = run_tests_sandboxed(code, str(test_source), timeout_s, memory_mb=memory_mb)
                out.append(1.0 if ok else 0.0)
        return out

    reward.__name__ = "unit_test_reward"
    return reward


# --------------------------------------------------------------------------------------
# Combining rewards, and measuring one before training with it
# --------------------------------------------------------------------------------------


def build_reward_functions(
    kind: str = "maths",
    use_format: bool = True,
    use_length: bool = True,
    target_tokens: int = 256,
    tolerance: float = 1e-6,
    fallback_to_last_number: bool = False,
) -> tuple[list[Callable[..., list[float]]], list[float]]:
    """The reward list and weights the training script passes to GRPOTrainer.

    TRL sums the functions after weighting them with `reward_weights`, so the
    weights are the statement of what matters: correctness dominates, format is a
    hint, length is a brake. Returning them together keeps the two lists the same
    length, which TRL requires.
    """
    functions: list[Callable[..., list[float]]] = []
    weights: list[float] = []

    if kind == "maths":
        functions.append(numeric_reward(tolerance, fallback_to_last_number))
        weights.append(1.0)
    elif kind == "exact":
        functions.append(exact_match_reward)
        weights.append(1.0)
    elif kind == "code":
        functions.append(unit_test_reward())
        weights.append(1.0)
    else:
        raise ValueError(f"unknown reward kind {kind!r}; expected maths, exact or code")

    if use_format:
        functions.append(format_reward())
        weights.append(1.0)
    if use_length:
        functions.append(length_reward(target_tokens))
        weights.append(1.0)
    return functions, weights


def score(functions: Sequence[Callable[..., list[float]]],
          weights: Sequence[float],
          completions: Sequence[Any],
          **columns: Any) -> list[float]:
    """Apply every reward function and return the weighted sum, one per completion."""
    totals = [0.0] * len(completions)
    for function, weight in zip(functions, weights):
        values = function(completions=list(completions), **columns)
        for i, value in enumerate(values):
            totals[i] += weight * float(value)
    return totals


def evaluate_reward(functions: Sequence[Callable[..., list[float]]],
                    weights: Sequence[float],
                    cases: Iterable[dict]) -> dict:
    """Score a labelled set of completions and report whether the reward separates them.

    Each case is {"completion": ..., "should_score": "high"|"low", extra columns...}.
    A reward you have not run against known-good and known-bad answers is a reward
    you are about to train against blind: the gap between the two means is the
    signal GRPO has to work with, and any overlap is a way for the model to score
    without doing the task.
    """
    cases = list(cases)
    rows = []
    for case in cases:
        columns = {k: [v] for k, v in case.items() if k not in ("completion", "should_score", "note")}
        total = score(functions, weights, [case["completion"]], **columns)[0]
        rows.append({"note": case.get("note", ""), "expected": case["should_score"], "reward": round(total, 4)})

    highs = [r["reward"] for r in rows if r["expected"] == "high"]
    lows = [r["reward"] for r in rows if r["expected"] == "low"]
    mean = lambda xs: sum(xs) / len(xs) if xs else float("nan")  # noqa: E731 - one line, read once
    overlap = [r for r in rows if r["expected"] == "low" and highs and r["reward"] >= min(highs)]
    return {
        "rows": rows,
        "mean_high": round(mean(highs), 4),
        "mean_low": round(mean(lows), 4),
        "separation": round(mean(highs) - mean(lows), 4),
        "overlapping_low_cases": [r["note"] for r in overlap],
    }


# --------------------------------------------------------------------------------------
# Worked examples, including the ones that cheat
# --------------------------------------------------------------------------------------

DEMO_CASES = [
    {
        "note": "correct, in the requested format",
        "completion": "<think>Six boxes of four pencils is 24 pencils. Nine are used, so 15 remain.</think>\nAnswer: 15",
        "answer": "15",
        "should_score": "high",
    },
    {
        "note": "correct answer, no format",
        "completion": "Six times four is twenty-four, minus nine leaves fifteen. The answer is 15.",
        "answer": "15",
        "should_score": "high",
    },
    {
        "note": "wrong answer, perfect format (the format-reward hack)",
        "completion": "<think>Six boxes of four pencils is 24 pencils.</think>\nAnswer: 24",
        "answer": "15",
        "should_score": "low",
    },
    {
        "note": "every number in sight, hoping one is right (the last-number hack)",
        "completion": "Maybe 6, or 4, or 24, or 9, or 13, or 14, or 15",
        "answer": "15",
        "should_score": "low",
    },
    {
        "note": "right answer buried in a wall of text (the length hack)",
        "completion": "<think>" + ("Let me reconsider the problem once more. " * 120) + "</think>\nAnswer: 15",
        "answer": "15",
        "should_score": "low",
    },
    {
        "note": "refuses to answer",
        "completion": "<think>This needs more information.</think>\nAnswer: 0",
        "answer": "15",
        "should_score": "low",
    },
]

SANDBOX_CASES = [
    {
        "note": "a correct function",
        "completion": "```python\ndef add(a, b):\n    return a + b\n```",
        "tests": "assert add(2, 2) == 4\n\nassert add(-1, 1) == 0\n",
        "should_score": "high",
    },
    {
        "note": "an off-by-one function",
        "completion": "```python\ndef add(a, b):\n    return a + b + 1\n```",
        "tests": "assert add(2, 2) == 4\n\nassert add(-1, 1) == 0\n",
        "should_score": "low",
    },
    {
        "note": "an infinite loop (stopped by the timeout)",
        "completion": "```python\ndef add(a, b):\n    while True:\n        pass\n```",
        "tests": "assert add(2, 2) == 4\n",
        "should_score": "low",
    },
    {
        "note": "prints the expected output instead of computing it",
        "completion": "```python\nprint('all tests passed')\ndef add(a, b):\n    return None\n```",
        "tests": "assert add(2, 2) == 4\n",
        "should_score": "low",
    },
]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--demo", action="store_true", help="score the built-in worked examples")
    parser.add_argument("--show-sandbox", action="store_true",
                        help="also run the unit-test reward on four candidate functions")
    parser.add_argument("--target-tokens", type=int, default=256)
    parser.add_argument("--json", action="store_true", help="print the report as JSON")
    args = parser.parse_args()

    if not args.demo:
        print(__doc__)
        return

    functions, weights = build_reward_functions(
        kind="maths", target_tokens=args.target_tokens, fallback_to_last_number=True
    )
    report = evaluate_reward(functions, weights, DEMO_CASES)

    if args.show_sandbox:
        code_functions, code_weights = build_reward_functions(kind="code", use_format=False, use_length=False)
        report["sandbox"] = evaluate_reward(code_functions, code_weights, SANDBOX_CASES)

    if args.json:
        print(json.dumps(report, indent=2))
        return

    print(f"{'reward':>8}  {'expected':<9} case")
    for row in report["rows"]:
        print(f"{row['reward']:>8.3f}  {row['expected']:<9} {row['note']}")
    print(f"\nmean reward, should score high: {report['mean_high']}")
    print(f"mean reward, should score low:  {report['mean_low']}")
    print(f"separation: {report['separation']}")
    if report["overlapping_low_cases"]:
        print("\nThese cases scored at least as well as the worst good answer:")
        for note in report["overlapping_low_cases"]:
            print(f"  - {note}")
        print("Each one is a way for the model to score without doing the task.")

    if "sandbox" in report:
        print("\nunit-test reward:")
        for row in report["sandbox"]["rows"]:
            print(f"{row['reward']:>8.3f}  {row['expected']:<9} {row['note']}")


if __name__ == "__main__":
    main()