"""Evolution engine for OpenClaw trading system.

This module provides the EvolutionEngine class that orchestrates
evolutionary algorithms for trading strategy optimization.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Any, Callable, Dict, Generic, List, Optional, Protocol, TypeVar, Union
from datetime import datetime

import numpy as np
from pydantic import BaseModel, Field

from openclaw.utils.logging import get_logger


class EvolutionAlgorithm(Enum):
    """Available evolution algorithms."""

    GA = "genetic_algorithm"
    GP = "genetic_programming"
    NSGA2 = "nsga2"


class EvolutionStatus(Enum):
    """Status of evolution process."""

    IDLE = auto()
    RUNNING = auto()
    PAUSED = auto()
    CONVERGED = auto()
    MAX_GENERATIONS = auto()
    STOPPED = auto()


@dataclass
class EvolutionConfig:
    """Configuration for evolution process.

    Attributes:
        population_size: Number of individuals in population
        max_generations: Maximum number of generations
        crossover_rate: Probability of crossover (0-1)
        mutation_rate: Probability of mutation (0-1)
        elite_size: Number of elite individuals to preserve
        convergence_threshold: Fitness improvement threshold for convergence
        stagnation_generations: Generations without improvement to trigger convergence
        random_seed: Random seed for reproducibility
    """

    population_size: int = 100
    max_generations: int = 500
    crossover_rate: float = 0.8
    mutation_rate: float = 0.1
    elite_size: int = 5
    convergence_threshold: float = 0.001
    stagnation_generations: int = 50
    random_seed: Optional[int] = None

    def __post_init__(self) -> None:
        """Validate configuration."""
        if self.population_size < 10:
            raise ValueError("Population size must be at least 10")
        if self.max_generations < 1:
            raise ValueError("Max generations must be at least 1")
        if not 0 <= self.crossover_rate <= 1:
            raise ValueError("Crossover rate must be between 0 and 1")
        if not 0 <= self.mutation_rate <= 1:
            raise ValueError("Mutation rate must be between 0 and 1")
        if self.elite_size >= self.population_size:
            raise ValueError("Elite size must be less than population size")


T = TypeVar("T")


class Individual(Protocol):
    """Protocol for evolution individuals."""

    fitness: float
    genes: Any

    def copy(self) -> Individual:
        """Create a deep copy of the individual."""
        ...


class EvolutionCallback(Protocol):
    """Protocol for evolution callback functions."""

    def __call__(
        self,
        generation: int,
        population: List[Individual],
        best_fitness: float,
        avg_fitness: float,
    ) -> None:
        """Called after each generation."""
        ...


class EvolutionMonitor(BaseModel):
    """Monitoring data for evolution process."""

    generation: int = 0
    best_fitness_history: List[float] = Field(default_factory=list)
    avg_fitness_history: List[float] = Field(default_factory=list)
    diversity_history: List[float] = Field(default_factory=list)
    generation_time_ms: List[float] = Field(default_factory=list)
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None
    converged: bool = False
    convergence_generation: Optional[int] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert monitor data to dictionary."""
        return {
            "generation": self.generation,
            "best_fitness_history": self.best_fitness_history,
            "avg_fitness_history": self.avg_fitness_history,
            "diversity_history": self.diversity_history,
            "total_generations": len(self.best_fitness_history),
            "converged": self.converged,
            "convergence_generation": self.convergence_generation,
        }


