import random
import numpy as np


class NSGA2:
    """NSGA-II算法类：实现多目标优化的非支配排序、选择、环境选择等核心操作"""

    def __init__(self, pop_size, objective_num):
        """
        初始化NSGA-II
        :param pop_size: 种群大小
        :param objective_num: 目标函数数量
        """
        self.pop_size = pop_size  # 种群大小
        self.objective_num = objective_num  # 目标数量

    def fast_non_dominated_sort(self, objective_values):
        """
        快速非支配排序：将个体按支配关系分层（前沿）
        :param objective_values: 所有个体的目标函数值列表
        :return: ranks（每个个体的层级）和 fronts（每层的个体索引列表）
        """
        pop_size = len(objective_values)
        domination_count = [0] * pop_size  # 每个个体被支配的次数
        dominated_solutions = [[] for _ in range(pop_size)]  # 每个个体支配的个体列表
        ranks = [0] * pop_size  # 每个个体的层级（0为最优）
        front = [[]]  # 第0层前沿（初始为空）

        # 1. 计算支配关系
        for p in range(pop_size):
            for q in range(pop_size):
                if p == q:
                    continue  # 个体不与自身比较
                # 判断p是否支配q
                if self._dominates(objective_values[p], objective_values[q]):
                    dominated_solutions[p].append(q)  # p支配q，记录q
                # 判断q是否支配p
                elif self._dominates(objective_values[q], objective_values[p]):
                    domination_count[p] += 1  # p被q支配，计数+1
            # 若p未被任何个体支配，加入第0层前沿
            if domination_count[p] == 0:
                ranks[p] = 0
                front[0].append(p)

        # 2. 生成后续层级的前沿
        i = 0
        while len(front[i]) > 0:  # 当前层非空时继续
            next_front = []  # 下一层前沿
            for p in front[i]:  # 遍历当前层的每个个体
                for q in dominated_solutions[p]:  # p支配的个体q
                    domination_count[q] -= 1  # q的被支配计数-1
                    if domination_count[q] == 0:  # q不再被任何个体支配
                        ranks[q] = i + 1  # 层级为当前层+1
                        next_front.append(q)  # 加入下一层
            i += 1
            front.append(next_front)  # 追加下一层

        return ranks, front[:-1]  # 返回层级和非空前沿

    def _dominates(self, a, b):
        """
        判断个体a是否支配个体b
        支配条件：a的所有目标函数值≤b，且至少有一个目标函数值<b
        :param a: 个体a的目标值
        :param b: 个体b的目标值
        :return: 布尔值（a支配b则为True）
        """
        all_le = all(a[i] <= b[i] for i in range(self.objective_num))  # 所有目标a≤b
        any_lt = any(a[i] < b[i] for i in range(self.objective_num))  # 至少一个目标a<b
        return all_le and any_lt

    def crowding_distance(self, objective_values, front):
        """
        计算前沿中个体的拥挤度（衡量个体在前沿中的分散程度）
        :param objective_values: 所有个体的目标值
        :param front: 当前前沿的个体索引列表
        :return: 每个个体的拥挤度（与front顺序对应）
        """
        pop_size = len(front)
        distance = [0.0] * pop_size  # 拥挤度初始化
        if pop_size <= 1:  # 前沿只有1个个体时，拥挤度为无穷大
            return distance

        # 对每个目标函数计算拥挤度
        for m in range(self.objective_num):
            # 按目标m的值对前沿个体排序
            sorted_indices = sorted(range(pop_size), key=lambda i: objective_values[front[i]][m])
            # 目标m的最大值和最小值
            max_val = objective_values[front[sorted_indices[-1]]][m]
            min_val = objective_values[front[sorted_indices[0]]][m]

            if max_val - min_val == 0:  # 目标值无差异，无需计算
                continue

            # 边界个体（最小和最大）的拥挤度设为无穷大（确保被保留）
            distance[sorted_indices[0]] = float('inf')
            distance[sorted_indices[-1]] = float('inf')

            # 中间个体的拥挤度：相邻个体的目标值差除以目标值范围
            for i in range(1, pop_size - 1):
                distance[sorted_indices[i]] += (
                                                       objective_values[front[sorted_indices[i + 1]]][m] -
                                                       objective_values[front[sorted_indices[i - 1]]][m]
                                               ) / (max_val - min_val)

        return distance

    def selection(self, population, objective_values):
        """
        锦标赛选择：从种群中选择父代（保留较优个体）
        :param population: 当前种群
        :param objective_values: 种群的目标函数值
        :return: 选择后的父代种群（与原种群大小相同）
        """
        pop_size = len(population)
        # 计算每个个体的层级和拥挤度
        ranks, _ = self.fast_non_dominated_sort(objective_values)
        distances = self._calculate_all_crowding_distances(objective_values, ranks)

        selected = []
        for _ in range(pop_size):
            # 随机选择两个个体进行锦标赛
            i = random.randint(0, pop_size - 1)
            j = random.randint(0, pop_size - 1)
            # 选择层级低（更优）或同层级拥挤度高（更分散）的个体
            if ranks[i] < ranks[j] or (ranks[i] == ranks[j] and distances[i] >= distances[j]):
                selected.append(population[i])
            else:
                selected.append(population[j])

        return np.array(selected)

    def _calculate_all_crowding_distances(self, objective_values, ranks):
        """
        计算所有个体的拥挤度（内部方法）
        :param objective_values: 目标函数值
        :param ranks: 每个个体的层级
        :return: 所有个体的拥挤度列表
        """
        max_rank = max(ranks) if ranks else 0  # 最大层级
        all_distances = [0.0] * len(objective_values)  # 所有个体的拥挤度

        # 按层级计算拥挤度
        for r in range(max_rank + 1):
            front = [i for i in range(len(objective_values)) if ranks[i] == r]  # 当前层级的个体
            if len(front) <= 1:  # 单个个体拥挤度为无穷大
                for i in front:
                    all_distances[i] = float('inf')
                continue
            # 计算当前层级的拥挤度
            distances = self.crowding_distance(objective_values, front)
            # 映射到全局索引
            for i, idx in enumerate(front):
                all_distances[idx] = distances[i]

        return all_distances

    def environmental_selection(self, population, objective_values):
        """
        环境选择：从合并的父代和子代中选择最优的pop_size个个体
        :param population: 合并的种群（父代+子代）
        :param objective_values: 合并种群的目标函数值
        :return: 选择后的种群和对应的目标值
        """
        # 非支配排序
        ranks, fronts = self.fast_non_dominated_sort(objective_values)
        selected = []  # 选中的个体索引

        # 按层级从优到劣选择，直到接近种群大小
        for front in fronts:
            if len(selected) + len(front) <= self.pop_size:
                selected.extend(front)  # 整层加入
            else:
                # 当前层无法全加入时，按拥挤度选择剩余个体
                distances = self.crowding_distance(objective_values, front)
                # 按拥挤度降序排序（优先选择分散的个体）
                sorted_front = [x for _, x in sorted(zip(distances, front), reverse=True)]
                # 选择剩余所需数量
                selected.extend(sorted_front[:self.pop_size - len(selected)])
                break

        # 返回选中的个体和对应的目标值
        return population[selected], [objective_values[i] for i in selected]