"""Tests for the reward functions, because a reward with a bug trains the bug in.

Purpose: prove that every function in rewards.py returns what the lesson says it
    returns, including on the completions that try to cheat it. A reward function is
    the only thing standing between a reinforcement-learning run and hours of
    optimising the wrong quantity, and it is ordinary code with ordinary bugs.
Platform: all (pure Python and the standard library's unittest; the sandbox tests
    are skipped automatically on a non-POSIX system)
Minimum memory: 8 GB
Assumes: Python 3.10 or newer, and rewards.py in the same directory.

Usage: python3 test-rewards.py
       python3 test-rewards.py -v        # one line per test
"""

from __future__ import annotations

import os
import unittest

import rewards


class TestAnswerExtraction(unittest.TestCase):
    def test_gsm8k_marker_wins(self):
        self.assertEqual(rewards.extract_final_answer("working 3 + 4\n#### 7"), "7")

    def test_boxed_marker(self):
        self.assertEqual(rewards.extract_final_answer("so \\boxed{42} is it"), "42")

    def test_answer_label(self):
        self.assertEqual(rewards.extract_final_answer("Answer: -12"), "-12")

    def test_thousands_separator_is_stripped(self):
        self.assertEqual(rewards.extract_final_answer("Answer: 1,250"), "1250")

    def test_currency_prefix_is_ignored(self):
        self.assertEqual(rewards.extract_final_answer("Answer: $18"), "18")

    def test_last_marked_answer_wins_over_earlier_ones(self):
        self.assertEqual(rewards.extract_final_answer("Answer: 3\nno wait\nAnswer: 5"), "5")

    def test_fallback_takes_the_last_number(self):
        self.assertEqual(rewards.extract_final_answer("6 times 4 is 24 minus 9 is 15"), "15")

    def test_fallback_can_be_disabled(self):
        self.assertIsNone(rewards.extract_final_answer("no marker here, just 15", fallback_to_last_number=False))

    def test_no_number_at_all(self):
        self.assertIsNone(rewards.extract_final_answer("I do not know"))


class TestNumericReward(unittest.TestCase):
    def setUp(self):
        self.reward = rewards.numeric_reward()

    def test_correct_answer_scores_one(self):
        self.assertEqual(self.reward(completions=["Answer: 15"], answer=["15"]), [1.0])

    def test_decimal_and_integer_forms_agree(self):
        self.assertEqual(self.reward(completions=["Answer: 15.0"], answer=["15"]), [1.0])

    def test_wrong_answer_scores_zero(self):
        self.assertEqual(self.reward(completions=["Answer: 16"], answer=["15"]), [0.0])

    def test_missing_answer_scores_zero(self):
        self.assertEqual(self.reward(completions=["I would rather not"], answer=["15"]), [0.0])

    def test_tolerance_is_absolute_and_respected(self):
        loose = rewards.numeric_reward(tolerance=0.01)
        self.assertEqual(loose(completions=["Answer: 15.005"], answer=["15"]), [1.0])
        self.assertEqual(loose(completions=["Answer: 15.5"], answer=["15"]), [0.0])

    def test_batch_is_scored_elementwise(self):
        got = self.reward(completions=["Answer: 1", "Answer: 2", "Answer: 4"], answer=["1", "3", "4"])
        self.assertEqual(got, [1.0, 0.0, 1.0])


class TestExactMatchReward(unittest.TestCase):
    def test_string_equality_is_strict(self):
        self.assertEqual(rewards.exact_match_reward(completions=["Answer: 15"], answer=["15"]), [1.0])
        self.assertEqual(rewards.exact_match_reward(completions=["Answer: 15.0"], answer=["15"]), [0.0])


class TestFormatReward(unittest.TestCase):
    def setUp(self):
        self.reward = rewards.format_reward()

    def test_the_requested_shape_scores(self):
        good = "<think>six fours are twenty-four</think>\nAnswer: 15"
        self.assertEqual(self.reward(completions=[good]), [0.2])

    def test_trailing_commentary_does_not_score(self):
        bad = "<think>working</think>\nAnswer: 15\nHope that helps!"
        self.assertEqual(self.reward(completions=[bad]), [0.0])

    def test_missing_working_block_does_not_score(self):
        self.assertEqual(self.reward(completions=["Answer: 15"]), [0.0])

    def test_a_wrong_answer_in_the_right_shape_still_scores(self):
        # This is the point of the reward-hacking example: on its own, the format
        # reward pays for presentation and knows nothing about correctness.
        self.assertEqual(self.reward(completions=["<think>x</think>\nAnswer: 999"]), [0.2])


class TestLengthReward(unittest.TestCase):
    def setUp(self):
        self.reward = rewards.length_reward(target_tokens=100, penalty=0.2)

    def test_short_completions_are_not_penalised(self):
        self.assertEqual(self.reward(completions=["a b c"]), [0.0])

    def test_the_penalty_saturates(self):
        very_long = " ".join(["word"] * 1000)
        self.assertAlmostEqual(self.reward(completions=[very_long])[0], -0.2)

    def test_the_penalty_is_monotone_in_length(self):
        short = self.reward(completions=[" ".join(["word"] * 100)])[0]
        longer = self.reward(completions=[" ".join(["word"] * 150)])[0]
        self.assertLessEqual(longer, short)