class EvolutionEngine:
    """Orchestrates evolutionary algorithms for strategy optimization.

    This engine provides a unified interface for running different
    evolutionary algorithms with configurable parameters and monitoring.

    Args:
        config: Evolution configuration
        algorithm: Type of evolution algorithm to use
        fitness_func: Function to evaluate individual fitness
    """

    def __init__(
        self,
        config: EvolutionConfig,
        algorithm: EvolutionAlgorithm,
        fitness_func: Callable[[Any], float],
    ):
        self.config = config
        self.algorithm = algorithm
        self.fitness_func = fitness_func
        self.status = EvolutionStatus.IDLE
        self.monitor = EvolutionMonitor()
        self.population: List[Individual] = []
        self.best_individual: Optional[Individual] = None
        self._callbacks: List[EvolutionCallback] = []
        self._stagnation_counter = 0
        self._last_best_fitness = float("-inf")

        # Set random seed if provided
        if config.random_seed is not None:
            np.random.seed(config.random_seed)

        self.logger = get_logger("evolution.engine")
        self.logger.info(
            f"EvolutionEngine initialized: {algorithm.value}, "
            f"population={config.population_size}, max_gen={config.max_generations}"
        )

    def register_callback(self, callback: EvolutionCallback) -> None:
        """Register a callback for generation updates.

        Args:
            callback: Function to call after each generation
        """
        self._callbacks.append(callback)

    def unregister_callback(self, callback: EvolutionCallback) -> None:
        """Unregister a callback."""
        if callback in self._callbacks:
            self._callbacks.remove(callback)

    def initialize_population(
        self,
        init_func: Callable[[], Individual],
    ) -> None:
        """Initialize the population.

        Args:
            init_func: Function to create a new individual
        """
        self.population = [init_func() for _ in range(self.config.population_size)]
        self._evaluate_population()
        self.logger.info(f"Population initialized with {len(self.population)} individuals")

    def _evaluate_population(self) -> None:
        """Evaluate fitness for all individuals."""
        for individual in self.population:
            individual.fitness = self.fitness_func(individual.genes)

        # Sort by fitness (descending)
        self.population.sort(key=lambda x: x.fitness, reverse=True)
        self.best_individual = self.population[0]

    def _select_parents(self) -> tuple[Individual, Individual]:
        """Select two parents using tournament selection."""
        tournament_size = 3

        def tournament() -> Individual:
            contestants = np.random.choice(self.population, tournament_size, replace=False)
            return max(contestants, key=lambda x: x.fitness)

        return tournament(), tournament()

    def _crossover(
        self,
        parent1: Individual,
        parent2: Individual,
    ) -> tuple[Individual, Individual]:
        """Perform crossover between two parents."""
        if np.random.random() > self.config.crossover_rate:
            return parent1.copy(), parent2.copy()

        # Generic crossover - override in specialized classes
        child1 = parent1.copy()
        child2 = parent2.copy()
        return child1, child2

    def _mutate(self, individual: Individual) -> Individual:
        """Mutate an individual."""
        if np.random.random() > self.config.mutation_rate:
            return individual

        # Generic mutation - override in specialized classes
        return individual.copy()

    def _create_next_generation(self) -> None:
        """Create the next generation."""
        new_population: List[Individual] = []

        # Elitism: keep best individuals
        elite = [ind.copy() for ind in self.population[: self.config.elite_size]]
        new_population.extend(elite)

        # Create offspring
        while len(new_population) < self.config.population_size:
            parent1, parent2 = self._select_parents()
            child1, child2 = self._crossover(parent1, parent2)
            child1 = self._mutate(child1)
            child2 = self._mutate(child2)

            new_population.append(child1)
            if len(new_population) < self.config.population_size:
                new_population.append(child2)

        self.population = new_population
        self._evaluate_population()

    def _calculate_diversity(self) -> float:
        """Calculate population diversity as average pairwise distance."""
        if len(self.population) < 2:
            return 0.0

        # Simple diversity: standard deviation of fitness values
        fitness_values = [ind.fitness for ind in self.population]
        return float(np.std(fitness_values))

    def _check_convergence(self) -> bool:
        """Check if evolution has converged."""
        if not self.population:
            return False

        current_best = self.population[0].fitness
        improvement = current_best - self._last_best_fitness

        if improvement < self.config.convergence_threshold:
            self._stagnation_counter += 1
        else:
            self._stagnation_counter = 0
            self._last_best_fitness = current_best

        return self._stagnation_counter >= self.config.stagnation_generations

    def _trigger_callbacks(self) -> None:
        """Trigger registered callbacks."""
        if not self.population:
            return

        best_fitness = self.population[0].fitness
        avg_fitness = np.mean([ind.fitness for ind in self.population])

        for callback in self._callbacks:
            try:
                callback(
                    generation=self.monitor.generation,
                    population=self.population,
                    best_fitness=best_fitness,
                    avg_fitness=avg_fitness,
                )
            except Exception as e:
                self.logger.error(f"Callback error: {e}")

    def run(self) -> EvolutionMonitor:
        """Run the evolution process.

        Returns:
            EvolutionMonitor with complete evolution history
        """
        if not self.population:
            raise ValueError("Population not initialized. Call initialize_population() first.")

        self.status = EvolutionStatus.RUNNING
        self.monitor.start_time = datetime.now()
        self.logger.info("Evolution started")

        for generation in range(self.config.max_generations):
            self.monitor.generation = generation

            # Create next generation
            self._create_next_generation()

            # Record statistics
            best_fitness = self.population[0].fitness
            avg_fitness = float(np.mean([ind.fitness for ind in self.population]))
            diversity = self._calculate_diversity()

            self.monitor.best_fitness_history.append(best_fitness)
            self.monitor.avg_fitness_history.append(avg_fitness)
            self.monitor.diversity_history.append(diversity)

            # Trigger callbacks
            self._trigger_callbacks()

            # Log progress
            if generation % 50 == 0 or generation < 5:
                self.logger.info(
                    f"Generation {generation}: best={best_fitness:.4f}, "
                    f"avg={avg_fitness:.4f}, diversity={diversity:.4f}"
                )

            # Check convergence
            if self._check_convergence():
                self.status = EvolutionStatus.CONVERGED
                self.monitor.converged = True
                self.monitor.convergence_generation = generation
                self.logger.info(f"Converged at generation {generation}")
                break

        else:
            self.status = EvolutionStatus.MAX_GENERATIONS
            self.logger.info(f"Reached max generations ({self.config.max_generations})")

        self.monitor.end_time = datetime.now()
        self.status = EvolutionStatus.STOPPED if self.status == EvolutionStatus.RUNNING else self.status

        return self.monitor

    def get_best_individual(self) -> Optional[Individual]:
        """Get the best individual from the current population."""
        return self.best_individual

    def get_population_stats(self) -> Dict[str, float]:
        """Get statistics about the current population."""
        if not self.population:
            return {}

        fitness_values = [ind.fitness for ind in self.population]
        return {
            "best": float(np.max(fitness_values)),
            "worst": float(np.min(fitness_values)),
            "mean": float(np.mean(fitness_values)),
            "median": float(np.median(fitness_values)),
            "std": float(np.std(fitness_values)),
        }

    def reset(self) -> None:
        """Reset the engine to initial state."""
        self.status = EvolutionStatus.IDLE
        self.population = []
        self.best_individual = None
        self.monitor = EvolutionMonitor()
        self._stagnation_counter = 0
        self._last_best_fitness = float("-inf")
        self.logger.info("EvolutionEngine reset")
