"""
EPD-GA: Hybrid Genetic Algorithm with Evolutionary Population Dynamics
========================================================================

Combines:
1. Genetic Operators (mutation & crossover) for exploration/exploitation
2. EPD for self-organized criticality and guided repositioning
3. Adaptive mechanisms for dynamic optimization

This hybrid approach leverages:
- Mutation: Random exploration
- Crossover: Exploitation through recombination
- EPD: Self-organized repositioning of poor solutions near best ones
"""

import numpy as np
from typing import Callable, Tuple
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 AdaptiveGeneticOperators
from mha_toolbox.evolutionary_population_dynamics import AdaptiveEPD
from mha_toolbox.base import BaseOptimizer


class EPD_GA(BaseOptimizer):
    """
    Hybrid EPD-GA Algorithm
    
    Integrates genetic operations with evolutionary population dynamics
    for enhanced exploration-exploitation balance.
    
    Key Features:
    - Mutation for exploration
    - Crossover for exploitation  
    - EPD for self-organized criticality
    - Adaptive parameter control
    - Elite preservation
    """
    
    def __init__(self,
                 population_size: int = 50,
                 max_iterations: int = 100,
                 mutation_rate: float = 0.15,
                 crossover_rate: float = 0.8,
                 epd_frequency: int = 5,
                 epd_removal_ratio: float = 0.2,
                 elite_ratio: float = 0.1,
                 mutation_type: str = 'gaussian',
                 crossover_type: str = 'simulated_binary',
                 **kwargs):
        """
        Initialize EPD-GA
        
        Args:
            population_size: Size of population
            max_iterations: Maximum generations
            mutation_rate: Initial mutation probability
            crossover_rate: Crossover probability
            epd_frequency: Apply EPD every N generations
            epd_removal_ratio: Fraction to reposition with EPD
            elite_ratio: Fraction of best to preserve
            mutation_type: Type of mutation operator
            crossover_type: Type of crossover operator
        """
        super().__init__(population_size=population_size, max_iterations=max_iterations, **kwargs)
        
        self.mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate
        self.epd_frequency = epd_frequency
        self.epd_removal_ratio = epd_removal_ratio
        self.elite_ratio = elite_ratio
        self.mutation_type = mutation_type
        self.crossover_type = crossover_type
        
        # Initialize components
        self.genetic_ops = AdaptiveGeneticOperators(mutation_rate, crossover_rate)
        self.epd = AdaptiveEPD(removal_ratio=epd_removal_ratio)
        
        # Tracking
        self.diversity_history = []
        self.epd_applications = []
        
    def _optimize(self, objective_function: Callable, **kwargs) -> Tuple:
        """
        Run EPD-GA 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 diversity
            diversity = self._calculate_diversity(population)
            self.diversity_history.append(diversity)
            
            # Adapt genetic operators
            self.genetic_ops.adapt_rates(iteration, self.max_iterations_, diversity)
            
            # Apply EPD periodically
            if iteration % self.epd_frequency == 0 and iteration > 0:
                population, fitness = self._apply_epd_step(
                    population, fitness, iteration, objective_function
                )
                self.epd_applications.append(iteration)
            
            # Sort 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 through genetic operations
            new_population = []
            new_fitness = []
            
            # Preserve elite
            for i in range(n_elite):
                new_population.append(elite_pop[i])
                new_fitness.append(elite_fit[i])
            
            # Generate offspring with genetic operators
            while len(new_population) < self.population_size_:
                # Tournament selection
                parent1 = self.genetic_ops.tournament_selection(
                    population, fitness, tournament_size=3
                )
                parent2 = self.genetic_ops.tournament_selection(
                    population, fitness, tournament_size=3
                )
                
                # Crossover (exploitation)
                offspring1, offspring2 = self._apply_crossover(parent1, parent2)
                
                # Mutation (exploration)
                offspring1 = self._apply_mutation(offspring1)
                offspring2 = self._apply_mutation(offspring2)
                
                # Evaluate offspring
                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 solution
            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())
            
            # Progress reporting
            if (iteration + 1) % 10 == 0:
                epd_status = "EPD Applied" if iteration in self.epd_applications else ""
                print(f"Generation {iteration + 1}/{self.max_iterations_}, "
                      f"Best: {best_fitness:.6e}, "
                      f"Diversity: {diversity:.3f} {epd_status}")
        
        return best_solution, best_fitness, convergence_curve, local_fitness, local_positions
    
    def _apply_epd_step(self,
                       population: np.ndarray,
                       fitness: np.ndarray,
                       iteration: int,
                       objective_function: Callable) -> Tuple[np.ndarray, np.ndarray]:
        """
        Apply EPD to reposition poor solutions
        
        This implements self-organized criticality by moving worst
        solutions near best ones, focusing exploitation on promising regions.
        """
        bounds = (self.lower_bound_, self.upper_bound_)
        
        # Apply EPD with elite preservation
        n_elite = max(1, int(self.population_size_ * self.elite_ratio))
        elite_indices = np.argsort(fitness)[:n_elite]
        elite_pop = population[elite_indices].copy()
        elite_fit = fitness[elite_indices].copy()
        
        # Apply EPD dynamics
        new_pop, new_fit = self.epd.apply_dynamics(
            population, fitness, bounds, iteration, self.max_iterations_
        )
        
        # Re-evaluate repositioned individuals
        for i in range(len(new_pop)):
            if not np.array_equal(new_pop[i], population[i]):
                new_fit[i] = objective_function(new_pop[i])
        
        # Restore elite (guarantee no degradation)
        new_pop[elite_indices] = elite_pop
        new_fit[elite_indices] = elite_fit
        
        return new_pop, new_fit
    
    def _apply_mutation(self, individual: np.ndarray) -> np.ndarray:
        """Apply mutation operator"""
        bounds = (self.lower_bound_, self.upper_bound_)
        
        if self.mutation_type == 'uniform':
            return self.genetic_ops.uniform_mutation(individual, bounds)
        elif self.mutation_type == 'gaussian':
            return self.genetic_ops.gaussian_mutation(individual, bounds)
        elif self.mutation_type == 'polynomial':
            return self.genetic_ops.polynomial_mutation(individual, bounds)
        elif self.mutation_type == 'boundary':
            return self.genetic_ops.boundary_mutation(individual, bounds)
        else:
            return self.genetic_ops.gaussian_mutation(individual, bounds)
    
    def _apply_crossover(self, parent1: np.ndarray,
                        parent2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Apply crossover operator"""
        bounds = (self.lower_bound_, self.upper_bound_)
        
        if self.crossover_type == 'single_point':
            return self.genetic_ops.single_point_crossover(parent1, parent2)
        elif self.crossover_type == 'two_point':
            return self.genetic_ops.two_point_crossover(parent1, parent2)
        elif self.crossover_type == 'uniform':
            return self.genetic_ops.uniform_crossover(parent1, parent2)
        elif self.crossover_type == 'arithmetic':
            return self.genetic_ops.arithmetic_crossover(parent1, parent2)
        elif self.crossover_type == 'simulated_binary':
            return self.genetic_ops.simulated_binary_crossover(parent1, parent2, bounds)
        else:
            return self.genetic_ops.simulated_binary_crossover(parent1, parent2, bounds)
    
    def _calculate_diversity(self, population: np.ndarray) -> float:
        """Calculate population diversity metric"""
        if len(population) < 2:
            return 1.0
        
        distances = []
        n_samples = min(len(population), 20)  # Sample for efficiency
        sample_indices = np.random.choice(len(population), n_samples, replace=False)
        
        for i in range(len(sample_indices)):
            for j in range(i + 1, len(sample_indices)):
                idx1, idx2 = sample_indices[i], sample_indices[j]
                dist = np.linalg.norm(population[idx1] - population[idx2])
                distances.append(dist)
        
        if not distances:
            return 1.0
        
        avg_distance = np.mean(distances)
        max_distance = np.linalg.norm(np.asarray(self.upper_bound_) - np.asarray(self.lower_bound_))
        
        return float(min(1.0, avg_distance / max_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_)
        )
    
    def get_algorithm_info(self) -> dict:
        """Get algorithm information"""
        return {
            'name': 'EPD-GA',
            'full_name': 'Evolutionary Population Dynamics Genetic Algorithm',
            'type': 'Hybrid Evolutionary',
            'components': {
                'mutation': self.mutation_type,
                'crossover': self.crossover_type,
                'population_dynamics': 'Adaptive EPD',
                'selection': 'Tournament'
            },
            'parameters': {
                'population_size': self.population_size_,
                'mutation_rate': self.mutation_rate,
                'crossover_rate': self.crossover_rate,
                'epd_frequency': self.epd_frequency,
                'epd_removal_ratio': self.epd_removal_ratio,
                'elite_ratio': self.elite_ratio
            },
            'features': [
                'Adaptive genetic operators',
                'Self-organized criticality (EPD)',
                'Elite preservation',
                'Diversity maintenance',
                'Exploration-exploitation balance'
            ]
        }


# Create variants with different configurations

class EPD_GA_Explorative(EPD_GA):
    """EPD-GA variant focused on exploration"""
    
    def __init__(self, **kwargs):
        super().__init__(
            mutation_rate=0.25,  # Higher mutation
            crossover_rate=0.6,  # Lower crossover
            epd_frequency=10,     # Less frequent EPD
            mutation_type='uniform',
            **kwargs
        )


class EPD_GA_Exploitative(EPD_GA):
    """EPD-GA variant focused on exploitation"""
    
    def __init__(self, **kwargs):
        super().__init__(
            mutation_rate=0.05,  # Lower mutation
            crossover_rate=0.9,  # Higher crossover
            epd_frequency=3,      # More frequent EPD
            epd_removal_ratio=0.3,
            mutation_type='gaussian',
            **kwargs
        )


class EPD_GA_Balanced(EPD_GA):
    """EPD-GA variant with balanced exploration-exploitation"""
    
    def __init__(self, **kwargs):
        super().__init__(
            mutation_rate=0.15,
            crossover_rate=0.8,
            epd_frequency=5,
            mutation_type='gaussian',
            crossover_type='simulated_binary',
            **kwargs
        )
