import numpy as np
import matplotlib.pyplot as plt
import random
import math

# Set random seed for reproducibility
random.seed(42)
np.random.seed(42)

def objective_function(x: float) -> float:
    """
    The objective function to maximize.
    Formula: f(x) = -x^2 + 4x + 10

    Args:
        x (float): The input value for the function.

    Returns:
        float: The calculated value of the function.
    """
    return -x**2 + 4*x + 10

# Define the search space boundaries
LOWER_BOUND = 0
UPPER_BOUND = 5

print(f"Objective function for x=2: {objective_function(2)}")
print(f"Objective function for x=0: {objective_function(0)}")
print(f"Objective function for x=5: {objective_function(5)}")

def initialize_population(population_size: int, lower_bound: float, upper_bound: float) -> list:
    """
    Initializes a population of random solutions within the given bounds.

    Args:
        population_size (int): The number of individuals in the population.
        lower_bound (float): The lower limit of the search space.
        upper_bound (float): The upper limit of the search space.

    Returns:
        list: A list of floats, where each float is a random solution.
    """
    return [random.uniform(lower_bound, upper_bound) for _ in range(population_size)]

# Example usage
initial_pop = initialize_population(population_size=10, lower_bound=LOWER_BOUND, upper_bound=UPPER_BOUND)
print(f"Initial population (first 5 individuals): {initial_pop[:5]}")

def calculate_fitness(population: list, objective_func: callable) -> list:
    """
    Calculates the fitness for each individual in the population.
    For maximization problems, fitness is often directly the objective function value.

    Args:
        population (list): A list of individual solutions.
        objective_func (callable): The function to evaluate the fitness of each individual.

    Returns:
        list: A list of fitness scores corresponding to each individual.
    """
    return [objective_func(individual) for individual in population]

# Example usage
fitness_scores = calculate_fitness(initial_pop, objective_function)
print(f"Fitness scores for initial population (first 5): {fitness_scores[:5]}")

def selection(population: list, fitness_scores: list, num_parents: int) -> list:
    """
    Performs tournament selection to choose parents for the next generation.
    Formula: Selects `num_parents` individuals by repeatedly picking `tournament_size`
             random individuals and choosing the best among them.

    Args:
        population (list): The current population of solutions.
        fitness_scores (list): The fitness scores for each individual in the population.
        num_parents (int): The number of parents to select.

    Returns:
        list: A list of selected parent individuals.
    """
    parents = []
    tournament_size = 3 # Hyperparameter for tournament selection

    for _ in range(num_parents):
        # Randomly select individuals for the tournament
        tournament_competitors_indices = random.sample(range(len(population)), tournament_size)

        # Find the best individual in the tournament
        best_competitor_index = tournament_competitors_indices[0]
        for i in tournament_competitors_indices:
            if fitness_scores[i] > fitness_scores[best_competitor_index]:
                best_competitor_index = i
        parents.append(population[best_competitor_index])
    return parents

