"""Genetic Algorithm implementation for OpenClaw trading system.

This module provides GA capabilities for optimizing trading strategy parameters
with various selection, crossover, and mutation operators.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple

import numpy as np

from openclaw.utils.logging import get_logger


class SelectionOperator(Enum):
    """Available selection operators."""

    ROULETTE = auto()
    TOURNAMENT = auto()
    RANK = auto()
    STOCHASTIC_UNIVERSAL = auto()


class CrossoverOperator(Enum):
    """Available crossover operators."""

    SINGLE_POINT = auto()
    TWO_POINT = auto()
    UNIFORM = auto()
    ARITHMETIC = auto()


class MutationOperator(Enum):
    """Available mutation operators."""

    GAUSSIAN = auto()
    UNIFORM = auto()
    BOUNDARY = auto()
    POLYNOMIAL = auto()


@dataclass
class Chromosome:
    """Chromosome for genetic algorithm.

    Attributes:
        genes: Array of gene values
        fitness: Fitness value (higher is better)
        generation: Generation when created
    """

    genes: np.ndarray
    fitness: float = 0.0
    generation: int = 0

    def copy(self) -> Chromosome:
        """Create a deep copy of the chromosome."""
        return Chromosome(
            genes=self.genes.copy(),
            fitness=self.fitness,
            generation=self.generation,
        )

    def __len__(self) -> int:
        """Return the number of genes."""
        return len(self.genes)


@dataclass
class GAConfig:
    """Configuration for Genetic Algorithm.

    Attributes:
        population_size: Number of chromosomes in population
        max_generations: Maximum number of generations
        crossover_rate: Probability of crossover (0-1)
        mutation_rate: Probability of mutation per gene (0-1)
        mutation_sigma: Standard deviation for Gaussian mutation
        elite_size: Number of elite chromosomes to preserve
        selection: Selection operator to use
        crossover: Crossover operator to use
        mutation: Mutation operator to use
        bounds: Tuple of (min, max) bounds for gene values
    """

    population_size: int = 100
    max_generations: int = 500
    crossover_rate: float = 0.8
    mutation_rate: float = 0.1
    mutation_sigma: float = 0.1
    elite_size: int = 5
    selection: SelectionOperator = SelectionOperator.TOURNAMENT
    crossover: CrossoverOperator = CrossoverOperator.UNIFORM
    mutation: MutationOperator = MutationOperator.GAUSSIAN
    bounds: Optional[Tuple[float, float]] = None

    def __post_init__(self) -> None:
        """Validate configuration."""
        if self.population_size < 10:
            raise ValueError("Population size must be at least 10")
        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")


class GeneticAlgorithm:
    """Genetic Algorithm for trading strategy optimization.

    This class implements a standard GA with various operators for
    optimizing trading strategy parameters.

    Args:
        config: GA configuration
        fitness_func: Function to evaluate chromosome fitness
        gene_init_func: Function to initialize a new gene array
    """

    def __init__(
        self,
        config: GAConfig,
        fitness_func: Callable[[np.ndarray], float],
        gene_init_func: Callable[[], np.ndarray],
    ):
        self.config = config
        self.fitness_func = fitness_func
        self.gene_init_func = gene_init_func
        self.population: List[Chromosome] = []
        self.generation = 0
        self.best_chromosome: Optional[Chromosome] = None
        self.fitness_history: List[float] = []

        # Set random seed for reproducibility
        np.random.seed()

        self.logger = get_logger("evolution.ga")
        self.logger.info(
            f"GeneticAlgorithm initialized: pop_size={config.population_size}, "
            f"selection={config.selection.name}, crossover={config.crossover.name}"
        )

    def initialize(self) -> None:
        """Initialize the population with random chromosomes."""
        self.population = []
        for _ in range(self.config.population_size):
            genes = self.gene_init_func()
            chromosome = Chromosome(genes=genes, generation=0)
            chromosome.fitness = self.fitness_func(genes)
            self.population.append(chromosome)

        self._sort_population()
        self.best_chromosome = self.population[0]
        self.logger.info(f"Population initialized with {len(self.population)} chromosomes")

    def _sort_population(self) -> None:
        """Sort population by fitness (descending)."""
        self.population.sort(key=lambda x: x.fitness, reverse=True)

    def _select_roulette(self) -> Chromosome:
        """Select a chromosome using roulette wheel selection."""
        fitnesses = np.array([c.fitness for c in self.population])
        min_fitness = fitnesses.min()

        # Shift fitnesses to be positive
        if min_fitness < 0:
            fitnesses = fitnesses - min_fitness + 1e-10

        total = fitnesses.sum()
        if total == 0:
            return np.random.choice(self.population)

        probabilities = fitnesses / total
        idx = np.random.choice(len(self.population), p=probabilities)
        return self.population[idx]

    def _select_tournament(self, tournament_size: int = 3) -> Chromosome:
        """Select a chromosome using tournament selection."""
        contestants = np.random.choice(self.population, tournament_size, replace=False)
        return max(contestants, key=lambda x: x.fitness)

    def _select_rank(self) -> Chromosome:
        """Select a chromosome using rank-based selection."""
        n = len(self.population)
        ranks = np.arange(n, 0, -1)  # Higher rank = better fitness
        total_rank = ranks.sum()
        probabilities = ranks / total_rank
        idx = np.random.choice(n, p=probabilities)
        return self.population[idx]

    def _select_stochastic_universal(self, num_selections: int) -> List[Chromosome]:
        """Select multiple chromosomes using Stochastic Universal Sampling."""
        fitnesses = np.array([c.fitness for c in self.population])
        min_fitness = fitnesses.min()

        if min_fitness < 0:
            fitnesses = fitnesses - min_fitness + 1e-10

        total = fitnesses.sum()
        if total == 0:
            return list(np.random.choice(self.population, num_selections, replace=True))

        probabilities = fitnesses / total
        pointers = np.linspace(0, 1, num_selections, endpoint=False)
        pointers += np.random.uniform(0, 1 / num_selections)

        selected = []
        cumsum = np.cumsum(probabilities)
        for ptr in pointers:
            idx = np.searchsorted(cumsum, ptr)
            selected.append(self.population[idx])

        return selected

    def _select_parent(self) -> Chromosome:
        """Select a parent using the configured selection operator."""
        if self.config.selection == SelectionOperator.ROULETTE:
            return self._select_roulette()
        elif self.config.selection == SelectionOperator.TOURNAMENT:
            return self._select_tournament()
        elif self.config.selection == SelectionOperator.RANK:
            return self._select_rank()
        else:
            # Default to tournament
            return self._select_tournament()

    def _crossover_single_point(
        self, parent1: Chromosome, parent2: Chromosome
    ) -> Tuple[Chromosome, Chromosome]:
        """Perform single-point crossover."""
        point = np.random.randint(1, len(parent1.genes))
        child1_genes = np.concatenate([parent1.genes[:point], parent2.genes[point:]])
        child2_genes = np.concatenate([parent2.genes[:point], parent1.genes[point:]])

        return (
            Chromosome(genes=child1_genes, generation=self.generation),
            Chromosome(genes=child2_genes, generation=self.generation),
        )

    def _crossover_two_point(
        self, parent1: Chromosome, parent2: Chromosome
    ) -> Tuple[Chromosome, Chromosome]:
        """Perform two-point crossover."""
        points = sorted(np.random.choice(len(parent1.genes) - 1, 2, replace=False) + 1)
        p1, p2 = points

        child1_genes = np.concatenate([
            parent1.genes[:p1],
            parent2.genes[p1:p2],
            parent1.genes[p2:],
        ])
        child2_genes = np.concatenate([
            parent2.genes[:p1],
            parent1.genes[p1:p2],
            parent2.genes[p2:],
        ])

        return (
            Chromosome(genes=child1_genes, generation=self.generation),
            Chromosome(genes=child2_genes, generation=self.generation),
        )

    def _crossover_uniform(
        self, parent1: Chromosome, parent2: Chromosome
    ) -> Tuple[Chromosome, Chromosome]:
        """Perform uniform crossover."""
        mask = np.random.random(len(parent1.genes)) < 0.5
        child1_genes = np.where(mask, parent1.genes, parent2.genes)
        child2_genes = np.where(mask, parent2.genes, parent1.genes)

        return (
            Chromosome(genes=child1_genes, generation=self.generation),
            Chromosome(genes=child2_genes, generation=self.generation),
        )

    def _crossover_arithmetic(
        self, parent1: Chromosome, parent2: Chromosome
    ) -> Tuple[Chromosome, Chromosome]:
        """Perform arithmetic crossover."""
        alpha = np.random.random()
        child1_genes = alpha * parent1.genes + (1 - alpha) * parent2.genes
        child2_genes = alpha * parent2.genes + (1 - alpha) * parent1.genes

        return (
            Chromosome(genes=child1_genes, generation=self.generation),
            Chromosome(genes=child2_genes, generation=self.generation),
        )

    def _crossover(
        self, parent1: Chromosome, parent2: Chromosome
    ) -> Tuple[Chromosome, Chromosome]:
        """Perform crossover using the configured operator."""
        if np.random.random() > self.config.crossover_rate:
            return parent1.copy(), parent2.copy()

        if self.config.crossover == CrossoverOperator.SINGLE_POINT:
            return self._crossover_single_point(parent1, parent2)
        elif self.config.crossover == CrossoverOperator.TWO_POINT:
            return self._crossover_two_point(parent1, parent2)
        elif self.config.crossover == CrossoverOperator.UNIFORM:
            return self._crossover_uniform(parent1, parent2)
        elif self.config.crossover == CrossoverOperator.ARITHMETIC:
            return self._crossover_arithmetic(parent1, parent2)
        else:
            return self._crossover_uniform(parent1, parent2)

    def _mutate_gaussian(self, chromosome: Chromosome) -> Chromosome:
        """Apply Gaussian mutation."""
        mask = np.random.random(len(chromosome.genes)) < self.config.mutation_rate
        noise = np.random.normal(0, self.config.mutation_sigma, len(chromosome.genes))
        mutated = chromosome.genes.copy()
        mutated[mask] += noise[mask]

        # Apply bounds if configured
        if self.config.bounds:
            min_val, max_val = self.config.bounds
            mutated = np.clip(mutated, min_val, max_val)

        return Chromosome(genes=mutated, generation=self.generation)

    def _mutate_uniform(self, chromosome: Chromosome) -> Chromosome:
        """Apply uniform mutation."""
        if self.config.bounds is None:
            return chromosome.copy()

        mask = np.random.random(len(chromosome.genes)) < self.config.mutation_rate
        mutated = chromosome.genes.copy()
        min_val, max_val = self.config.bounds
        mutated[mask] = np.random.uniform(min_val, max_val, mask.sum())

        return Chromosome(genes=mutated, generation=self.generation)

    def _mutate_boundary(self, chromosome: Chromosome) -> Chromosome:
        """Apply boundary mutation."""
        if self.config.bounds is None:
            return chromosome.copy()

        mask = np.random.random(len(chromosome.genes)) < self.config.mutation_rate
        mutated = chromosome.genes.copy()
        min_val, max_val = self.config.bounds

        for i in range(len(mutated)):
            if mask[i]:
                mutated[i] = min_val if np.random.random() < 0.5 else max_val

        return Chromosome(genes=mutated, generation=self.generation)

    def _mutate_polynomial(self, chromosome: Chromosome, eta: float = 20.0) -> Chromosome:
        """Apply polynomial mutation."""
        if self.config.bounds is None:
            return self._mutate_gaussian(chromosome)

        mask = np.random.random(len(chromosome.genes)) < self.config.mutation_rate
        mutated = chromosome.genes.copy()
        min_val, max_val = self.config.bounds

        for i in range(len(mutated)):
            if mask[i]:
                x = mutated[i]
                delta1 = (x - min_val) / (max_val - min_val)
                delta2 = (max_val - x) / (max_val - min_val)
                rand = np.random.random()
                mut_pow = 1.0 / (eta + 1.0)

                if rand <= 0.5:
                    xy = 1.0 - delta1
                    val = 2.0 * rand + (1.0 - 2.0 * rand) * (xy ** (eta + 1))
                    delta_q = val ** mut_pow - 1.0
                else:
                    xy = 1.0 - delta2
                    val = 2.0 * (1.0 - rand) + 2.0 * (rand - 0.5) * (xy ** (eta + 1))
                    delta_q = 1.0 - val ** mut_pow

                mutated[i] = x + delta_q * (max_val - min_val)
                mutated[i] = np.clip(mutated[i], min_val, max_val)

        return Chromosome(genes=mutated, generation=self.generation)

    def _mutate(self, chromosome: Chromosome) -> Chromosome:
        """Apply mutation using the configured operator."""
        if self.config.mutation == MutationOperator.GAUSSIAN:
            return self._mutate_gaussian(chromosome)
        elif self.config.mutation == MutationOperator.UNIFORM:
            return self._mutate_uniform(chromosome)
        elif self.config.mutation == MutationOperator.BOUNDARY:
            return self._mutate_boundary(chromosome)
        elif self.config.mutation == MutationOperator.POLYNOMIAL:
            return self._mutate_polynomial(chromosome)
        else:
            return self._mutate_gaussian(chromosome)

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

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

        # Create offspring
        while len(new_population) < self.config.population_size:
            parent1 = self._select_parent()
            parent2 = self._select_parent()

            child1, child2 = self._crossover(parent1, parent2)
            child1 = self._mutate(child1)
            child2 = self._mutate(child2)

            # Evaluate fitness
            child1.fitness = self.fitness_func(child1.genes)
            child2.fitness = self.fitness_func(child2.genes)

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

        self.population = new_population
        self._sort_population()
        self.best_chromosome = self.population[0]

    def step(self) -> float:
        """Execute one generation step.

        Returns:
            Best fitness value after the step
        """
        if not self.population:
            raise ValueError("Population not initialized. Call initialize() first.")

        self.generation += 1
        self._create_next_generation()

        best_fitness = self.population[0].fitness
        self.fitness_history.append(best_fitness)

        return best_fitness

    def run(self, callback: Optional[Callable[[int, float], None]] = None) -> Chromosome:
        """Run the genetic algorithm.

        Args:
            callback: Optional callback function (generation, best_fitness) -> None

        Returns:
            Best chromosome found
        """
        if not self.population:
            self.initialize()

        self.logger.info(f"GA started for {self.config.max_generations} generations")

        for gen in range(self.config.max_generations):
            best_fitness = self.step()

            if callback:
                callback(gen, best_fitness)

            if gen % 50 == 0 or gen < 5:
                self.logger.info(f"Generation {gen}: best_fitness={best_fitness:.4f}")

        self.logger.info(
            f"GA completed: best_fitness={self.best_chromosome.fitness:.4f}"
        )
        return self.best_chromosome

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

        fitnesses = [c.fitness for c in self.population]
        return {
            "generation": self.generation,
            "best": float(np.max(fitnesses)),
            "worst": float(np.min(fitnesses)),
            "mean": float(np.mean(fitnesses)),
            "median": float(np.median(fitnesses)),
            "std": float(np.std(fitnesses)),
            "diversity": float(np.std([c.genes.std() for c in self.population])),
        }
