"""
Evolutionary Prompt Optimization for Toxo.

This module implements genetic algorithms and evolutionary strategies
for optimizing soft prompts, complementing the SPSA approach.
"""

import numpy as np
import asyncio
from typing import List, Dict, Any, Optional, Tuple, Callable
from dataclasses import dataclass, field
import random
from concurrent.futures import ThreadPoolExecutor
import copy

from ..core.config import TrainingConfig
from ..utils.logger import get_logger
from ..utils.exceptions import TrainingError


@dataclass
class Individual:
    """Individual in the evolutionary population."""
    genome: np.ndarray  # Soft prompt vector
    fitness: float = 0.0
    age: int = 0
    diversity_score: float = 0.0
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def __post_init__(self):
        """Initialize individual metadata."""
        if not self.metadata:
            self.metadata = {
                "birth_generation": 0,
                "parent_fitness": None,
                "mutation_count": 0,
                "crossover_count": 0
            }


@dataclass
class EvolutionaryParams:
    """Parameters for evolutionary optimization."""
    population_size: int = 50
    elite_ratio: float = 0.2
    mutation_rate: float = 0.1
    crossover_rate: float = 0.8
    diversity_pressure: float = 0.1
    max_generations: int = 100
    stagnation_limit: int = 20
    tournament_size: int = 3
    adaptive_mutation: bool = True
    niching: bool = True
    

