"""
Standard Genetic Algorithm (GA)
=================================

Implements classic genetic algorithm with:
- Mutation for exploration
- Crossover for exploitation
- Tournament selection
- Elitism

Reference: Holland, J. H. (1992). Genetic algorithms. Scientific American, 267(1), 66-73.
"""

import numpy as np
from typing import Callable, Tuple, Optional
import sys
import os

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from mha_toolbox.genetic_operators import GeneticOperators, AdaptiveGeneticOperators
from mha_toolbox.base import BaseOptimizer


class GeneticAlgorithm(BaseOptimizer):
    """
    Standard Genetic Algorithm
    
    Features:
    - Multiple mutation strategies
    - Multiple crossover strategies
    - Tournament selection
    - Elitism
    - Adaptive operators
    """
    
    def __init__(self,
                 population_size: int = 50,
                 max_iterations: int = 100,
                 mutation_rate: float = 0.1,
                 crossover_rate: float = 0.8,
                 mutation_type: str = 'gaussian',
                 crossover_type: str = 'simulated_binary',
                 elite_ratio: float = 0.1,
                 tournament_size: int = 3,
                 adaptive: bool = True,
                 **kwargs):
        """
        Initialize Genetic Algorithm
        
        Args:
            population_size: Size of population
            max_iterations: Maximum generations
            mutation_rate: Initial mutation probability
            crossover_rate: Crossover probability
            mutation_type: 'uniform', 'gaussian', 'polynomial', 'boundary'
            crossover_type: 'single_point', 'two_point', 'uniform', 'arithmetic', 'simulated_binary'
            elite_ratio: Fraction of best to preserve
            tournament_size: Size of tournament selection
            adaptive: Use adaptive operators
        """
        super().__init__(population_size=population_size, max_iterations=max_iterations, **kwargs)
        
        self.mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate
        self.mutation_type = mutation_type
        self.crossover_type = crossover_type
        self.elite_ratio = elite_ratio
        self.tournament_size = tournament_size
        self.adaptive = adaptive
        
        # Initialize operators
        if adaptive:
            self.operators = AdaptiveGeneticOperators(mutation_rate, crossover_rate)
        else:
            self.operators = GeneticOperators(mutation_rate, crossover_rate)
        
        # Store for analysis
        self.diversity_history = []
        
    def _optimize(self, objective_function: Callable, **kwargs) -> Tuple:
        """
        Run genetic algorithm optimization
        
        Args:
            objective_function: Function to optimize
            
        Returns:
            Tuple of (best_solution, best_fitness, convergence_curve, local_fitness, local_positions)
        """
        # Initialize population
        population = self._initialize_population()
        fitness = np.array([objective_function(ind) for ind in population])
        
        # Track best
        best_idx = np.argmin(fitness)
        best_solution = population[best_idx].copy()
        best_fitness = fitness[best_idx]
        convergence_curve = [best_fitness]
        local_fitness = [fitness.copy()]
        local_positions = [population.copy()]
        
        n_elite = max(1, int(self.population_size_ * self.elite_ratio))
        
        # Main evolution loop
        for iteration in range(self.max_iterations_):
            # Calculate and store diversity
            diversity = self._calculate_diversity(population)
            self.diversity_history.append(diversity)
            
            # Adapt operators if enabled
            if self.adaptive:
                self.operators.adapt_rates(iteration, self.max_iterations_, diversity)
            
            # Sort population by fitness
            sorted_indices = np.argsort(fitness)
            elite_pop = population[sorted_indices[:n_elite]].copy()
            elite_fit = fitness[sorted_indices[:n_elite]].copy()
            
            # Generate new population
            new_population = []
            new_fitness = []
            
            # Keep elite
            for i in range(n_elite):
                new_population.append(elite_pop[i])
                new_fitness.append(elite_fit[i])
            
            # Generate offspring
            while len(new_population) < self.population_size_:
                # Selection
                parent1 = self.operators.tournament_selection(
                    population, fitness, self.tournament_size
                )
                parent2 = self.operators.tournament_selection(
                    population, fitness, self.tournament_size
                )
                
                # Crossover
                offspring1, offspring2 = self._apply_crossover(parent1, parent2)
                
                # Mutation
                offspring1 = self._apply_mutation(offspring1)
                offspring2 = self._apply_mutation(offspring2)
                
                # Evaluate
                fit1 = objective_function(offspring1)
                fit2 = objective_function(offspring2)
                
                new_population.append(offspring1)
                new_fitness.append(fit1)
                
                if len(new_population) < self.population_size_:
                    new_population.append(offspring2)
                    new_fitness.append(fit2)
            
            # Update population
            population = np.array(new_population[:self.population_size_])
            fitness = np.array(new_fitness[:self.population_size_])
            
            # Update best
            current_best_idx = np.argmin(fitness)
            if fitness[current_best_idx] < best_fitness:
                best_fitness = fitness[current_best_idx]
                best_solution = population[current_best_idx].copy()
            
            convergence_curve.append(best_fitness)
            local_fitness.append(fitness.copy())
            local_positions.append(population.copy())
            
            # Optional: Print progress
            if (iteration + 1) % 10 == 0:
                print(f"Generation {iteration + 1}/{self.max_iterations_}, "
                      f"Best Fitness: {best_fitness:.6e}, "
                      f"Diversity: {diversity:.3f}")
        
        return best_solution, best_fitness, convergence_curve, local_fitness, local_positions
    
    def _apply_mutation(self, individual: np.ndarray) -> np.ndarray:
        """Apply selected mutation operator"""
        lower = float(self.lower_bound_[0]) if isinstance(self.lower_bound_, np.ndarray) else float(self.lower_bound_)
        upper = float(self.upper_bound_[0]) if isinstance(self.upper_bound_, np.ndarray) else float(self.upper_bound_)
        bounds = (lower, upper)
        
        if self.mutation_type == 'uniform':
            return self.operators.uniform_mutation(individual, bounds)
        elif self.mutation_type == 'gaussian':
            return self.operators.gaussian_mutation(individual, bounds)
        elif self.mutation_type == 'polynomial':
            return self.operators.polynomial_mutation(individual, bounds)
        elif self.mutation_type == 'boundary':
            return self.operators.boundary_mutation(individual, bounds)
        else:
            return self.operators.gaussian_mutation(individual, bounds)
    
    def _apply_crossover(self, parent1: np.ndarray, 
                        parent2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Apply selected crossover operator"""
        lower = float(self.lower_bound_[0]) if isinstance(self.lower_bound_, np.ndarray) else float(self.lower_bound_)
        upper = float(self.upper_bound_[0]) if isinstance(self.upper_bound_, np.ndarray) else float(self.upper_bound_)
        bounds = (lower, upper)
        
        if self.crossover_type == 'single_point':
            return self.operators.single_point_crossover(parent1, parent2)
        elif self.crossover_type == 'two_point':
            return self.operators.two_point_crossover(parent1, parent2)
        elif self.crossover_type == 'uniform':
            return self.operators.uniform_crossover(parent1, parent2)
        elif self.crossover_type == 'arithmetic':
            return self.operators.arithmetic_crossover(parent1, parent2)
        elif self.crossover_type == 'simulated_binary':
            return self.operators.simulated_binary_crossover(parent1, parent2, bounds)
        else:
            return self.operators.simulated_binary_crossover(parent1, parent2, bounds)
    
    def _calculate_diversity(self, population: np.ndarray) -> float:
        """Calculate population diversity"""
        if len(population) < 2:
            return 1.0
        
        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)
        # Handle case where bounds are arrays (per-dimension)
        bound_range = np.asarray(self.upper_bound_) - np.asarray(self.lower_bound_)
        if isinstance(bound_range, np.ndarray) and bound_range.ndim > 0:
            max_possible_distance = np.sqrt(np.sum(bound_range ** 2))
        else:
            max_possible_distance = np.sqrt(self.dimensions_) * float(bound_range)
        
        return float(min(1.0, avg_distance / max_possible_distance))
    
    def _initialize_population(self) -> np.ndarray:
        """Initialize random population"""
        return np.random.uniform(
            self.lower_bound_,
            self.upper_bound_,
            (self.population_size_, self.dimensions_)
        )


# Alias for convenience
GA = GeneticAlgorithm
