"""NSGA-II (Non-dominated Sorting Genetic Algorithm II) implementation.

This module provides NSGA-II for multi-objective optimization of trading strategies,
handling objectives like profit, risk, and drawdown simultaneously.
"""

from __future__ import annotations

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

import numpy as np

from openclaw.utils.logging import get_logger


class DominanceRelation(Enum):
    """Dominance relations between solutions."""

    DOMINATES = 1
    DOMINATED = 2
    NON_DOMINATED = 3
    IDENTICAL = 4


@dataclass
class ObjectiveValue:
    """Value of a single objective.

    Attributes:
        name: Objective name
        value: Objective value
        minimize: Whether to minimize (True) or maximize (False)
        weight: Weight for this objective
    """

    name: str
    value: float
    minimize: bool = False
    weight: float = 1.0

    def normalized_value(self) -> float:
        """Get normalized value (higher is always better)."""
        return -self.value if self.minimize else self.value


@dataclass
class Individual:
    """Individual for NSGA-II multi-objective optimization.

    Attributes:
        genes: Chromosome representation
        objectives: Dictionary of objective values
        rank: Non-domination rank
        crowding_distance: Crowding distance for diversity
        domination_count: Number of individuals dominating this one
        dominated_solutions: List of individuals this one dominates
    """

    genes: np.ndarray
    objectives: Dict[str, ObjectiveValue] = field(default_factory=dict)
    rank: int = 0
    crowding_distance: float = 0.0
    domination_count: int = 0
    dominated_solutions: List["Individual"] = field(default_factory=list)

    def copy(self) -> Individual:
        """Create a deep copy of the individual."""
        new_individual = Individual(
            genes=self.genes.copy(),
            objectives={k: copy.copy(v) for k, v in self.objectives.items()},
            rank=self.rank,
            crowding_distance=self.crowding_distance,
        )
        return new_individual

    def dominates(self, other: Individual) -> bool:
        """Check if this individual dominates another.

        A dominates B if A is no worse than B in all objectives
        and strictly better in at least one.

        Args:
            other: Individual to compare against

        Returns:
            True if this individual dominates other
        """
        if not self.objectives or not other.objectives:
            return False

        strictly_better = False
        for key, obj in self.objectives.items():
            if key not in other.objectives:
                return False

            other_obj = other.objectives[key]
            self_val = obj.normalized_value()
            other_val = other_obj.normalized_value()

            if self_val < other_val:
                return False
            if self_val > other_val:
                strictly_better = True

        return strictly_better

    def get_dominance_relation(self, other: Individual) -> DominanceRelation:
        """Get the dominance relation between two individuals.

        Args:
            other: Individual to compare against

        Returns:
            Dominance relation
        """
        if self.dominates(other):
            return DominanceRelation.DOMINATES
        if other.dominates(self):
            return DominanceRelation.DOMINATED

        # Check if identical
        if len(self.objectives) != len(other.objectives):
            return DominanceRelation.NON_DOMINATED

        identical = all(
            abs(self.objectives[k].normalized_value() - other.objectives[k].normalized_value()) < 1e-10
            for k in self.objectives
            if k in other.objectives
        )

        if identical:
            return DominanceRelation.IDENTICAL

        return DominanceRelation.NON_DOMINATED

    def get_objective_vector(self) -> np.ndarray:
        """Get objective values as numpy array (normalized)."""
        return np.array([obj.normalized_value() for obj in self.objectives.values()])

    def __eq__(self, other: object) -> bool:
        """Check equality with another Individual by comparing genes."""
        if not isinstance(other, Individual):
            return False
        return np.array_equal(self.genes, other.genes)

    def __hash__(self) -> int:
        """Hash based on genes for use in sets/dicts."""
        return hash(self.genes.tobytes())


@dataclass
class ParetoFront:
    """A Pareto front (set of non-dominated solutions).

    Attributes:
        individuals: List of individuals in this front
        rank: Front rank (0 is best)
    """

    individuals: List[Individual] = field(default_factory=list)
    rank: int = 0

    def __len__(self) -> int:
        return len(self.individuals)

    def get_crowding_distances(self) -> List[float]:
        """Get crowding distances for all individuals."""
        return [ind.crowding_distance for ind in self.individuals]