# Example usage
selected_parents = selection(initial_pop, fitness_scores, num_parents=initial_pop.size // 2 if hasattr(initial_pop, 'size') else len(initial_pop) // 2)
print(f"Selected parents (first 3): {selected_parents[:3]}")

def crossover(parents: list, offspring_size: int, lower_bound: float, upper_bound: float) -> list:
    """
    Performs single-point crossover to create offspring from selected parents.
    Formula: If `crossover_point` is chosen between parent1 and parent2, then
             offspring1 = parent1[:crossover_point] + parent2[crossover_point:]
             offspring2 = parent2[:crossover_point] + parent1[crossover_point:]
             (For real-valued genes, it's typically a weighted average).
             Here, we use a blend crossover for real-valued genes.

    Args:
        parents (list): A list of parent individuals.
        offspring_size (int): The desired number of offspring.
        lower_bound (float): The lower limit of the search space.
        upper_bound (float): The upper limit of the search space.

    Returns:
        list: A list of new offspring individuals.
    """
    offspring = []
    num_parents = len(parents)

    for _ in range(offspring_size):
        # Randomly select two parents
        parent1 = random.choice(parents)
        parent2 = random.choice(parents)

        # Blend Crossover for real-valued genes
        alpha = random.uniform(0, 1) # Blending factor
        child = alpha * parent1 + (1 - alpha) * parent2

        # Ensure child stays within bounds
        child = max(lower_bound, min(child, upper_bound))
        offspring.append(child)

    return offspring

# Example usage (assuming selected_parents has enough members)
num_offspring = len(initial_pop) - len(selected_parents) # Fill remaining slots
if num_offspring < 0: # Handle cases where more parents than population size
    num_offspring = len(initial_pop)

created_offspring = crossover(selected_parents, num_offspring, LOWER_BOUND, UPPER_BOUND)
print(f"Created offspring (first 3): {created_offspring[:3]}")

def mutate(offspring: list, mutation_rate: float, lower_bound: float, upper_bound: float, mutation_strength: float = 0.1) -> list:
    """
    Applies random mutation to the offspring.
    Formula: If a gene is selected for mutation, its value is perturbed by adding a small random value.
             Here, `individual = individual + random_normal_deviation * mutation_strength`.

    Args:
        offspring (list): A list of offspring individuals.
        mutation_rate (float): The probability of an individual undergoing mutation.
        lower_bound (float): The lower limit of the search space.
        upper_bound (float): The upper limit of the search space.
        mutation_strength (float): Controls the magnitude of the mutation.

    Returns:
        list: The offspring population after potential mutations.
    """
    mutated_offspring = []
    for individual in offspring:
        if random.random() < mutation_rate:
            # Apply a small perturbation (e.g., Gaussian mutation)
            perturbation = np.random.normal(0, mutation_strength * (upper_bound - lower_bound))
            individual += perturbation

            # Ensure individual stays within bounds after mutation
            individual = max(lower_bound, min(individual, upper_bound))
        mutated_offspring.append(individual)
    return mutated_offspring

# Example usage
mutated_offspring_example = mutate(created_offspring, mutation_rate=0.1,
                                   lower_bound=LOWER_BOUND, upper_bound=UPPER_BOUND)
print(f"Mutated offspring (first 3): {mutated_offspring_example[:3]}")

def genetic_algorithm(objective_func: callable,
                      population_size: int,
                      lower_bound: float,
                      upper_bound: float,
                      generations: int,
                      mutation_rate: float,
                      mutation_strength: float) -> tuple:
    """
    Executes the main genetic algorithm optimization process.

    Args:
        objective_func (callable): The function to maximize.
        population_size (int): The number of individuals in each generation.
        lower_bound (float): The lower limit of the search space.
        upper_bound (float): The upper limit of the search space.
        generations (int): The total number of generations to run the GA.
        mutation_rate (float): The probability of an individual undergoing mutation.
        mutation_strength (float): Controls the magnitude of the mutation.

    Returns:
        tuple: A tuple containing:
            - list: The final population of solutions.
            - list: The best individual found across all generations.
            - list: The fitness of the best individual across all generations.
            - list: A list of the best fitness score at each generation.
    """
    population = initialize_population(population_size, lower_bound, upper_bound)
    best_individual_overall = None
    best_fitness_overall = -np.inf # For maximization
    best_fitness_per_generation = []

    for gen in range(generations):
        fitness_scores = calculate_fitness(population, objective_func)

        # Find the best individual in the current generation
        current_best_idx = np.argmax(fitness_scores)
        current_best_individual = population[current_best_idx]
        current_best_fitness = fitness_scores[current_best_idx]

        # Update overall best if current generation's best is better
        if current_best_fitness > best_fitness_overall:
            best_fitness_overall = current_best_fitness
            best_individual_overall = current_best_individual

        best_fitness_per_generation.append(best_fitness_overall)

        # Selection: Choose parents
        num_parents = population_size // 2 # Arbitrarily select half as parents
        parents = selection(population, fitness_scores, num_parents)

        # Crossover: Create offspring to fill the new population
        offspring_size = population_size - len(parents)
        offspring = crossover(parents, offspring_size, lower_bound, upper_bound)

        # Mutation: Introduce diversity in offspring
        mutated_offspring = mutate(offspring, mutation_rate, lower_bound, upper_bound, mutation_strength)

        # Combine parents and mutated offspring to form the new population
        # A common strategy is to keep the elite (best individuals) and replace the rest
        # For simplicity, we just combine parents and offspring here.
        # In more advanced GAs, elitism ensures the best solution isn't lost.
        new_population = parents + mutated_offspring

        # If for some reason the new population size doesn't match, adjust
        if len(new_population) < population_size:
            # Fill remaining slots with random individuals or duplicates
            new_population.extend(initialize_population(population_size - len(new_population), lower_bound, upper_bound))
        elif len(new_population) > population_size:
            new_population = new_population[:population_size]

        population = new_population

        # Optional: print progress every few generations
        if gen % (generations // 10 if generations >= 10 else 1) == 0 or gen == generations -1:
            print(f"Generation {gen+1}/{generations}: Best Fitness = {best_fitness_overall:.4f}, Best Individual = {best_individual_overall:.4f}")

    return population, best_individual_overall, best_fitness_overall, best_fitness_per_generation

# GA Parameters
POPULATION_SIZE = 50
GENERATIONS = 100
MUTATION_RATE = 0.1
MUTATION_STRENGTH = 0.05

print("Running Genetic Algorithm...")
final_population, best_solution, best_fitness, fitness_history = genetic_algorithm(
    objective_function, POPULATION_SIZE, LOWER_BOUND, UPPER_BOUND,
    GENERATIONS, MUTATION_RATE, MUTATION_STRENGTH
)

print("\nGenetic Algorithm Finished!")
print(f"Optimized solution (x): {best_solution:.4f}")
print(f"Maximum fitness (f(x)): {best_fitness:.4f}")

# Compare with actual optimum
actual_optimum_x = 2
actual_optimum_fitness = objective_function(actual_optimum_x)
print(f"\nActual optimum (x=2): {actual_optimum_fitness:.4f}")

fig, axes = plt.subplots(1, 2, figsize=(15, 6))

# --- First Plot: Fitness vs. Generations (Cumulative Best Fitness) ---
axes[0].plot(range(1, GENERATIONS + 1), fitness_history, color='skyblue', linewidth=2)
axes[0].set_title('Evolution of Best Fitness Over Generations')
axes[0].set_xlabel('Generation')
axes[0].set_ylabel('Best Fitness')
axes[0].grid(True, linestyle='--', alpha=0.7)
axes[0].axhline(y=actual_optimum_fitness, color='red', linestyle=':', label=f'True Optimum ({actual_optimum_fitness:.2f})')
axes[0].legend()

# --- Second Plot: Distribution of Solutions in Final Population vs. Objective Function ---
# Generate x values for plotting the objective function
x_vals = np.linspace(LOWER_BOUND, UPPER_BOUND, 100)
y_vals = [objective_function(x) for x in x_vals]

axes[1].plot(x_vals, y_vals, label='Objective Function $f(x) = -x^2 + 4x + 10$', color='blue', alpha=0.7)
axes[1].scatter(final_population, [objective_function(x) for x in final_population],
            color='green', alpha=0.6, s=50, label='Final Population Individuals')
axes[1].scatter(best_solution, best_fitness, color='red', s=200, marker='*',
            label=f'Best Solution ({best_solution:.2f}, {best_fitness:.2f})', zorder=5)
axes[1].set_title('Final Population Distribution and Objective Function')
axes[1].set_xlabel('x value')
axes[1].set_ylabel('f(x) (Fitness)')
axes[1].grid(True, linestyle='--', alpha=0.7)
axes[1].legend()

plt.tight_layout()
plt.show()