class EvolutionaryOptimizer:
    """
    Evolutionary optimizer for soft prompt training.
    
    Uses genetic algorithms with advanced features like:
    - Elitism and diversity preservation
    - Adaptive mutation rates
    - Niching for maintaining population diversity
    - Multi-objective optimization
    """
    
    def __init__(
        self,
        dimension: int,
        length: int,
        config: TrainingConfig,
        params: Optional[EvolutionaryParams] = None
    ):
        """
        Initialize evolutionary optimizer.
        
        Args:
            dimension: Embedding dimension
            length: Soft prompt length
            config: Training configuration
            params: Evolutionary parameters
        """
        self.dimension = dimension
        self.length = length
        self.config = config
        self.params = params or EvolutionaryParams()
        self.logger = get_logger("toxo.evolutionary")
        
        # Evolution state
        self.population: List[Individual] = []
        self.generation = 0
        self.best_individual: Optional[Individual] = None
        self.fitness_history: List[float] = []
        self.diversity_history: List[float] = []
        self.stagnation_counter = 0
        
        # Statistics
        self.total_evaluations = 0
        self.convergence_data = []
        
        self.logger.info(f"Evolutionary optimizer initialized: pop_size={self.params.population_size}")
    
    def initialize_population(self, seed_individuals: Optional[List[np.ndarray]] = None) -> None:
        """
        Initialize the population with random or seeded individuals.
        
        Args:
            seed_individuals: Optional list of seed genomes
        """
        self.population = []
        
        # Add seed individuals if provided
        if seed_individuals:
            for i, genome in enumerate(seed_individuals[:self.params.population_size]):
                individual = Individual(
                    genome=genome.copy(),
                    metadata={"birth_generation": 0, "source": "seed", "seed_index": i}
                )
                self.population.append(individual)
        
        # Fill remaining slots with random individuals
        while len(self.population) < self.params.population_size:
            genome = self._create_random_genome()
            individual = Individual(
                genome=genome,
                metadata={"birth_generation": 0, "source": "random"}
            )
            self.population.append(individual)
        
        self.logger.info(f"Population initialized with {len(self.population)} individuals")
    
    def _create_random_genome(self) -> np.ndarray:
        """Create a random genome (soft prompt vector)."""
        return np.random.normal(0, 0.01, (self.length, self.dimension)).astype(np.float32)
    
    async def evolve(
        self,
        fitness_function: Callable[[np.ndarray], float],
        max_generations: Optional[int] = None
    ) -> Individual:
        """
        Run the evolutionary optimization process.
        
        Args:
            fitness_function: Function to evaluate fitness of individuals
            max_generations: Maximum number of generations
            
        Returns:
            Best individual found
        """
        max_gens = max_generations or self.params.max_generations
        
        self.logger.info(f"Starting evolution for {max_gens} generations")
        
        # Initialize if not done
        if not self.population:
            self.initialize_population()
        
        # Evaluate initial population
        await self._evaluate_population(fitness_function)
        
        for generation in range(max_gens):
            self.generation = generation
            
            # Create next generation
            new_population = await self._create_next_generation()
            
            # Evaluate new individuals
            await self._evaluate_new_individuals(new_population, fitness_function)
            
            # Update population
            self.population = new_population
            
            # Update statistics
            self._update_statistics()
            
            # Check convergence
            if self._check_convergence():
                self.logger.info(f"Convergence reached at generation {generation}")
                break
            
            # Log progress
            if generation % 10 == 0:
                self._log_generation_stats(generation)
        
        self.logger.info(f"Evolution completed after {self.generation + 1} generations")
        return self.best_individual
    
    async def _evaluate_population(self, fitness_function: Callable[[np.ndarray], float]) -> None:
        """Evaluate fitness for all individuals in population."""
        tasks = []
        for individual in self.population:
            if individual.fitness == 0.0:  # Not evaluated yet
                task = self._evaluate_individual(individual, fitness_function)
                tasks.append(task)
        
        if tasks:
            await asyncio.gather(*tasks)
            self.total_evaluations += len(tasks)
    
    async def _evaluate_individual(
        self, 
        individual: Individual, 
        fitness_function: Callable[[np.ndarray], float]
    ) -> None:
        """Evaluate a single individual."""
        try:
            # Run fitness evaluation
            individual.fitness = await asyncio.get_event_loop().run_in_executor(
                None, fitness_function, individual.genome
            )
            
            # Calculate diversity score
            individual.diversity_score = self._calculate_diversity_score(individual)
            
        except Exception as e:
            self.logger.warning(f"Fitness evaluation failed: {str(e)}")
            individual.fitness = 0.0  # Assign minimum fitness on failure
    
    def _calculate_diversity_score(self, individual: Individual) -> float:
        """Calculate diversity score relative to population."""
        if len(self.population) <= 1:
            return 1.0
        
        distances = []
        for other in self.population:
            if other is not individual:
                distance = np.linalg.norm(individual.genome - other.genome)
                distances.append(distance)
        
        return np.mean(distances) if distances else 1.0
    
    async def _create_next_generation(self) -> List[Individual]:
        """Create the next generation through selection, crossover, and mutation."""
        new_population = []
        
        # Elitism: Keep best individuals
        elite_count = int(self.params.population_size * self.params.elite_ratio)
        elite_individuals = self._select_elite(elite_count)
        
        # Age elite individuals
        for individual in elite_individuals:
            individual.age += 1
            new_population.append(individual)
        
        # Generate offspring to fill remaining slots
        while len(new_population) < self.params.population_size:
            if random.random() < self.params.crossover_rate:
                # Crossover
                parent1 = self._tournament_selection()
                parent2 = self._tournament_selection()
                offspring1, offspring2 = self._crossover(parent1, parent2)
                
                # Mutate offspring
                if random.random() < self._get_adaptive_mutation_rate():
                    self._mutate(offspring1)
                if random.random() < self._get_adaptive_mutation_rate():
                    self._mutate(offspring2)
                
                new_population.extend([offspring1, offspring2])
            else:
                # Direct mutation of selected individual
                parent = self._tournament_selection()
                offspring = self._clone_individual(parent)
                self._mutate(offspring)
                new_population.append(offspring)
        
        # Trim to exact population size
        return new_population[:self.params.population_size]
    
    def _select_elite(self, count: int) -> List[Individual]:
        """Select elite individuals based on fitness and diversity."""
        if self.params.niching:
            # Multi-objective selection considering fitness and diversity
            scored_individuals = []
            for individual in self.population:
                score = (individual.fitness + 
                        self.params.diversity_pressure * individual.diversity_score)
                scored_individuals.append((score, individual))
            
            scored_individuals.sort(key=lambda x: x[0], reverse=True)
            return [individual for _, individual in scored_individuals[:count]]
        else:
            # Simple fitness-based selection
            sorted_pop = sorted(self.population, key=lambda x: x.fitness, reverse=True)
            return sorted_pop[:count]
    
    def _tournament_selection(self) -> Individual:
        """Select individual using tournament selection."""
        tournament_size = min(self.params.tournament_size, len(self.population))
        candidates = random.sample(self.population, tournament_size)
        return max(candidates, key=lambda x: x.fitness)
    
    def _crossover(self, parent1: Individual, parent2: Individual) -> Tuple[Individual, Individual]:
        """Create offspring through crossover."""
        # Blend crossover (BLX-α)
        alpha = 0.1
        
        offspring1_genome = np.zeros_like(parent1.genome)
        offspring2_genome = np.zeros_like(parent2.genome)
        
        for i in range(self.length):
            for j in range(self.dimension):
                p1_val = parent1.genome[i, j]
                p2_val = parent2.genome[i, j]
                
                # Calculate range
                d = abs(p1_val - p2_val)
                min_val = min(p1_val, p2_val) - alpha * d
                max_val = max(p1_val, p2_val) + alpha * d
                
                # Generate offspring values
                offspring1_genome[i, j] = random.uniform(min_val, max_val)
                offspring2_genome[i, j] = random.uniform(min_val, max_val)
        
        # Create offspring individuals
        offspring1 = Individual(
            genome=offspring1_genome,
            metadata={
                "birth_generation": self.generation + 1,
                "parent1_fitness": parent1.fitness,
                "parent2_fitness": parent2.fitness,
                "crossover_count": 1
            }
        )
        
        offspring2 = Individual(
            genome=offspring2_genome,
            metadata={
                "birth_generation": self.generation + 1,
                "parent1_fitness": parent1.fitness,
                "parent2_fitness": parent2.fitness,
                "crossover_count": 1
            }
        )
        
        return offspring1, offspring2
    
    def _mutate(self, individual: Individual) -> None:
        """Apply mutation to an individual."""
        mutation_strength = self._get_adaptive_mutation_strength()
        
        # Gaussian mutation
        mutation = np.random.normal(0, mutation_strength, individual.genome.shape)
        individual.genome += mutation
        
        # Update metadata
        individual.metadata["mutation_count"] = individual.metadata.get("mutation_count", 0) + 1
    
    def _clone_individual(self, individual: Individual) -> Individual:
        """Create a copy of an individual."""
        return Individual(
            genome=individual.genome.copy(),
            metadata={
                "birth_generation": self.generation + 1,
                "parent_fitness": individual.fitness,
                "cloned_from": individual.metadata.get("source", "unknown")
            }
        )
    
    def _get_adaptive_mutation_rate(self) -> float:
        """Get adaptive mutation rate based on population diversity."""
        if not self.params.adaptive_mutation:
            return self.params.mutation_rate
        
        # Increase mutation rate if population is converging
        if len(self.diversity_history) > 5:
            recent_diversity = np.mean(self.diversity_history[-5:])
            if recent_diversity < 0.1:  # Low diversity
                return min(self.params.mutation_rate * 2.0, 0.5)
        
        return self.params.mutation_rate
    
    def _get_adaptive_mutation_strength(self) -> float:
        """Get adaptive mutation strength."""
        base_strength = 0.01
        
        # Decrease mutation strength as we find better solutions
        if self.best_individual and self.best_individual.fitness > 0:
            return base_strength * (1.0 / (1.0 + self.best_individual.fitness))
        
        return base_strength
    
    async def _evaluate_new_individuals(
        self, 
        population: List[Individual], 
        fitness_function: Callable[[np.ndarray], float]
    ) -> None:
        """Evaluate only new individuals in the population."""
        new_individuals = [ind for ind in population if ind.fitness == 0.0]
        
        if new_individuals:
            tasks = []
            for individual in new_individuals:
                task = self._evaluate_individual(individual, fitness_function)
                tasks.append(task)
            
            await asyncio.gather(*tasks)
            self.total_evaluations += len(tasks)
    
    def _update_statistics(self) -> None:
        """Update evolution statistics."""
        if not self.population:
            return
        
        # Find best individual
        current_best = max(self.population, key=lambda x: x.fitness)
        if self.best_individual is None or current_best.fitness > self.best_individual.fitness:
            self.best_individual = copy.deepcopy(current_best)
            self.stagnation_counter = 0
        else:
            self.stagnation_counter += 1
        
        # Update history
        self.fitness_history.append(current_best.fitness)
        
        # Calculate population diversity
        if len(self.population) > 1:
            diversities = [ind.diversity_score for ind in self.population]
            avg_diversity = np.mean(diversities)
            self.diversity_history.append(avg_diversity)
        
        # Store convergence data
        self.convergence_data.append({
            "generation": self.generation,
            "best_fitness": current_best.fitness,
            "avg_fitness": np.mean([ind.fitness for ind in self.population]),
            "diversity": self.diversity_history[-1] if self.diversity_history else 0.0,
            "evaluations": self.total_evaluations
        })
    
    def _check_convergence(self) -> bool:
        """Check if evolution has converged."""
        # Stagnation-based convergence
        if self.stagnation_counter >= self.params.stagnation_limit:
            return True
        
        # Diversity-based convergence
        if len(self.diversity_history) > 10:
            recent_diversity = np.mean(self.diversity_history[-10:])
            if recent_diversity < 0.01:  # Very low diversity
                return True
        
        return False
    
    def _log_generation_stats(self, generation: int) -> None:
        """Log statistics for current generation."""
        if not self.population:
            return
        
        fitnesses = [ind.fitness for ind in self.population]
        best_fitness = max(fitnesses)
        avg_fitness = np.mean(fitnesses)
        diversity = self.diversity_history[-1] if self.diversity_history else 0.0
        
        self.logger.info(
            f"Gen {generation}: Best={best_fitness:.4f}, "
            f"Avg={avg_fitness:.4f}, Div={diversity:.4f}, "
            f"Stag={self.stagnation_counter}, Evals={self.total_evaluations}"
        )
    
    def get_population_stats(self) -> Dict[str, Any]:
        """Get detailed population statistics."""
        if not self.population:
            return {}
        
        fitnesses = [ind.fitness for ind in self.population]
        ages = [ind.age for ind in self.population]
        
        return {
            "generation": self.generation,
            "population_size": len(self.population),
            "best_fitness": max(fitnesses),
            "worst_fitness": min(fitnesses),
            "avg_fitness": np.mean(fitnesses),
            "fitness_std": np.std(fitnesses),
            "avg_age": np.mean(ages),
            "max_age": max(ages),
            "diversity": self.diversity_history[-1] if self.diversity_history else 0.0,
            "stagnation_counter": self.stagnation_counter,
            "total_evaluations": self.total_evaluations,
            "convergence_data": self.convergence_data
        }
    
    def get_best_genome(self) -> Optional[np.ndarray]:
        """Get the genome of the best individual."""
        return self.best_individual.genome if self.best_individual else None
    
    def save_state(self) -> Dict[str, Any]:
        """Save the current state of the optimizer."""
        population_data = []
        for ind in self.population:
            population_data.append({
                "genome": ind.genome.tolist(),
                "fitness": ind.fitness,
                "age": ind.age,
                "diversity_score": ind.diversity_score,
                "metadata": ind.metadata
            })
        
        return {
            "dimension": self.dimension,
            "length": self.length,
            "generation": self.generation,
            "population": population_data,
            "best_individual": {
                "genome": self.best_individual.genome.tolist(),
                "fitness": self.best_individual.fitness,
                "metadata": self.best_individual.metadata
            } if self.best_individual else None,
            "fitness_history": self.fitness_history,
            "diversity_history": self.diversity_history,
            "stagnation_counter": self.stagnation_counter,
            "total_evaluations": self.total_evaluations,
            "convergence_data": self.convergence_data
        }
    
    def load_state(self, state: Dict[str, Any]) -> None:
        """Load state from saved data."""
        self.generation = state["generation"]
        self.stagnation_counter = state["stagnation_counter"]
        self.total_evaluations = state["total_evaluations"]
        self.fitness_history = state["fitness_history"]
        self.diversity_history = state["diversity_history"]
        self.convergence_data = state.get("convergence_data", [])
        
        # Reconstruct population
        self.population = []
        for ind_data in state["population"]:
            individual = Individual(
                genome=np.array(ind_data["genome"]),
                fitness=ind_data["fitness"],
                age=ind_data["age"],
                diversity_score=ind_data["diversity_score"],
                metadata=ind_data["metadata"]
            )
            self.population.append(individual)
        
        # Reconstruct best individual
        if state["best_individual"]:
            best_data = state["best_individual"]
            self.best_individual = Individual(
                genome=np.array(best_data["genome"]),
                fitness=best_data["fitness"],
                metadata=best_data["metadata"]
            )
        
        self.logger.info(f"Loaded evolutionary state at generation {self.generation}")