@dataclass
class NSGA2Config:
    """Configuration for NSGA-II.

    Attributes:
        population_size: Number of individuals in population
        max_generations: Maximum number of generations
        crossover_rate: Probability of crossover
        mutation_rate: Probability of mutation per gene
        mutation_sigma: Standard deviation for Gaussian mutation
        tournament_size: Size of tournament for selection
        num_objectives: Number of objectives to optimize
    """

    population_size: int = 100
    max_generations: int = 500
    crossover_rate: float = 0.9
    mutation_rate: float = 0.1
    mutation_sigma: float = 0.1
    tournament_size: int = 2
    num_objectives: int = 2

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


class NSGA2:
    """NSGA-II multi-objective optimization algorithm.

    This class implements NSGA-II for optimizing trading strategies
    with multiple conflicting objectives (e.g., profit vs risk).

    Args:
        config: NSGA-II configuration
        objective_funcs: List of objective evaluation functions
        gene_init_func: Function to initialize genes
        bounds: Optional bounds for gene values
    """

    def __init__(
        self,
        config: NSGA2Config,
        objective_funcs: Dict[str, Callable[[np.ndarray], float]],
        gene_init_func: Callable[[], np.ndarray],
        bounds: Optional[Tuple[np.ndarray, np.ndarray]] = None,
    ):
        self.config = config
        self.objective_funcs = objective_funcs
        self.gene_init_func = gene_init_func
        self.bounds = bounds
        self.population: List[Individual] = []
        self.generation = 0
        self.pareto_fronts: List[ParetoFront] = []

        self.logger = get_logger("evolution.nsga2")
        self.logger.info(
            f"NSGA2 initialized: pop_size={config.population_size}, "
            f"objectives={list(objective_funcs.keys())}"
        )

    def initialize(self) -> None:
        """Initialize population with random individuals."""
        self.population = []
        for _ in range(self.config.population_size):
            genes = self.gene_init_func()
            individual = Individual(genes=genes)
            self._evaluate_objectives(individual)
            self.population.append(individual)

        self.logger.info(f"Population initialized with {len(self.population)} individuals")

    def _evaluate_objectives(self, individual: Individual) -> None:
        """Evaluate all objectives for an individual."""
        for name, func in self.objective_funcs.items():
            value = func(individual.genes)
            individual.objectives[name] = ObjectiveValue(name=name, value=value)

    def _fast_non_dominated_sort(self) -> List[ParetoFront]:
        """Perform fast non-dominated sorting.

        Returns:
            List of Pareto fronts in order of rank
        """
        fronts: List[ParetoFront] = []

        # Reset domination counts
        for p in self.population:
            p.domination_count = 0
            p.dominated_solutions = []

        # First front
        first_front = ParetoFront(rank=0)

        for i, p in enumerate(self.population):
            for j, q in enumerate(self.population):
                if i == j:
                    continue

                if p.dominates(q):
                    if q not in p.dominated_solutions:
                        p.dominated_solutions.append(q)
                elif q.dominates(p):
                    p.domination_count += 1

            if p.domination_count == 0:
                p.rank = 0
                first_front.individuals.append(p)

        fronts.append(first_front)

        # Subsequent fronts
        i = 0
        while i < len(fronts) and fronts[i].individuals:
            next_front = ParetoFront(rank=i + 1)

            for p in fronts[i].individuals:
                for q in p.dominated_solutions:
                    q.domination_count -= 1
                    if q.domination_count == 0:
                        q.rank = i + 1
                        next_front.individuals.append(q)

            if next_front.individuals:
                fronts.append(next_front)
            i += 1

        return fronts

    def _calculate_crowding_distance(self, front: ParetoFront) -> None:
        """Calculate crowding distance for individuals in a front.

        Args:
            front: Pareto front to calculate distances for
        """
        if len(front.individuals) <= 2:
            for ind in front.individuals:
                ind.crowding_distance = float("inf")
            return

        # Initialize distances
        for ind in front.individuals:
            ind.crowding_distance = 0.0

        # Get objective names
        if not front.individuals:
            return

        objective_names = list(front.individuals[0].objectives.keys())

        for obj_name in objective_names:
            # Sort by this objective
            front.individuals.sort(key=lambda x: x.objectives[obj_name].normalized_value())

            # Set boundary distances to infinity
            front.individuals[0].crowding_distance = float("inf")
            front.individuals[-1].crowding_distance = float("inf")

            # Get min and max values
            min_val = front.individuals[0].objectives[obj_name].normalized_value()
            max_val = front.individuals[-1].objectives[obj_name].normalized_value()
            obj_range = max_val - min_val

            if obj_range == 0:
                continue

            # Calculate distances for intermediate individuals
            for i in range(1, len(front.individuals) - 1):
                prev_val = front.individuals[i - 1].objectives[obj_name].normalized_value()
                next_val = front.individuals[i + 1].objectives[obj_name].normalized_value()
                front.individuals[i].crowding_distance += (next_val - prev_val) / obj_range

    def _crowded_comparison_operator(self, ind1: Individual, ind2: Individual) -> bool:
        """Crowded comparison operator for sorting.

        Returns:
            True if ind1 is better than ind2
        """
        if ind1.rank < ind2.rank:
            return True
        if ind1.rank == ind2.rank and ind1.crowding_distance > ind2.crowding_distance:
            return True
        return False

    def _make_new_population(self) -> List[Individual]:
        """Create new population using selection, crossover, and mutation."""
        new_population: List[Individual] = []

        while len(new_population) < self.config.population_size:
            # Tournament selection
            parent1 = self._tournament_select()
            parent2 = self._tournament_select()

            # Crossover
            child1, child2 = self._crossover(parent1, parent2)

            # Mutation
            child1 = self._mutate(child1)
            child2 = self._mutate(child2)

            # Evaluate objectives
            self._evaluate_objectives(child1)
            self._evaluate_objectives(child2)

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

        return new_population

    def _tournament_select(self) -> Individual:
        """Select individual using binary tournament."""
        candidates = np.random.choice(self.population, self.config.tournament_size, replace=False)
        best = candidates[0]

        for candidate in candidates[1:]:
            if self._crowded_comparison_operator(candidate, best):
                best = candidate

        return best

    def _crossover(self, parent1: Individual, parent2: Individual) -> Tuple[Individual, Individual]:
        """Simulated binary crossover (SBX)."""
        if np.random.random() > self.config.crossover_rate:
            return parent1.copy(), parent2.copy()

        child1 = parent1.copy()
        child2 = parent2.copy()

        # Uniform crossover for real-valued genes
        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)

        # Apply bounds if specified
        if self.bounds:
            min_bounds, max_bounds = self.bounds
            child1.genes = np.clip(child1.genes, min_bounds, max_bounds)
            child2.genes = np.clip(child2.genes, min_bounds, max_bounds)

        return child1, child2

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

        mutated = individual.copy()

        # Gaussian mutation
        mask = np.random.random(len(mutated.genes)) < self.config.mutation_rate
        noise = np.random.normal(0, self.config.mutation_sigma, len(mutated.genes))
        mutated.genes[mask] += noise[mask]

        # Apply bounds if specified
        if self.bounds:
            min_bounds, max_bounds = self.bounds
            mutated.genes = np.clip(mutated.genes, min_bounds, max_bounds)

        return mutated

    def step(self) -> List[ParetoFront]:
        """Execute one generation step.

        Returns:
            List of Pareto fronts
        """
        if not self.population:
            raise ValueError("Population not initialized. Call initialize() first.")

        self.generation += 1

        # Create offspring population
        offspring = self._make_new_population()

        # Combine parent and offspring
        combined = self.population + offspring

        # Non-dominated sorting
        fronts = self._fast_non_dominated_sort_on_population(combined)

        # Calculate crowding distances
        for front in fronts:
            self._calculate_crowding_distance(front)

        # Select next generation
        new_population: List[Individual] = []
        front_idx = 0

        while len(new_population) + len(fronts[front_idx].individuals) <= self.config.population_size:
            new_population.extend(fronts[front_idx].individuals)
            front_idx += 1
            if front_idx >= len(fronts):
                break

        # Fill remaining slots from next front using crowding distance
        if front_idx < len(fronts) and len(new_population) < self.config.population_size:
            remaining = self.config.population_size - len(new_population)
            last_front = fronts[front_idx]

            # Sort by crowding distance (descending)
            last_front.individuals.sort(key=lambda x: x.crowding_distance, reverse=True)
            new_population.extend(last_front.individuals[:remaining])

        self.population = new_population
        self.pareto_fronts = self._fast_non_dominated_sort()

        return self.pareto_fronts

    def _fast_non_dominated_sort_on_population(
        self, population: List[Individual]
    ) -> List[ParetoFront]:
        """Perform fast non-dominated sorting on a given population."""
        # Reset domination counts
        for p in population:
            p.domination_count = 0
            p.dominated_solutions = []

        fronts: List[ParetoFront] = []
        first_front = ParetoFront(rank=0)

        for i, p in enumerate(population):
            for j, q in enumerate(population):
                if i == j:
                    continue

                if p.dominates(q):
                    if q not in p.dominated_solutions:
                        p.dominated_solutions.append(q)
                elif q.dominates(p):
                    p.domination_count += 1

            if p.domination_count == 0:
                p.rank = 0
                first_front.individuals.append(p)

        fronts.append(first_front)

        i = 0
        while i < len(fronts) and fronts[i].individuals:
            next_front = ParetoFront(rank=i + 1)

            for p in fronts[i].individuals:
                for q in p.dominated_solutions:
                    q.domination_count -= 1
                    if q.domination_count == 0:
                        q.rank = i + 1
                        next_front.individuals.append(q)

            if next_front.individuals:
                fronts.append(next_front)
            i += 1

        return fronts

    def run(
        self,
        callback: Optional[Callable[[int, List[ParetoFront]], None]] = None,
    ) -> List[ParetoFront]:
        """Run the NSGA-II algorithm.

        Args:
            callback: Optional callback (generation, pareto_fronts) -> None

        Returns:
            List of Pareto fronts (first front is optimal)
        """
        if not self.population:
            self.initialize()

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

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

            if callback:
                callback(gen, fronts)

            if gen % 50 == 0 or gen < 5:
                first_front_size = len(fronts[0]) if fronts else 0
                self.logger.info(
                    f"Generation {gen}: {len(fronts)} fronts, "
                    f"first front size={first_front_size}"
                )

        self.logger.info(
            f"NSGA2 completed: {len(self.pareto_fronts)} fronts, "
            f"first front has {len(self.pareto_fronts[0]) if self.pareto_fronts else 0} solutions"
        )

        return self.pareto_fronts

    def get_pareto_front_solutions(self) -> List[Individual]:
        """Get solutions from the first Pareto front."""
        if not self.pareto_fronts:
            self.pareto_fronts = self._fast_non_dominated_sort()

        return self.pareto_fronts[0].individuals if self.pareto_fronts else []

    def get_hypervolume(self, reference_point: Optional[np.ndarray] = None) -> float:
        """Calculate hypervolume indicator (simplified for 2D).

        Args:
            reference_point: Reference point for hypervolume calculation

        Returns:
            Hypervolume value
        """
        solutions = self.get_pareto_front_solutions()
        if not solutions:
            return 0.0

        objective_vectors = np.array([s.get_objective_vector() for s in solutions])

        if reference_point is None:
            # Use worst values as reference
            reference_point = objective_vectors.min(axis=0) - 1.0

        if len(objective_vectors[0]) == 2:
            # 2D hypervolume calculation
            sorted_indices = np.argsort(objective_vectors[:, 0])
            sorted_points = objective_vectors[sorted_indices]

            volume = 0.0
            for i, point in enumerate(sorted_points):
                if i == 0:
                    width = point[0] - reference_point[0]
                else:
                    width = point[0] - sorted_points[i - 1][0]

                height = point[1] - reference_point[1]
                volume += width * height

            return volume

        # For higher dimensions, return 0 (would need proper HV algorithm)
        return 0.0

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

        stats = {
            "generation": self.generation,
            "population_size": len(self.population),
            "num_fronts": len(self.pareto_fronts),
        }

        if self.pareto_fronts:
            first_front = self.pareto_fronts[0]
            stats["pareto_front_size"] = len(first_front)

            # Objective statistics on Pareto front
            for obj_name in self.objective_funcs.keys():
                values = [
                    ind.objectives[obj_name].value
                    for ind in first_front.individuals
                    if obj_name in ind.objectives
                ]
                if values:
                    stats[f"pareto_{obj_name}_mean"] = float(np.mean(values))
                    stats[f"pareto_{obj_name}_std"] = float(np.std(values))

        return stats
