import re
import math
import random

random.seed(42)

#####################
# Program Database  #
#####################

class ProgramDatabase:
    """
    A database of programs and their evaluation scores.

    Stores (program_str, score) pairs and supports fitness-proportional
    sampling for the evolutionary loop.
    """

    def __init__(self):
        self._programs: list[str] = []
        self._scores: list[float] = []

    def add(self, program: str, score: float) -> None:
        """Add a program and its score to the database."""
        # Fill in

    def sample(self, n: int = 1, temperature: float = 1.0) -> list[tuple[str, float]]:
        """
        Sample n (program, score) pairs with probability proportional to score.

        Uses softmax weighting with a temperature parameter:
          weight_i = exp(score_i / temperature) / sum_j(exp(score_j / temperature))

        Samples WITH replacement. If the database is empty, returns [].

        Args:
            n: Number of (program, score) pairs to sample.
            temperature: Softmax temperature (higher = more uniform sampling).

        Returns:
            List of (program, score) tuples.
        """
        # Fill in

    def best(self) -> tuple[str, float]:
        """Return the (program, score) pair with the highest score."""
        # Fill in

    def top_k(self, k: int) -> list[tuple[str, float]]:
        """
        Return the k highest-scoring (program, score) pairs, sorted descending.

        If k >= len(database), return all entries sorted by score (highest first).

        Args:
            k: Number of top entries to return.

        Returns:
            List of (program, score) tuples sorted by score descending.
        """
        # Fill in

    def __len__(self) -> int:
        return len(self._programs)


##########################
# Prompt Construction    #
##########################

def build_prompt(
    task_description: str,
    inspiration_programs: list[tuple[str, float]],
) -> str:
    """
    Build a prompt for the LLM that asks it to improve upon existing programs.

    The prompt must include:
      1. An introductory instruction line.
      2. The task description.
      3. Each inspiration program with its score, formatted as a Python code block.
      4. A closing instruction asking the LLM to output an improved program
         inside a ```python ... ``` block.

    Args:
        task_description: A string describing the optimization task.
        inspiration_programs: A list of (program_str, score) tuples from the database.

    Returns:
        A formatted prompt string.
    """
    # Fill in


#######################
# Program Parsing     #
#######################

def parse_program(response: str) -> str:
    """
    Extract a Python program from an LLM response.

    Looks for content between ```python and ``` markers (case-insensitive).
    If no code block is found, returns the full response stripped of whitespace.

    Args:
        response: Raw string returned by the LLM.

    Returns:
        Extracted program string.
    """
    # Fill in


##############################
# Main Evolutionary Loop     #
##############################

def evolve(
    task_description: str,
    seed_program: str,
    evaluator: callable,
    llm: callable,
    n_iterations: int = 10,
    n_inspirations: int = 2,
    db: ProgramDatabase | None = None,
) -> tuple[str, float]:
    """
    Run the OpenEvolve evolutionary loop.

    Algorithm:
      1. Create a new ProgramDatabase (or use the provided one if given).
      2. Evaluate the seed program and add it to the database.
      3. For each iteration:
           a. Sample n_inspirations (program, score) pairs from the database.
           b. Build a prompt using the task description and sampled programs.
           c. Query the LLM with the prompt to get a response string.
           d. Parse the program from the LLM response.
           e. Evaluate the new program.
              - If the evaluator raises an exception, use a score of 0.0.
           f. Add the new program and score to the database.
      4. Return the best (program, score) found.

    Args:
        task_description: Description of the optimization task.
        seed_program: Initial program string to seed the database.
        evaluator: Callable that takes a program string and returns a float score.
        llm: Callable that takes a prompt string and returns a response string.
        n_iterations: Number of evolutionary iterations to run.
        n_inspirations: Number of programs to sample per iteration.
        db: Optional pre-initialized ProgramDatabase. Creates a new one if None.

    Returns:
        Tuple of (best_program_string, best_score).
    """
    # Fill in


################################
# Temperature Scheduling      #
################################

def temperature_schedule(
    iteration: int,
    n_iterations: int,
    start_temp: float = 2.0,
    end_temp: float = 0.1,
) -> float:
    """
    Return the softmax temperature at a given iteration using linear annealing.

    Linearly interpolates from start_temp down to end_temp:

        temp = start_temp - (start_temp - end_temp) * iteration / (n_iterations - 1)

    At iteration=0 the temperature is start_temp (high, encourages exploration).
    At iteration=n_iterations-1 it is end_temp (low, encourages exploitation).

    Args:
        iteration:    Current iteration index (0-based).
        n_iterations: Total number of iterations.
        start_temp:   Temperature at the first iteration.
        end_temp:     Temperature at the last iteration.

    Returns:
        Temperature (float) for the given iteration.

    Note: behaviour is undefined for n_iterations <= 1 or iteration out of range.
    """
    # Fill in


##############################
# Island Evolution           #
##############################

def island_evolve(
    task_description: str,
    seed_program: str,
    evaluator: callable,
    llm: callable,
    n_islands: int = 3,
    n_iterations_per_island: int = 5,
    n_inspirations: int = 2,
) -> tuple[str, float]:
    """
    Run OpenEvolve independently on multiple populations (islands) and return
    the best result found across all islands.

    For each island (0 to n_islands-1):
      1. Create a fresh ProgramDatabase.
      2. Call evolve() with n_iterations_per_island iterations, using the
         provided seed_program as the starting point.

    After all islands have finished, return the (program, score) pair with
    the highest score across all islands.

    Args:
        task_description:       Description of the optimization task.
        seed_program:           Seed program shared by all islands.
        evaluator:              Callable that scores a program string.
        llm:                    Callable that takes a prompt and returns a response.
        n_islands:              Number of independent evolution runs.
        n_iterations_per_island: Iterations passed to each evolve() call.
        n_inspirations:         n_inspirations passed to each evolve() call.

    Returns:
        Tuple of (best_program_string, best_score) across all islands.
    """
    # Fill in


##################
# Demo           #
##################

if __name__ == "__main__":
    # Simple demo: evolve a function f() that returns sum(i**2 for i in range(5)) = 30.
    # A mock LLM cycles through increasingly better implementations.

    TARGET = sum(i**2 for i in range(5))  # 30

    def demo_evaluator(program: str) -> float:
        """Score a program by executing it and measuring distance from TARGET."""
        env: dict = {}
        exec(program, env)
        result = env["f"]()
        return -abs(result - TARGET)  # higher (closer to 0) is better

    llm_responses = [
        "```python\ndef f():\n    return 0\n```",
        "```python\ndef f():\n    return sum(range(5))\n```",
        "```python\ndef f():\n    return sum(i**2 for i in range(5))\n```",
    ]
    call_idx = [0]

    def demo_llm(prompt: str) -> str:
        idx = min(call_idx[0], len(llm_responses) - 1)
        call_idx[0] += 1
        return llm_responses[idx]

    seed = "def f():\n    return -999"
    task = "Write a Python function f() that returns the sum of squares of 0 through 4."

    best_prog, best_score = evolve(task, seed, demo_evaluator, demo_llm, n_iterations=5)
    print(f"Best score: {best_score}")
    print(f"Best program:\n{best_prog}")
