"""Evolutionary strategy optimizer.

Production-grade evolutionary optimization engine inspired by the Darwin SDK
and the ProFiT framework (Siper et al., "Program Search for Financial Trading").

Features:
- 6 selection methods (Tournament, Uniform, Roulette, Rank, UCB1, Epsilon-Greedy)
- NSGA-II multi-objective optimization with Pareto fronts
- Island model with Ring and FullyConnected migration topologies
- Convergence detection with patience and diversity collapse
- Adaptive mutation rates via Rechenberg's 1/5 success rule
- Walk-forward validation to prevent overfitting
- Per-generation statistics tracking
"""

from __future__ import annotations

import logging
from typing import Any, Callable

from horizon._horizon import (
    EvolutionConfig,
    EvolutionResult,
    GenerationStats,
    Individual,
    IslandConfig,
    MigrationTopology,
    NsgaIndividual,
    NsgaResult,
    ParamBound,
    SelectionMethod,
    ConvergenceState,
    AdaptiveMutationState,
    ObjectiveDirection,
    # Selection functions
    tournament_select,
    uniform_select,
    roulette_select,
    rank_select,
    ucb1_select,
    epsilon_greedy_select,
    # Genetic operators
    gaussian_mutate,
    uniform_crossover,
    init_population,
    # NSGA-II
    non_dominated_sort,
    crowding_distance_assign,
    nsga2_select,
    pareto_front,
    # Island model
    island_migrate,
    island_selection_methods,
    # Convergence
    check_convergence,
    # Adaptive mutation
    adapt_mutation_sigma,
    # Diversity & stats
    population_diversity,
    compute_generation_stats,
)

logger = logging.getLogger("horizon")


def _select_parents(
    fitnesses: list[float],
    method: SelectionMethod,
    count: int,
    seed: int,
    tournament_size: int = 3,
    selection_counts: list[int] | None = None,
    total_selections: int = 0,
    epsilon: float = 0.1,
) -> list[int]:
    """Select parent indices using the specified selection method."""
    if method == SelectionMethod.Tournament:
        return tournament_select(fitnesses, tournament_size, count, seed)
    elif method == SelectionMethod.Uniform:
        return uniform_select(fitnesses, count, seed)
    elif method == SelectionMethod.RouletteWheel:
        return roulette_select(fitnesses, count, seed)
    elif method == SelectionMethod.Rank:
        return rank_select(fitnesses, count, seed)
    elif method == SelectionMethod.Ucb1:
        counts = selection_counts or [0] * len(fitnesses)
        return ucb1_select(fitnesses, counts, total_selections, count)
    elif method == SelectionMethod.EpsilonGreedy:
        return epsilon_greedy_select(fitnesses, epsilon, count, seed)
    else:
        return tournament_select(fitnesses, tournament_size, count, seed)


