"""
Genetic Operators Module
========================

Implements standard genetic algorithm operators:
- Mutation: Introduces random changes for exploration
- Crossover: Combines features of good solutions for exploitation

These operators are fundamental to evolutionary programming methods
and enable both exploration and exploitation in optimization.
"""

import numpy as np
from typing import Tuple, List, Optional


class GeneticOperators:
    """Collection of genetic operators for evolutionary algorithms"""
    
    def __init__(self, mutation_rate: float = 0.1, crossover_rate: float = 0.8):
        """
        Initialize genetic operators
        
        Args:
            mutation_rate: Probability of mutation (0 to 1)
            crossover_rate: Probability of crossover (0 to 1)
        """
        self.mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate
    
    # ==================== MUTATION OPERATORS ====================
    
    def uniform_mutation(self, individual: np.ndarray, bounds: Tuple[float, float], 
                        mutation_rate: Optional[float] = None) -> np.ndarray:
        """
        Uniform mutation: Random replacement within bounds
        Best for: Exploration, escaping local optima
        
        Args:
            individual: Solution vector to mutate
            bounds: (lower_bound, upper_bound) tuple
            mutation_rate: Override default mutation rate
            
        Returns:
            Mutated individual
        """
        rate = mutation_rate if mutation_rate is not None else self.mutation_rate
        lower, upper = bounds
        lower = np.broadcast_to(np.asarray(lower, dtype=float), len(individual)).copy()
        upper = np.broadcast_to(np.asarray(upper, dtype=float), len(individual)).copy()
        
        mutant = individual.copy()
        mask = np.random.random(len(individual)) < rate
        
        # Replace selected genes with random values
        mutant[mask] = np.random.uniform(lower[mask], upper[mask])
        
        return mutant
    
    def gaussian_mutation(self, individual: np.ndarray, bounds: Tuple[float, float],
                         mutation_rate: Optional[float] = None, 
                         sigma: float = 0.1) -> np.ndarray:
        """
        Gaussian mutation: Add normally distributed noise
        Best for: Fine-tuning, local search
        
        Args:
            individual: Solution vector to mutate
            bounds: (lower_bound, upper_bound) tuple
            mutation_rate: Override default mutation rate
            sigma: Standard deviation of Gaussian noise (relative to range)
            
        Returns:
            Mutated individual
        """
        rate = mutation_rate if mutation_rate is not None else self.mutation_rate
        lower, upper = bounds
        lower = np.broadcast_to(np.asarray(lower, dtype=float), len(individual)).copy()
        upper = np.broadcast_to(np.asarray(upper, dtype=float), len(individual)).copy()
        range_size = upper - lower
        
        mutant = individual.copy()
        mask = np.random.random(len(individual)) < rate
        
        # Add Gaussian noise to selected genes
        noise = np.random.normal(0, sigma * range_size[mask])
        mutant[mask] += noise
        
        # Ensure bounds
        mutant = np.clip(mutant, lower, upper)
        
        return mutant
    
    def polynomial_mutation(self, individual: np.ndarray, bounds: Tuple[float, float],
                           mutation_rate: Optional[float] = None,
                           eta: float = 20.0) -> np.ndarray:
        """
        Polynomial mutation: Self-adaptive mutation with distribution index
        Best for: Balanced exploration-exploitation
        
        Args:
            individual: Solution vector to mutate
            bounds: (lower_bound, upper_bound) tuple
            mutation_rate: Override default mutation rate
            eta: Distribution index (higher = less spread)
            
        Returns:
            Mutated individual
        """
        rate = mutation_rate if mutation_rate is not None else self.mutation_rate
        lower, upper = bounds
        lower = np.broadcast_to(np.asarray(lower, dtype=float), len(individual)).copy()
        upper = np.broadcast_to(np.asarray(upper, dtype=float), len(individual)).copy()
        
        mutant = individual.copy()
        
        for i in range(len(individual)):
            if np.random.random() < rate:
                y = mutant[i]
                delta1 = (y - lower[i]) / (upper[i] - lower[i])
                delta2 = (upper[i] - y) / (upper[i] - lower[i])
                
                rand = np.random.random()
                mut_pow = 1.0 / (eta + 1.0)
                
                if rand < 0.5:
                    xy = 1.0 - delta1
                    val = 2.0 * rand + (1.0 - 2.0 * rand) * xy ** (eta + 1.0)
                    deltaq = val ** mut_pow - 1.0
                else:
                    xy = 1.0 - delta2
                    val = 2.0 * (1.0 - rand) + 2.0 * (rand - 0.5) * xy ** (eta + 1.0)
                    deltaq = 1.0 - val ** mut_pow
                
                y = y + deltaq * (upper[i] - lower[i])
                mutant[i] = np.clip(y, lower[i], upper[i])
        
        return mutant
    
    def boundary_mutation(self, individual: np.ndarray, bounds: Tuple[float, float],
                         mutation_rate: Optional[float] = None) -> np.ndarray:
        """
        Boundary mutation: Set selected genes to boundary values
        Best for: Exploring extreme solutions
        
        Args:
            individual: Solution vector to mutate
            bounds: (lower_bound, upper_bound) tuple
            mutation_rate: Override default mutation rate
            
        Returns:
            Mutated individual
        """
        rate = mutation_rate if mutation_rate is not None else self.mutation_rate
        lower, upper = bounds
        lower = np.broadcast_to(np.asarray(lower, dtype=float), len(individual)).copy()
        upper = np.broadcast_to(np.asarray(upper, dtype=float), len(individual)).copy()
        
        mutant = individual.copy()
        mask = np.random.random(len(individual)) < rate
        
        # Set to random boundary
        boundary_values = np.where(np.random.random(np.sum(mask)) < 0.5, lower[mask], upper[mask])
        mutant[mask] = boundary_values
        
        return mutant
    
    # ==================== CROSSOVER OPERATORS ====================
    
    def single_point_crossover(self, parent1: np.ndarray, parent2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Single-point crossover: Exchange genes at one random point
        Best for: Simple problems, maintaining building blocks
        
        Args:
            parent1: First parent solution
            parent2: Second parent solution
            
        Returns:
            Tuple of (offspring1, offspring2)
        """
        if np.random.random() > self.crossover_rate:
            return parent1.copy(), parent2.copy()
        
        point = np.random.randint(1, len(parent1))
        
        offspring1 = np.concatenate([parent1[:point], parent2[point:]])
        offspring2 = np.concatenate([parent2[:point], parent1[point:]])
        
        return offspring1, offspring2
    
    def two_point_crossover(self, parent1: np.ndarray, parent2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Two-point crossover: Exchange genes between two random points
        Best for: Preserving gene adjacency
        
        Args:
            parent1: First parent solution
            parent2: Second parent solution
            
        Returns:
            Tuple of (offspring1, offspring2)
        """
        if np.random.random() > self.crossover_rate:
            return parent1.copy(), parent2.copy()
        
        points = sorted(np.random.choice(len(parent1), 2, replace=False))
        
        offspring1 = parent1.copy()
        offspring2 = parent2.copy()
        
        offspring1[points[0]:points[1]] = parent2[points[0]:points[1]]
        offspring2[points[0]:points[1]] = parent1[points[0]:points[1]]
        
        return offspring1, offspring2
    
    def uniform_crossover(self, parent1: np.ndarray, parent2: np.ndarray,
                         mix_ratio: float = 0.5) -> Tuple[np.ndarray, np.ndarray]:
        """
        Uniform crossover: Each gene randomly inherited from either parent
        Best for: Maximum genetic diversity
        
        Args:
            parent1: First parent solution
            parent2: Second parent solution
            mix_ratio: Probability of inheriting from parent1
            
        Returns:
            Tuple of (offspring1, offspring2)
        """
        if np.random.random() > self.crossover_rate:
            return parent1.copy(), parent2.copy()
        
        mask = np.random.random(len(parent1)) < mix_ratio
        
        offspring1 = np.where(mask, parent1, parent2)
        offspring2 = np.where(mask, parent2, parent1)
        
        return offspring1, offspring2
    
    def arithmetic_crossover(self, parent1: np.ndarray, parent2: np.ndarray,
                            alpha: Optional[float] = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        Arithmetic crossover: Linear combination of parents
        Best for: Continuous optimization, exploitation
        
        Args:
            parent1: First parent solution
            parent2: Second parent solution
            alpha: Blending coefficient (None = random)
            
        Returns:
            Tuple of (offspring1, offspring2)
        """
        if np.random.random() > self.crossover_rate:
            return parent1.copy(), parent2.copy()
        
        if alpha is None:
            alpha = np.random.random()
        
        offspring1 = alpha * parent1 + (1 - alpha) * parent2
        offspring2 = (1 - alpha) * parent1 + alpha * parent2
        
        return offspring1, offspring2
    
    def simulated_binary_crossover(self, parent1: np.ndarray, parent2: np.ndarray,
                                   bounds: Tuple[float, float],
                                   eta: float = 20.0) -> Tuple[np.ndarray, np.ndarray]:
        """
        Simulated Binary Crossover (SBX): Mimics single-point crossover in binary
        Best for: Real-valued problems, balanced exploration-exploitation
        
        Args:
            parent1: First parent solution
            parent2: Second parent solution
            bounds: (lower_bound, upper_bound) tuple
            eta: Distribution index (higher = more like parents)
            
        Returns:
            Tuple of (offspring1, offspring2)
        """
        if np.random.random() > self.crossover_rate:
            return parent1.copy(), parent2.copy()
        
        lower, upper = bounds
        lower = np.broadcast_to(np.asarray(lower, dtype=float), len(parent1)).copy()
        upper = np.broadcast_to(np.asarray(upper, dtype=float), len(parent1)).copy()
        offspring1 = parent1.copy()
        offspring2 = parent2.copy()
        
        for i in range(len(parent1)):
            if np.random.random() < 0.5:
                if abs(parent1[i] - parent2[i]) > 1e-14:
                    y1 = min(parent1[i], parent2[i])
                    y2 = max(parent1[i], parent2[i])
                    
                    rand = np.random.random()
                    beta = 1.0 + (2.0 * (y1 - lower[i]) / (y2 - y1))
                    alpha = 2.0 - beta ** -(eta + 1.0)
                    
                    if rand <= 1.0 / alpha:
                        betaq = (rand * alpha) ** (1.0 / (eta + 1.0))
                    else:
                        betaq = (1.0 / (2.0 - rand * alpha)) ** (1.0 / (eta + 1.0))
                    
                    c1 = 0.5 * ((y1 + y2) - betaq * (y2 - y1))
                    
                    beta = 1.0 + (2.0 * (upper[i] - y2) / (y2 - y1))
                    alpha = 2.0 - beta ** -(eta + 1.0)
                    
                    if rand <= 1.0 / alpha:
                        betaq = (rand * alpha) ** (1.0 / (eta + 1.0))
                    else:
                        betaq = (1.0 / (2.0 - rand * alpha)) ** (1.0 / (eta + 1.0))
                    
                    c2 = 0.5 * ((y1 + y2) + betaq * (y2 - y1))
                    
                    offspring1[i] = np.clip(c1, lower[i], upper[i])
                    offspring2[i] = np.clip(c2, lower[i], upper[i])
        
        return offspring1, offspring2
    
    # ==================== SELECTION OPERATORS ====================
    
    def tournament_selection(self, population: np.ndarray, fitness: np.ndarray,
                           tournament_size: int = 3) -> np.ndarray:
        """
        Tournament selection: Select best from random subset
        
        Args:
            population: Population array
            fitness: Fitness values
            tournament_size: Size of tournament
            
        Returns:
            Selected individual
        """
        indices = np.random.choice(len(population), tournament_size, replace=False)
        tournament_fitness = fitness[indices]
        winner_idx = indices[np.argmin(tournament_fitness)]
        return population[winner_idx].copy()
    
    def roulette_wheel_selection(self, population: np.ndarray, 
                                fitness: np.ndarray) -> np.ndarray:
        """
        Roulette wheel selection: Probability proportional to fitness
        
        Args:
            population: Population array
            fitness: Fitness values (minimization)
            
        Returns:
            Selected individual
        """
        # Convert to maximization (invert for minimization)
        inv_fitness = 1.0 / (fitness + 1e-10)
        probabilities = inv_fitness / np.sum(inv_fitness)
        
        idx = np.random.choice(len(population), p=probabilities)
        return population[idx].copy()


class AdaptiveGeneticOperators(GeneticOperators):
    """Adaptive genetic operators that adjust rates based on progress"""
    
    def __init__(self, initial_mutation_rate: float = 0.1, 
                 initial_crossover_rate: float = 0.8):
        super().__init__(initial_mutation_rate, initial_crossover_rate)
        self.initial_mutation_rate = initial_mutation_rate
        self.initial_crossover_rate = initial_crossover_rate
        
    def adapt_rates(self, iteration: int, max_iterations: int, 
                    diversity: float) -> None:
        """
        Adapt mutation and crossover rates based on progress and diversity
        
        Args:
            iteration: Current iteration
            max_iterations: Maximum iterations
            diversity: Population diversity metric (0 to 1)
        """
        progress = iteration / max_iterations
        
        # Increase mutation when diversity is low (stuck in local optima)
        if diversity < 0.1:
            self.mutation_rate = min(0.5, self.initial_mutation_rate * 2)
        else:
            # Decrease mutation rate as optimization progresses
            self.mutation_rate = self.initial_mutation_rate * (1 - 0.5 * progress)
        
        # Adjust crossover rate inversely
        self.crossover_rate = self.initial_crossover_rate * (0.5 + 0.5 * (1 - progress))


def calculate_diversity(population: np.ndarray) -> float:
    """
    Calculate population diversity
    
    Args:
        population: Population array
        
    Returns:
        Diversity metric (0 to 1)
    """
    if len(population) < 2:
        return 1.0
    
    # Calculate average pairwise distance
    distances = []
    for i in range(len(population)):
        for j in range(i + 1, len(population)):
            dist = np.linalg.norm(population[i] - population[j])
            distances.append(dist)
    
    if not distances:
        return 1.0
    
    avg_distance = np.mean(distances)
    max_possible_distance = np.sqrt(len(population[0]))
    
    return min(1.0, avg_distance / max_possible_distance)