class HybridOptimizer:
    """
    Hybrid optimizer combining SPSA and evolutionary approaches.
    
    Uses evolutionary algorithms for global exploration and SPSA for local refinement.
    """
    
    def __init__(
        self,
        dimension: int,
        length: int,
        config: TrainingConfig,
        evolutionary_params: Optional[EvolutionaryParams] = None
    ):
        """Initialize hybrid optimizer."""
        self.dimension = dimension
        self.length = length
        self.config = config
        self.logger = get_logger("toxo.hybrid_optimizer")
        
        # Initialize components
        self.evolutionary = EvolutionaryOptimizer(dimension, length, config, evolutionary_params)
        
        # Hybrid parameters
        self.global_phase_generations = 30
        self.local_refinement_iterations = 20
        
    async def optimize(
        self,
        fitness_function: Callable[[np.ndarray], float],
        max_iterations: int = 100
    ) -> np.ndarray:
        """
        Run hybrid optimization.
        
        Args:
            fitness_function: Function to evaluate fitness
            max_iterations: Maximum total iterations
            
        Returns:
            Best genome found
        """
        self.logger.info("Starting hybrid optimization (Evolutionary + SPSA)")
        
        # Phase 1: Global exploration with evolutionary algorithm
        self.logger.info("Phase 1: Global exploration with evolutionary algorithm")
        best_individual = await self.evolutionary.evolve(
            fitness_function, 
            self.global_phase_generations
        )
        
        global_best = best_individual.genome
        self.logger.info(f"Global phase complete. Best fitness: {best_individual.fitness:.4f}")
        
        # Phase 2: Local refinement with SPSA
        self.logger.info("Phase 2: Local refinement with SPSA")
        refined_genome = await self._spsa_refinement(
            global_best, 
            fitness_function, 
            self.local_refinement_iterations
        )
        
        self.logger.info("Hybrid optimization complete")
        return refined_genome
    
    async def _spsa_refinement(
        self,
        initial_genome: np.ndarray,
        fitness_function: Callable[[np.ndarray], float],
        iterations: int
    ) -> np.ndarray:
        """Refine genome using SPSA."""
        current_genome = initial_genome.copy()
        best_genome = initial_genome.copy()
        best_fitness = await asyncio.get_event_loop().run_in_executor(
            None, fitness_function, best_genome
        )
        
        learning_rate = 0.01
        perturbation_scale = 0.01
        
        for iteration in range(iterations):
            # Generate perturbation
            perturbation = np.random.choice([-1, 1], size=current_genome.shape).astype(np.float32)
            
            # Evaluate at perturbed points
            genome_plus = current_genome + perturbation_scale * perturbation
            genome_minus = current_genome - perturbation_scale * perturbation
            
            fitness_plus = await asyncio.get_event_loop().run_in_executor(
                None, fitness_function, genome_plus
            )
            fitness_minus = await asyncio.get_event_loop().run_in_executor(
                None, fitness_function, genome_minus
            )
            
            # Estimate gradient
            gradient = ((fitness_plus - fitness_minus) / (2 * perturbation_scale)) * perturbation
            
            # Update genome
            current_genome += learning_rate * gradient
            
            # Evaluate current genome
            current_fitness = await asyncio.get_event_loop().run_in_executor(
                None, fitness_function, current_genome
            )
            
            # Update best
            if current_fitness > best_fitness:
                best_fitness = current_fitness
                best_genome = current_genome.copy()
                self.logger.debug(f"SPSA iteration {iteration}: New best fitness {best_fitness:.4f}")
        
        return best_genome 