def evolve(
    data: Any,
    pipeline_factory: Callable,
    param_bounds: list[dict[str, Any]],
    pop_size: int = 50,
    generations: int = 100,
    mutation_rate: float = 0.1,
    crossover_rate: float = 0.7,
    tournament_size: int = 3,
    elite_count: int = 2,
    sigma: float = 0.1,
    fraction_train: float = 0.7,
    seed: int = 42,
    fitness_fn: Callable | None = None,
    selection: str | SelectionMethod = "tournament",
    early_stopping: bool = True,
    patience: int = 20,
    adaptive_mutation: bool = False,
    epsilon: float = 0.1,
) -> EvolutionResult:
    """Run evolutionary optimization on a strategy.

    Args:
        data: Historical data for backtesting.
        pipeline_factory: Function(params_dict) -> pipeline for hz.backtest().
        param_bounds: List of {"name": str, "min": float, "max": float, "discrete": bool}.
        selection: Selection method - "tournament", "uniform", "roulette", "rank", "ucb1",
            or "epsilon_greedy". Also accepts SelectionMethod enum.
        early_stopping: Enable convergence detection (patience + diversity collapse).
        patience: Generations without improvement before stopping.
        adaptive_mutation: Enable Rechenberg's 1/5 rule for adaptive sigma.
        epsilon: Epsilon value for epsilon-greedy selection.
        fitness_fn: Custom fitness function(backtest_result) -> float. Default uses Sharpe ratio.
    """
    from horizon._horizon import auth_require_ultra
    auth_require_ultra()

    bounds_min = [pb["min"] for pb in param_bounds]
    bounds_max = [pb["max"] for pb in param_bounds]
    is_discrete = [pb.get("discrete", False) for pb in param_bounds]
    param_names = [pb["name"] for pb in param_bounds]

    # Resolve selection method
    method = _resolve_selection_method(selection)

    # Split data for walk-forward
    if isinstance(data, list):
        split = int(len(data) * fraction_train)
        train_data = data[:split]
        val_data = data[split:]
    else:
        train_data = data
        val_data = data

    fit_fn = fitness_fn or _default_fitness

    def evaluate(genome: list[float], use_data: Any) -> float:
        params = {param_names[i]: genome[i] for i in range(len(param_names))}
        try:
            pipeline = pipeline_factory(params)
            from horizon import backtest
            result = backtest(pipeline, data=use_data)
            return fit_fn(result)
        except Exception as e:
            logger.debug("Evolution fitness eval failed: %s", e)
            return float("-inf")

    # Initialize population
    population = init_population(bounds_min, bounds_max, is_discrete, pop_size, seed)
    fitnesses = [evaluate(genome, train_data) for genome in population]

    best_genome = population[0]
    best_fitness = fitnesses[0]
    fitness_history = []

    # UCB1 tracking
    selection_counts = [0] * pop_size
    total_selections = 0

    # Adaptive mutation state
    current_sigma = sigma
    success_count = 0
    total_count = 0

    # Convergence tracking
    gens_without_improvement = 0
    prev_best = float("-inf")

    for gen in range(generations):
        # Track best
        for i, f in enumerate(fitnesses):
            if f > best_fitness:
                best_fitness = f
                best_genome = population[i].copy()
        fitness_history.append(best_fitness)

        # Early stopping check
        if early_stopping:
            diversity = population_diversity(population)
            conv = check_convergence(
                best_fitness, diversity, prev_best,
                gens_without_improvement, patience, 1e-4, 0.01,
            )
            gens_without_improvement = conv.generations_without_improvement
            prev_best = conv.best_fitness
            if conv.converged:
                logger.info("Evolution converged at generation %d: %s", gen, conv.reason)
                break

        # Adaptive mutation
        if adaptive_mutation and total_count > 0 and gen % 10 == 0:
            state = adapt_mutation_sigma(current_sigma, success_count, total_count)
            current_sigma = state.sigma
            success_count = 0
            total_count = 0

        # Sort by fitness for elitism
        ranked = sorted(range(len(fitnesses)), key=lambda i: fitnesses[i], reverse=True)
        new_pop = [population[ranked[i]] for i in range(min(elite_count, len(ranked)))]

        # Generate offspring
        while len(new_pop) < pop_size:
            offset = gen * 1000 + len(new_pop)
            parents = _select_parents(
                fitnesses, method, 2, seed + offset,
                tournament_size=tournament_size,
                selection_counts=selection_counts,
                total_selections=total_selections,
                epsilon=epsilon,
            )
            total_selections += 2
            for p_idx in parents:
                if p_idx < len(selection_counts):
                    selection_counts[p_idx] += 1

            p1 = population[parents[0]]
            p2 = population[parents[1]] if len(parents) > 1 else p1

            child = uniform_crossover(p1, p2, crossover_rate, seed + gen * 2000 + len(new_pop))
            child = gaussian_mutate(
                child, bounds_min, bounds_max, is_discrete,
                current_sigma, mutation_rate, seed + gen * 3000 + len(new_pop),
            )
            new_pop.append(child)

        population = new_pop[:pop_size]
        old_fitnesses = fitnesses
        fitnesses = [evaluate(genome, train_data) for genome in population]

        # Track mutation success for adaptive mutation
        if adaptive_mutation:
            for i in range(elite_count, len(fitnesses)):
                total_count += 1
                if fitnesses[i] > (old_fitnesses[ranked[0]] if ranked else 0):
                    success_count += 1

        # Reset selection counts for new population
        selection_counts = [0] * len(population)

        if gen % 10 == 0:
            logger.info("Generation %d: best_fitness=%.6f, sigma=%.4f", gen, best_fitness, current_sigma)

    # Validate on held-out data
    val_fitness = evaluate(best_genome, val_data) if val_data is not data else best_fitness

    return EvolutionResult(
        best_genome=best_genome,
        best_fitness=val_fitness,
        generations_run=len(fitness_history),
        fitness_history=fitness_history,
    )