class TestCombinedReward(unittest.TestCase):
    def setUp(self):
        self.functions, self.weights = rewards.build_reward_functions(
            kind="maths", target_tokens=100, fallback_to_last_number=True
        )

    def test_correct_and_well_formatted_beats_correct_alone(self):
        formatted = rewards.score(self.functions, self.weights,
                                  ["<think>w</think>\nAnswer: 15"], answer=["15"])[0]
        plain = rewards.score(self.functions, self.weights, ["the answer is 15"], answer=["15"])[0]
        self.assertGreater(formatted, plain)

    def test_correctness_outweighs_format(self):
        right_ugly = rewards.score(self.functions, self.weights, ["the answer is 15"], answer=["15"])[0]
        wrong_pretty = rewards.score(self.functions, self.weights,
                                     ["<think>w</think>\nAnswer: 24"], answer=["15"])[0]
        self.assertGreater(right_ugly, wrong_pretty)

    def test_the_demo_set_rewards_good_answers_more_on_average(self):
        report = rewards.evaluate_reward(self.functions, self.weights, rewards.DEMO_CASES)
        self.assertGreater(report["mean_high"], report["mean_low"])

    def test_the_demo_set_still_lets_two_hacks_through(self):
        # The whole point of the worked examples: this reward is better on average
        # and it is not safe. Two bad completions score as well as a good one, and
        # the test says so rather than the page hoping the reader noticed.
        report = rewards.evaluate_reward(self.functions, self.weights, rewards.DEMO_CASES)
        self.assertEqual(len(report["overlapping_low_cases"]), 2)

    def test_turning_off_the_last_number_fallback_closes_one_of_them(self):
        functions, weights = rewards.build_reward_functions(
            kind="maths", target_tokens=100, fallback_to_last_number=False
        )
        spray = "Maybe 6, or 4, or 24, or 9, or 13, or 14, or 15"
        self.assertEqual(rewards.score(functions, weights, [spray], answer=["15"])[0], 0.0)

    def test_the_same_change_also_zeroes_a_correct_but_unmarked_answer(self):
        # The price of closing that hole, and the reason the format reward exists:
        # with no fallback, an answer the extractor cannot find is an answer that
        # did not happen, however right it was.
        functions, weights = rewards.build_reward_functions(
            kind="maths", target_tokens=100, fallback_to_last_number=False
        )
        marked = "Six times four is twenty-four, minus nine. Answer: 15"
        self.assertEqual(rewards.score(functions, weights, [marked], answer=["15"])[0], 1.0)
        unmarked = "Six times four is twenty-four, minus nine leaves 15."
        self.assertEqual(rewards.score(functions, weights, [unmarked], answer=["15"])[0], 0.0)


@unittest.skipUnless(os.name == "posix", "the sandbox uses POSIX resource limits")
class TestUnitTestReward(unittest.TestCase):
    def setUp(self):
        self.reward = rewards.unit_test_reward(timeout_s=10)

    def test_correct_code_passes(self):
        completion = "```python\ndef add(a, b):\n    return a + b\n```"
        self.assertEqual(self.reward(completions=[completion], tests=["assert add(2, 2) == 4\n"]), [1.0])

    def test_incorrect_code_fails(self):
        completion = "```python\ndef add(a, b):\n    return a - b\n```"
        self.assertEqual(self.reward(completions=[completion], tests=["assert add(2, 2) == 4\n"]), [0.0])

    def test_an_infinite_loop_is_stopped(self):
        completion = "```python\ndef add(a, b):\n    while True:\n        pass\n```"
        self.assertEqual(self.reward(completions=[completion], tests=["assert add(2, 2) == 4\n"]), [0.0])

    def test_printing_the_answer_does_not_pass(self):
        completion = "```python\nprint('ok')\ndef add(a, b):\n    return None\n```"
        self.assertEqual(self.reward(completions=[completion], tests=["assert add(2, 2) == 4\n"]), [0.0])

    def test_code_outside_a_fence_is_still_run(self):
        self.assertEqual(self.reward(completions=["def add(a, b):\n    return a + b\n"],
                                     tests=["assert add(2, 2) == 4\n"]), [1.0])

    def test_partial_credit_scores_the_fraction_of_blocks(self):
        partial = rewards.unit_test_reward(timeout_s=10, partial_credit=True)
        completion = "```python\ndef add(a, b):\n    return abs(a) + abs(b)\n```"
        tests = ["assert add(2, 2) == 4\n\nassert add(-1, 1) == 0\n"]
        self.assertAlmostEqual(partial(completions=[completion], tests=tests)[0], 0.5)


class TestMessageShapedCompletions(unittest.TestCase):
    """TRL hands conversational datasets back as message lists, not strings."""

    def test_a_message_list_is_flattened(self):
        completion = [{"role": "assistant", "content": "Answer: 15"}]
        self.assertEqual(rewards.numeric_reward()(completions=[completion], answer=["15"]), [1.0])


if __name__ == "__main__":
    unittest.main()