import random
import unittest

from gradescope_utils.autograder_utils.decorators import weight, number

try:
    from programming.open_evolve import (
        ProgramDatabase, build_prompt, parse_program, evolve,
        temperature_schedule, island_evolve,
    )
except ImportError:
    from open_evolve import (
        ProgramDatabase, build_prompt, parse_program, evolve,
        temperature_schedule, island_evolve,
    )


class TestProgramDatabase(unittest.TestCase):

    @weight(2)
    @number("2.1")
    def test_add_and_len(self):
        """Test that add() appends programs and __len__ returns the correct count."""
        db = ProgramDatabase()
        self.assertEqual(len(db), 0)
        db.add("prog1", 1.0)
        self.assertEqual(len(db), 1)
        db.add("prog2", 2.0)
        self.assertEqual(len(db), 2)
        db.add("prog3", 0.5)
        self.assertEqual(len(db), 3)

    @weight(2)
    @number("2.2")
    def test_best_returns_highest_score(self):
        """Test that best() returns the (program, score) with the highest score."""
        db = ProgramDatabase()
        db.add("prog_low", 0.5)
        db.add("prog_high", 9.0)
        db.add("prog_mid", 3.0)
        best_prog, best_score = db.best()
        self.assertEqual(best_prog, "prog_high")
        self.assertAlmostEqual(best_score, 9.0)

    @weight(2)
    @number("2.3")
    def test_sample_returns_correct_count(self):
        """Test that sample(n) returns exactly n (program, score) tuples."""
        db = ProgramDatabase()
        db.add("a", 1.0)
        db.add("b", 2.0)
        db.add("c", 3.0)
        result = db.sample(2)
        self.assertEqual(len(result), 2)
        for prog, score in result:
            self.assertIsInstance(prog, str)
            self.assertIsInstance(score, float)

    @weight(2)
    @number("2.4")
    def test_sample_higher_score_sampled_more(self):
        """Test that higher-scored programs are sampled more often."""
        random.seed(0)
        db = ProgramDatabase()
        db.add("low", 0.01)
        db.add("high", 100.0)
        counts = {"low": 0, "high": 0}
        for _ in range(500):
            prog, _ = db.sample(1)[0]
            counts[prog] += 1
        self.assertGreater(counts["high"], counts["low"],
                           msg="Higher-scored program should be sampled more often.")

    @weight(2)
    @number("2.5")
    def test_sample_empty_database(self):
        """Test that sample() on an empty database returns an empty list."""
        db = ProgramDatabase()
        result = db.sample(1)
        self.assertEqual(result, [])


class TestBuildPrompt(unittest.TestCase):

    @weight(2)
    @number("2.6")
    def test_prompt_contains_task_description(self):
        """Test that build_prompt includes the task description."""
        task = "maximize the sum of squares"
        prompt = build_prompt(task, [("def f(): return 1", 0.5)])
        self.assertIn(task, prompt)

    @weight(2)
    @number("2.7")
    def test_prompt_contains_programs_and_scores(self):
        """Test that build_prompt includes program code and formatted scores."""
        prog = "def f(): return 42"
        prompt = build_prompt("some task", [(prog, 3.14)])
        self.assertIn(prog, prompt)
        self.assertIn("3.14", prompt)

    @weight(2)
    @number("2.8")
    def test_prompt_multiple_programs(self):
        """Test that build_prompt includes all provided inspiration programs."""
        programs = [("prog_a", 1.0), ("prog_b", 2.0), ("prog_c", 3.0)]
        prompt = build_prompt("some task", programs)
        for prog, _ in programs:
            self.assertIn(prog, prompt)


class TestParseProgram(unittest.TestCase):

    @weight(2)
    @number("2.9")
    def test_parse_code_block(self):
        """Test that parse_program extracts code from ```python ... ``` markers."""
        response = "Here is the improved version:\n```python\ndef f(): return 42\n```\nDone."
        result = parse_program(response)
        self.assertIn("def f(): return 42", result)

    @weight(2)
    @number("2.10")
    def test_parse_fallback_no_block(self):
        """Test that parse_program returns stripped response when no code block found."""
        response = "  def f(): return 1  "
        result = parse_program(response)
        self.assertEqual(result, "def f(): return 1")


class TestEvolve(unittest.TestCase):

    def _make_mock_llm(self, responses):
        """Returns a mock LLM that steps through a list of responses."""
        state = [0]
        def mock_llm(prompt):
            idx = min(state[0], len(responses) - 1)
            state[0] += 1
            return responses[idx]
        return mock_llm

    @weight(8)
    @number("2.11")
    def test_evolve_returns_best(self):
        """Test that evolve() returns the highest-scoring program found."""
        responses = [
            "```python\nversion_1\n```",
            "```python\nversion_2\n```",
            "```python\nversion_3\n```",
        ]
        score_map = {"seed": 0.0, "version_1": 0.3, "version_2": 0.7, "version_3": 0.5}
        evaluator = lambda p: score_map.get(p, 0.0)
        llm = self._make_mock_llm(responses)

        best_prog, best_score = evolve("task", "seed", evaluator, llm, n_iterations=3)
        self.assertAlmostEqual(best_score, 0.7)
        self.assertEqual(best_prog, "version_2")

    @weight(4)
    @number("2.12")
    def test_evolve_seed_is_evaluated(self):
        """Test that evolve() evaluates and stores the seed program."""
        evaluated = []

        def tracking_evaluator(p):
            evaluated.append(p)
            return 1.0

        llm = self._make_mock_llm(["```python\nnew_prog\n```"])
        evolve("task", "my_seed", tracking_evaluator, llm, n_iterations=1)
        self.assertIn("my_seed", evaluated)

    @weight(4)
    @number("2.13")
    def test_evolve_handles_evaluator_exception(self):
        """Test that evolve() catches evaluator exceptions and assigns score 0.0."""
        def bad_evaluator(p):
            if p == "bad_program":
                raise ValueError("bad program!")
            return 5.0

        llm = self._make_mock_llm(["```python\nbad_program\n```"])
        best_prog, best_score = evolve("task", "good_seed", bad_evaluator, llm, n_iterations=1)
        # good_seed scored 5.0; bad_program raised and should score 0.0
        self.assertAlmostEqual(best_score, 5.0)
        self.assertEqual(best_prog, "good_seed")