def evolve_nsga2(
    data: Any,
    pipeline_factory: Callable,
    param_bounds: list[dict[str, Any]],
    objectives: list[dict[str, str]],
    pop_size: int = 50,
    generations: int = 100,
    mutation_rate: float = 0.1,
    crossover_rate: float = 0.7,
    sigma: float = 0.1,
    fraction_train: float = 0.7,
    seed: int = 42,
    early_stopping: bool = True,
    patience: int = 20,
) -> dict[str, Any]:
    """Multi-objective evolution using NSGA-II.

    Optimizes multiple objectives simultaneously (e.g., maximize Sharpe AND
    minimize max drawdown). Returns the Pareto front of non-dominated solutions.

    Based on Deb et al., "A Fast and Elitist Multiobjective Genetic Algorithm: NSGA-II" (2002).

    Args:
        data: Historical data for backtesting.
        pipeline_factory: Function(params_dict) -> pipeline.
        param_bounds: Parameter bounds.
        objectives: List of {"metric": str, "direction": "maximize"|"minimize"}.
            Supported metrics: "sharpe_ratio", "total_pnl", "max_drawdown", "win_rate",
            "profit_factor", "sortino_ratio", "calmar_ratio".
        early_stopping: Enable convergence detection.
        patience: Generations without improvement before stopping.
    """
    from horizon._horizon import auth_require_ultra
    auth_require_ultra()

    bounds_min = [pb["min"] for pb in param_bounds]
    bounds_max = [pb["max"] for pb in param_bounds]
    is_discrete = [pb.get("discrete", False) for pb in param_bounds]
    param_names = [pb["name"] for pb in param_bounds]

    # Split data
    if isinstance(data, list):
        split = int(len(data) * fraction_train)
        train_data = data[:split]
    else:
        train_data = data

    def evaluate_multi(genome: list[float]) -> list[float]:
        params = {param_names[i]: genome[i] for i in range(len(param_names))}
        try:
            pipeline = pipeline_factory(params)
            from horizon import backtest
            result = backtest(pipeline, data=train_data)

            obj_values = []
            for obj in objectives:
                metric = obj["metric"]
                direction = obj.get("direction", "maximize")
                raw = getattr(result, metric, 0.0)
                if direction == "minimize":
                    raw = -raw
                obj_values.append(raw if raw == raw else 0.0)  # NaN guard
            return obj_values
        except Exception as e:
            logger.debug("NSGA-II fitness eval failed: %s", e)
            return [float("-inf")] * len(objectives)

    # Initialize population
    population = init_population(bounds_min, bounds_max, is_discrete, pop_size, seed)
    all_objectives = [evaluate_multi(g) for g in population]

    gens_without_improvement = 0
    prev_best_front_size = 0

    for gen in range(generations):
        # NSGA-II selection
        result = nsga2_select(all_objectives, pop_size)
        selected = result.selected_indices

        # Build next generation
        new_pop = [population[i] for i in selected]
        new_obj = [all_objectives[i] for i in selected]

        # Generate offspring via crossover + mutation
        offspring = []
        offspring_obj = []
        rng_offset = gen * 10000
        while len(offspring) < pop_size:
            idx_a = selected[rng_offset % len(selected)]
            idx_b = selected[(rng_offset + 1) % len(selected)]
            rng_offset += 2

            child = uniform_crossover(
                population[idx_a], population[idx_b],
                crossover_rate, seed + gen * 5000 + len(offspring),
            )
            child = gaussian_mutate(
                child, bounds_min, bounds_max, is_discrete,
                sigma, mutation_rate, seed + gen * 6000 + len(offspring),
            )
            offspring.append(child)
            offspring_obj.append(evaluate_multi(child))

        # Merge parents + offspring
        combined_pop = new_pop + offspring
        combined_obj = new_obj + offspring_obj

        # Select top pop_size from combined using NSGA-II
        result = nsga2_select(combined_obj, pop_size)
        population = [combined_pop[i] for i in result.selected_indices]
        all_objectives = [combined_obj[i] for i in result.selected_indices]

        # Early stopping
        if early_stopping:
            front = pareto_front(all_objectives)
            if len(front) == prev_best_front_size:
                gens_without_improvement += 1
            else:
                gens_without_improvement = 0
                prev_best_front_size = len(front)
            if gens_without_improvement >= patience:
                logger.info("NSGA-II converged at generation %d", gen)
                break

        if gen % 10 == 0:
            front = pareto_front(all_objectives)
            logger.info("Generation %d: pareto_front_size=%d", gen, len(front))

    # Extract Pareto front
    front_indices = pareto_front(all_objectives)
    pareto_genomes = [population[i] for i in front_indices]
    pareto_objectives = [all_objectives[i] for i in front_indices]

    return {
        "pareto_front": [
            {
                "genome": genome,
                "params": {param_names[j]: genome[j] for j in range(len(param_names))},
                "objectives": {
                    objectives[k]["metric"]: (
                        obj[k] if objectives[k].get("direction", "maximize") == "maximize"
                        else -obj[k]
                    )
                    for k in range(len(objectives))
                },
            }
            for genome, obj in zip(pareto_genomes, pareto_objectives)
        ],
        "generations_run": gen + 1 if 'gen' in dir() else generations,
        "pareto_size": len(front_indices),
        "objective_names": [o["metric"] for o in objectives],
    }


