import numpy as np
from typing import Dict, Any, Tuple, List, Optional
from ..optimizer import Optimizer
from ..privacy_mechanism import PrivacyMechanism
from ..risk_meter import RiskMeter
from ..utility_evaluator import UtilityEvaluator
from ..param_utils import is_valid_combo, sample_valid_combo


class GeneticOptimizer(Optimizer):
    """
    遗传算法优化器。
    支持连续、整数和类别参数，并支持参数依赖。

    参数:
        population_size: int，种群大小（默认 20）
        generations: int，进化代数（默认 10）
        mutation_rate: float，变异概率（默认 0.1）
        crossover_rate: float，交叉概率（默认 0.8）
        random_state: int，随机种子（默认 42）
        tournament_size: int，锦标赛选择大小（默认 3）
    """

    def __init__(self,
                 population_size: int = 20,
                 generations: int = 10,
                 mutation_rate: float = 0.1,
                 crossover_rate: float = 0.8,
                 random_state: int = 42,
                 tournament_size: int = 3):
        self.population_size = population_size
        self.generations = generations
        self.mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate
        self.random_state = random_state
        self.tournament_size = tournament_size
        np.random.seed(random_state)

    def _encode_param_space(self, param_space: Dict[str, Any]) -> Dict:
        """
        编码参数空间，返回每个参数的类型、范围等信息。
        """
        encoded = {}
        self.param_names = list(param_space.keys())
        for name, spec in param_space.items():
            if isinstance(spec, tuple) and len(spec) == 3:
                low, high, typ = spec
                encoded[name] = {'type': typ, 'low': low, 'high': high}
            elif isinstance(spec, list):
                encoded[name] = {'type': 'categorical', 'choices': spec}
            elif isinstance(spec, dict):
                # 带依赖的字典，保留所有信息
                encoded[name] = spec.copy()
            else:
                raise ValueError(f"参数空间格式不支持: {name}: {spec}")
        return encoded

    def _random_individual(self, encoded: Dict, param_space: Dict) -> Optional[Dict[str, Any]]:
        """
        生成一个随机合法的个体。优先使用 sample_valid_combo。
        如果失败，则尝试逐个生成并检查（最多重试100次）。
        """
        # 尝试用 sample_valid_combo 生成
        ind = sample_valid_combo(param_space)
        if ind is not None:
            return ind

        # 备选方案：逐个参数随机生成并检查
        for _ in range(100):
            cand = {}
            for name, info in encoded.items():
                if info['type'] == 'int':
                    low, high = info['range'] if 'range' in info else (info['low'], info['high'])
                    cand[name] = np.random.randint(low, high + 1)
                elif info['type'] == 'float':
                    low, high = info['range'] if 'range' in info else (info['low'], info['high'])
                    cand[name] = np.random.uniform(low, high)
                else:  # categorical
                    choices = info['choices'] if 'choices' in info else info.get('choices', [])
                    cand[name] = np.random.choice(choices)
            if is_valid_combo(cand, param_space):
                return cand
        return None

    def _mutate(self, ind: Dict[str, Any], encoded: Dict, param_space: Dict) -> Dict[str, Any]:
        """
        变异个体，并确保结果合法。如果变异后非法，则重试最多100次。
        """
        for _ in range(100):
            new_ind = ind.copy()
            for name, info in encoded.items():
                if np.random.rand() < self.mutation_rate:
                    if info['type'] == 'int':
                        low, high = info['range'] if 'range' in info else (info['low'], info['high'])
                        new_ind[name] = np.random.randint(low, high + 1)
                    elif info['type'] == 'float':
                        low, high = info['range'] if 'range' in info else (info['low'], info['high'])
                        new_ind[name] = np.random.uniform(low, high)
                    else:  # categorical
                        choices = info['choices'] if 'choices' in info else info.get('choices', [])
                        new_ind[name] = np.random.choice(choices)
            if is_valid_combo(new_ind, param_space):
                return new_ind
        # 多次重试失败，返回原个体
        return ind

    def _crossover(self, p1: Dict[str, Any], p2: Dict[str, Any], param_space: Dict) -> Tuple[
        Dict[str, Any], Dict[str, Any]]:
        """
        均匀交叉，并确保子代合法。如果交叉后非法，则重试最多100次。
        """
        for _ in range(100):
            c1, c2 = {}, {}
            for name in p1.keys():
                if np.random.rand() < 0.5:
                    c1[name] = p1[name]
                    c2[name] = p2[name]
                else:
                    c1[name] = p2[name]
                    c2[name] = p1[name]
            if is_valid_combo(c1, param_space) and is_valid_combo(c2, param_space):
                return c1, c2
        # 多次重试失败，返回父代副本
        return p1.copy(), p2.copy()

    def _tournament_selection(self, population: List, fitnesses: List, k: int) -> List:
        """锦标赛选择"""
        selected = []
        for _ in range(k):
            idx = np.random.choice(len(population), self.tournament_size, replace=False)
            best_idx = idx[np.argmax([fitnesses[i] for i in idx])]
            selected.append(population[best_idx])
        return selected

    def search(self,
               original_data: np.ndarray,
               mechanism: PrivacyMechanism,
               risk_meter: RiskMeter,
               utility_evaluator: UtilityEvaluator,
               param_space: Dict[str, Any],
               risk_constraint: Tuple[float, float],
               max_iter: int = 100) -> Tuple[Dict[str, Any], float, float]:
        """
        遗传算法搜索。
        """
        encoded = self._encode_param_space(param_space)

        # 初始化种群（确保所有个体合法）
        population = []
        while len(population) < self.population_size:
            ind = self._random_individual(encoded, param_space)
            if ind is not None:
                population.append(ind)

        best_individual = None
        best_fitness = -np.inf
        best_risk = None

        # 进化代数
        n_generations = min(self.generations, max_iter // self.population_size)

        for gen in range(n_generations):
            # 评估适应度
            fitnesses = []
            for ind in population:
                perturbed = mechanism.apply(original_data, ind)
                risk = risk_meter.measure(original_data, perturbed, params=ind)
                if risk_constraint[0] <= risk <= risk_constraint[1]:
                    utility = utility_evaluator.evaluate(original_data, perturbed)
                    fitness = utility
                else:
                    fitness = -np.inf  # 不可行解
                fitnesses.append(fitness)

                # 更新全局最优
                if fitness > best_fitness:
                    best_fitness = fitness
                    best_individual = ind.copy()
                    best_risk = risk

            # 选择父代
            selected = self._tournament_selection(population, fitnesses, self.population_size)

            # 生成下一代
            next_population = []
            while len(next_population) < self.population_size:
                # 选择两个父代
                parents = np.random.choice(len(selected), 2, replace=False)
                p1, p2 = selected[parents[0]], selected[parents[1]]

                # 交叉
                if np.random.rand() < self.crossover_rate:
                    c1, c2 = self._crossover(p1, p2, param_space)
                else:
                    c1, c2 = p1.copy(), p2.copy()

                # 变异
                c1 = self._mutate(c1, encoded, param_space)
                c2 = self._mutate(c2, encoded, param_space)

                next_population.extend([c1, c2])

            population = next_population[:self.population_size]

        # 如果未找到可行解，返回空
        if best_individual is None:
            return {}, -np.inf, np.inf

        return best_individual, best_fitness, best_risk