class TestTopK(unittest.TestCase):

    @weight(2)
    @number("2.14")
    def test_top_k_sorted_descending(self):
        """top_k returns entries sorted by score, highest first."""
        db = ProgramDatabase()
        db.add("a", 1.0)
        db.add("b", 5.0)
        db.add("c", 3.0)
        result = db.top_k(2)
        self.assertEqual(len(result), 2)
        self.assertEqual(result[0][0], "b")
        self.assertAlmostEqual(result[0][1], 5.0)
        self.assertEqual(result[1][0], "c")

    @weight(2)
    @number("2.15")
    def test_top_k_larger_than_db(self):
        """top_k returns all entries when k >= database size."""
        db = ProgramDatabase()
        db.add("x", 2.0)
        db.add("y", 4.0)
        result = db.top_k(10)
        self.assertEqual(len(result), 2)
        self.assertGreaterEqual(result[0][1], result[1][1])


class TestTemperatureSchedule(unittest.TestCase):

    @weight(3)
    @number("2.16")
    def test_start_and_end_values(self):
        """Temperature equals start_temp at iteration 0 and end_temp at the last."""
        self.assertAlmostEqual(temperature_schedule(0, 10, start_temp=2.0, end_temp=0.1), 2.0)
        self.assertAlmostEqual(temperature_schedule(9, 10, start_temp=2.0, end_temp=0.1), 0.1)

    @weight(3)
    @number("2.17")
    def test_monotone_decreasing(self):
        """Temperature decreases monotonically across iterations."""
        temps = [temperature_schedule(i, 20) for i in range(20)]
        for i in range(len(temps) - 1):
            self.assertGreaterEqual(temps[i], temps[i + 1])


class TestIslandEvolve(unittest.TestCase):

    def _make_mock_llm(self, responses):
        state = [0]
        def mock_llm(prompt):
            idx = min(state[0], len(responses) - 1)
            state[0] += 1
            return responses[idx]
        return mock_llm

    @weight(12)
    @number("2.18")
    def test_island_evolve_returns_best_across_islands(self):
        """island_evolve returns the single best result found across all islands."""
        # Island 0 will find score 0.9, island 1 will find 0.5, island 2 will find 0.7
        call_count = [0]
        island_peaks = {0: ("island0_best", 0.9), 1: ("island1_best", 0.5), 2: ("island2_best", 0.7)}
        island_idx = [0]

        def mock_evaluator(p):
            if p == "seed":
                return 0.0
            island = island_idx[0] % 3
            name, score = island_peaks[island]
            island_idx[0] += 1
            return score if p == name else 0.1

        responses_per_island = [
            f"```python\nisland{i}_best\n```" for i in range(3)
        ]
        resp_state = [0]
        def mock_llm(prompt):
            r = responses_per_island[resp_state[0] % 3]
            resp_state[0] += 1
            return r

        best_prog, best_score = island_evolve(
            "task", "seed", mock_evaluator, mock_llm,
            n_islands=3, n_iterations_per_island=1,
        )
        self.assertAlmostEqual(best_score, 0.9)
        self.assertEqual(best_prog, "island0_best")

    @weight(6)
    @number("2.19")
    def test_island_evolve_single_island_matches_evolve(self):
        """With n_islands=1, island_evolve result matches a single evolve() call."""
        responses = ["```python\nbetter_prog\n```"]
        score_map = {"seed": 0.0, "better_prog": 1.0}
        evaluator = lambda p: score_map.get(p, 0.0)

        state = [0]
        def mock_llm(prompt):
            r = responses[min(state[0], len(responses) - 1)]
            state[0] += 1
            return r

        best_prog, best_score = island_evolve(
            "task", "seed", evaluator, mock_llm,
            n_islands=1, n_iterations_per_island=1,
        )
        self.assertAlmostEqual(best_score, 1.0)
        self.assertEqual(best_prog, "better_prog")

    @weight(6)
    @number("2.20")
    def test_island_evolve_fresh_db_per_island(self):
        """Each island starts from a fresh database (seed re-evaluated each time)."""
        eval_calls = []
        def tracking_evaluator(p):
            eval_calls.append(p)
            return 0.0

        mock_llm = lambda p: "```python\nsome_prog\n```"
        island_evolve("task", "seed", tracking_evaluator, mock_llm,
                      n_islands=3, n_iterations_per_island=1)

        # seed must be evaluated once per island
        seed_count = eval_calls.count("seed")
        self.assertEqual(seed_count, 3)


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