"""
Evolutionary Population Dynamics (EPD)
=======================================

Implements population management based on self-organized criticality.
The EPD method removes the least optimal solutions and repositions them
near the best ones, focusing on exploitation of promising regions.

Key Features:
- Self-organized criticality
- Exploitation-focused repositioning
- Dynamic population management
- Maintains diversity while converging
"""

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


class EvolutionaryPopulationDynamics:
    """
    EPD manager for evolutionary algorithms
    
    Implements self-organized criticality by:
    1. Identifying worst-performing individuals
    2. Removing them from poor regions
    3. Repositioning near best solutions
    4. Adding controlled randomness
    """
    
    def __init__(self, 
                 removal_ratio: float = 0.2,
                 repositioning_radius: float = 0.1,
                 adaptive: bool = True):
        """
        Initialize EPD manager
        
        Args:
            removal_ratio: Fraction of worst solutions to remove/reposition
            repositioning_radius: How close to reposition near best (relative)
            adaptive: Whether to adapt parameters based on progress
        """
        self.removal_ratio = removal_ratio
        self.repositioning_radius = repositioning_radius
        self.adaptive = adaptive
        self.initial_radius = repositioning_radius
        
    def apply_dynamics(self,
                      population: np.ndarray,
                      fitness: np.ndarray,
                      bounds: Tuple[float, float],
                      iteration: int = 0,
                      max_iterations: int = 100) -> Tuple[np.ndarray, np.ndarray]:
        """
        Apply evolutionary population dynamics
        
        Args:
            population: Current population
            fitness: Fitness values
            bounds: (lower_bound, upper_bound)
            iteration: Current iteration
            max_iterations: Maximum iterations
            
        Returns:
            Tuple of (new_population, new_fitness)
        """
        pop_size = len(population)
        n_remove = max(1, int(pop_size * self.removal_ratio))
        
        # Adapt repositioning radius if enabled
        if self.adaptive:
            progress = iteration / max_iterations
            # Start with larger radius, decrease for exploitation
            self.repositioning_radius = self.initial_radius * (1 - 0.7 * progress)
        
        # Identify worst individuals
        worst_indices = np.argsort(fitness)[-n_remove:]
        
        # Identify best individuals
        best_indices = np.argsort(fitness)[:max(1, n_remove // 2)]
        
        # Create new population (copy to avoid modifying original)
        new_population = population.copy()
        
        # Reposition worst near best with self-organized criticality
        for idx in worst_indices:
            new_population[idx] = self._reposition_individual(
                population[best_indices],
                bounds,
                iteration,
                max_iterations
            )
        
        return new_population, fitness.copy()
    
    def _reposition_individual(self,
                              elite_population: np.ndarray,
                              bounds: Tuple[float, float],
                              iteration: int,
                              max_iterations: int) -> np.ndarray:
        """
        Reposition an individual near elite solutions
        
        Uses self-organized criticality principle:
        - Near best solutions (exploitation)
        - With controlled randomness (exploration)
        - Adaptive based on progress
        
        Args:
            elite_population: Best individuals to reposition near
            bounds: Search space bounds
            iteration: Current iteration
            max_iterations: Maximum iterations
            
        Returns:
            Repositioned individual
        """
        lower, upper = bounds
        range_size = upper - lower
        
        # Select random elite as anchor
        anchor = elite_population[np.random.randint(len(elite_population))]
        
        # Generate position near anchor with Gaussian distribution
        # Radius decreases with iterations for increased exploitation
        radius = self.repositioning_radius * range_size
        
        # Add Gaussian noise around anchor
        noise = np.random.normal(0, radius, len(anchor))
        new_individual = anchor + noise
        
        # Add small uniform component for diversity
        uniform_component = np.random.uniform(-radius * 0.2, radius * 0.2, len(anchor))
        new_individual += uniform_component
        
        # Ensure bounds
        new_individual = np.clip(new_individual, lower, upper)
        
        return new_individual
    
    def apply_criticality_dynamics(self,
                                  population: np.ndarray,
                                  fitness: np.ndarray,
                                  bounds: Tuple[float, float],
                                  threshold_percentile: float = 70) -> Tuple[np.ndarray, np.ndarray]:
        """
        Apply self-organized criticality dynamics
        
        Removes solutions below fitness threshold and repositions
        near best, creating avalanche-like improvement effects
        
        Args:
            population: Current population
            fitness: Fitness values
            bounds: Search space bounds
            threshold_percentile: Fitness percentile threshold
            
        Returns:
            Tuple of (new_population, new_fitness)
        """
        threshold = np.percentile(fitness, threshold_percentile)
        
        # Identify solutions below threshold (critical state)
        critical_indices = np.where(fitness > threshold)[0]
        
        if len(critical_indices) == 0:
            return population.copy(), fitness.copy()
        
        # Identify elite solutions
        elite_indices = np.where(fitness <= np.percentile(fitness, 30))[0]
        
        if len(elite_indices) == 0:
            return population.copy(), fitness.copy()
        
        new_population = population.copy()
        
        # Reposition critical solutions near elite
        for idx in critical_indices:
            new_population[idx] = self._reposition_individual(
                population[elite_indices],
                bounds,
                0,
                100
            )
        
        return new_population, fitness.copy()


class AdaptiveEPD(EvolutionaryPopulationDynamics):
    """
    Adaptive EPD with dynamic parameter adjustment
    """
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.stagnation_counter = 0
        self.best_fitness_history = []
        
    def apply_dynamics(self,
                      population: np.ndarray,
                      fitness: np.ndarray,
                      bounds: Tuple[float, float],
                      iteration: int = 0,
                      max_iterations: int = 100) -> Tuple[np.ndarray, np.ndarray]:
        """
        Apply adaptive EPD with stagnation detection
        """
        # Track best fitness
        current_best = np.min(fitness)
        self.best_fitness_history.append(current_best)
        
        # Detect stagnation
        if len(self.best_fitness_history) > 10:
            recent_improvement = (
                self.best_fitness_history[-10] - self.best_fitness_history[-1]
            )
            
            if recent_improvement < 1e-6:
                self.stagnation_counter += 1
            else:
                self.stagnation_counter = 0
            
            # Increase removal ratio during stagnation
            if self.stagnation_counter > 5:
                self.removal_ratio = min(0.4, self.removal_ratio * 1.1)
                self.repositioning_radius = min(0.3, self.repositioning_radius * 1.2)
            else:
                # Reset to initial values
                self.removal_ratio = 0.2
                self.repositioning_radius = self.initial_radius
        
        return super().apply_dynamics(
            population, fitness, bounds, iteration, max_iterations
        )


class HybridEPD:
    """
    Hybrid EPD combining multiple strategies
    """
    
    def __init__(self):
        self.standard_epd = EvolutionaryPopulationDynamics()
        self.adaptive_epd = AdaptiveEPD()
        self.use_adaptive = False
        
    def apply_dynamics(self,
                      population: np.ndarray,
                      fitness: np.ndarray,
                      bounds: Tuple[float, float],
                      iteration: int = 0,
                      max_iterations: int = 100) -> Tuple[np.ndarray, np.ndarray]:
        """
        Apply hybrid EPD strategy
        
        Switches between standard and adaptive based on progress
        """
        progress = iteration / max_iterations
        
        # Use adaptive EPD in middle phase when stagnation likely
        if 0.3 < progress < 0.8:
            self.use_adaptive = True
            return self.adaptive_epd.apply_dynamics(
                population, fitness, bounds, iteration, max_iterations
            )
        else:
            self.use_adaptive = False
            return self.standard_epd.apply_dynamics(
                population, fitness, bounds, iteration, max_iterations
            )


def apply_epd_with_elitism(population: np.ndarray,
                          fitness: np.ndarray,
                          bounds: Tuple[float, float],
                          epd: EvolutionaryPopulationDynamics,
                          elite_ratio: float = 0.1) -> Tuple[np.ndarray, np.ndarray]:
    """
    Apply EPD while preserving elite solutions
    
    Args:
        population: Current population
        fitness: Fitness values
        bounds: Search space bounds
        epd: EPD instance
        elite_ratio: Fraction of best to preserve
        
    Returns:
        Tuple of (new_population, new_fitness)
    """
    n_elite = max(1, int(len(population) * elite_ratio))
    
    # Identify elite
    elite_indices = np.argsort(fitness)[:n_elite]
    elite_pop = population[elite_indices].copy()
    elite_fit = fitness[elite_indices].copy()
    
    # Apply EPD to rest
    new_pop, new_fit = epd.apply_dynamics(population, fitness, bounds)
    
    # Restore elite
    new_pop[:n_elite] = elite_pop
    new_fit[:n_elite] = elite_fit
    
    return new_pop, new_fit