def evolve_islands(
    data: Any,
    pipeline_factory: Callable,
    param_bounds: list[dict[str, Any]],
    pop_size: int = 50,
    generations: int = 100,
    mutation_rate: float = 0.1,
    crossover_rate: float = 0.7,
    sigma: float = 0.1,
    fraction_train: float = 0.7,
    seed: int = 42,
    fitness_fn: Callable | None = None,
    n_islands: int = 4,
    migration_interval: int = 5,
    migration_count: int = 2,
    topology: str | MigrationTopology = "ring",
    early_stopping: bool = True,
    patience: int = 20,
) -> dict[str, Any]:
    """Island model evolution with parallel populations.

    Multiple populations evolve independently using different selection methods.
    Periodically, top individuals migrate between islands to share genetic material.

    Inspired by the Darwin SDK's island model and the ProFiT framework's approach
    of maintaining diverse search strategies via per-island selection specialization.

    Args:
        data: Historical data for backtesting.
        pipeline_factory: Function(params_dict) -> pipeline.
        param_bounds: Parameter bounds.
        n_islands: Number of parallel populations.
        migration_interval: Generations between migration events.
        migration_count: Number of top individuals to migrate.
        topology: "ring" or "fully_connected".
        early_stopping: Enable convergence detection.
        patience: Generations without improvement before stopping.
    """
    from horizon._horizon import auth_require_ultra
    auth_require_ultra()

    bounds_min = [pb["min"] for pb in param_bounds]
    bounds_max = [pb["max"] for pb in param_bounds]
    is_discrete = [pb.get("discrete", False) for pb in param_bounds]
    param_names = [pb["name"] for pb in param_bounds]

    # Resolve topology
    if isinstance(topology, str):
        topo = MigrationTopology.Ring if topology.lower() == "ring" else MigrationTopology.FullyConnected
    else:
        topo = topology

    config = IslandConfig(n_islands, migration_interval, migration_count, topo)

    # Assign different selection methods per island
    methods = island_selection_methods(n_islands)

    # Split data
    if isinstance(data, list):
        split = int(len(data) * fraction_train)
        train_data = data[:split]
        val_data = data[split:]
    else:
        train_data = data
        val_data = data

    fit_fn = fitness_fn or _default_fitness

    def evaluate(genome: list[float], use_data: Any) -> float:
        params = {param_names[i]: genome[i] for i in range(len(param_names))}
        try:
            pipeline = pipeline_factory(params)
            from horizon import backtest
            result = backtest(pipeline, data=use_data)
            return fit_fn(result)
        except Exception as e:
            logger.debug("Island evolution fitness eval failed: %s", e)
            return float("-inf")

    # Initialize islands
    island_size = pop_size // n_islands
    islands: list[list[list[float]]] = []
    island_fitnesses: list[list[float]] = []

    for i in range(n_islands):
        pop = init_population(bounds_min, bounds_max, is_discrete, island_size, seed + i * 1000)
        fits = [evaluate(g, train_data) for g in pop]
        islands.append(pop)
        island_fitnesses.append(fits)

    global_best_fitness = float("-inf")
    global_best_genome: list[float] = []
    fitness_history: list[float] = []
    gens_without_improvement = 0
    prev_best = float("-inf")

    for gen in range(generations):
        # Evolve each island
        for island_idx in range(n_islands):
            method = methods[island_idx]
            pop = islands[island_idx]
            fits = island_fitnesses[island_idx]

            if not pop:
                continue

            # Sort for elitism
            ranked = sorted(range(len(fits)), key=lambda i: fits[i], reverse=True)
            elite = min(2, len(ranked))
            new_pop = [pop[ranked[i]] for i in range(elite)]

            # Generate offspring
            while len(new_pop) < island_size:
                offset = gen * 10000 + island_idx * 1000 + len(new_pop)
                parents = _select_parents(
                    fits, method, 2, seed + offset,
                    tournament_size=3,
                )
                p1 = pop[parents[0]]
                p2 = pop[parents[1]] if len(parents) > 1 else p1

                child = uniform_crossover(p1, p2, crossover_rate, seed + offset + 5000)
                child = gaussian_mutate(
                    child, bounds_min, bounds_max, is_discrete,
                    sigma, mutation_rate, seed + offset + 7000,
                )
                new_pop.append(child)

            islands[island_idx] = new_pop[:island_size]
            island_fitnesses[island_idx] = [evaluate(g, train_data) for g in islands[island_idx]]

        # Migration
        if gen > 0 and gen % migration_interval == 0:
            islands = island_migrate(islands, island_fitnesses, config)
            # Trim islands back to size and re-evaluate migrants
            for i in range(n_islands):
                if len(islands[i]) > island_size:
                    fits = [evaluate(g, train_data) for g in islands[i]]
                    ranked = sorted(range(len(fits)), key=lambda j: fits[j], reverse=True)
                    islands[i] = [islands[i][j] for j in ranked[:island_size]]
                    island_fitnesses[i] = [fits[j] for j in ranked[:island_size]]

        # Track global best
        for i in range(n_islands):
            for j, f in enumerate(island_fitnesses[i]):
                if f > global_best_fitness:
                    global_best_fitness = f
                    global_best_genome = islands[i][j].copy()

        fitness_history.append(global_best_fitness)

        # Early stopping
        if early_stopping:
            all_genomes = [g for island in islands for g in island]
            diversity = population_diversity(all_genomes)
            conv = check_convergence(
                global_best_fitness, diversity, prev_best,
                gens_without_improvement, patience, 1e-4, 0.01,
            )
            gens_without_improvement = conv.generations_without_improvement
            prev_best = conv.best_fitness
            if conv.converged:
                logger.info("Island evolution converged at generation %d: %s", gen, conv.reason)
                break

        if gen % 10 == 0:
            logger.info("Island gen %d: best=%.6f, islands=%d", gen, global_best_fitness, n_islands)

    # Validate on held-out data
    val_fitness = evaluate(global_best_genome, val_data) if val_data is not data else global_best_fitness

    # Collect top strategies from all islands
    all_genomes = [(g, f) for i in range(n_islands) for g, f in zip(islands[i], island_fitnesses[i])]
    all_genomes.sort(key=lambda x: x[1], reverse=True)
    top_strategies = [
        {
            "genome": g,
            "params": {param_names[j]: g[j] for j in range(len(param_names))},
            "fitness": f,
        }
        for g, f in all_genomes[:10]
    ]

    return {
        "best_genome": global_best_genome,
        "best_params": {param_names[j]: global_best_genome[j] for j in range(len(param_names))} if global_best_genome else {},
        "best_fitness": val_fitness,
        "generations_run": len(fitness_history),
        "fitness_history": fitness_history,
        "top_strategies": top_strategies,
        "n_islands": n_islands,
        "topology": str(topo),
    }


def _default_fitness(result: Any) -> float:
    """Default fitness function using Sharpe ratio."""
    if hasattr(result, "sharpe_ratio"):
        return result.sharpe_ratio
    if hasattr(result, "total_pnl"):
        return result.total_pnl
    return 0.0


def _resolve_selection_method(selection: str | SelectionMethod) -> SelectionMethod:
    """Convert string to SelectionMethod enum."""
    if isinstance(selection, SelectionMethod):
        return selection
    mapping = {
        "tournament": SelectionMethod.Tournament,
        "uniform": SelectionMethod.Uniform,
        "roulette": SelectionMethod.RouletteWheel,
        "roulette_wheel": SelectionMethod.RouletteWheel,
        "rank": SelectionMethod.Rank,
        "ucb1": SelectionMethod.Ucb1,
        "epsilon_greedy": SelectionMethod.EpsilonGreedy,
        "epsilon-greedy": SelectionMethod.EpsilonGreedy,
    }
    return mapping.get(selection.lower(), SelectionMethod.Tournament)
