import random
import numpy as np

class NSGA2:
    def __init__(self, pop_size, obj_num):
        self.pop_size = pop_size
        self.obj_num = obj_num

    def fast_non_dominated_sort(self, pop_obj):
        """快速非支配排序"""
        pop_size = len(pop_obj)
        dominated = [[] for _ in range(pop_size)]  # 被支配个体列表
        rank = [0] * pop_size  # 个体的非支配等级
        n = [0] * pop_size  # 支配该个体的个体数量

        # 计算每个个体的支配关系
        for p in range(pop_size):
            for q in range(pop_size):
                if p != q:
                    # p支配q
                    if all(pop_obj[p][i] <= pop_obj[q][i] for i in range(self.obj_num)) and \
                            any(pop_obj[p][i] < pop_obj[q][i] for i in range(self.obj_num)):
                        dominated[p].append(q)
                    # q支配p
                    elif all(pop_obj[q][i] <= pop_obj[p][i] for i in range(self.obj_num)) and \
                            any(pop_obj[q][i] < pop_obj[p][i] for i in range(self.obj_num)):
                        n[p] += 1

            # 找到等级为0的个体
            if n[p] == 0:
                rank[p] = 0

        # 计算其他等级
        current_rank = 0
        while True:
            next_rank = []
            for p in range(pop_size):
                if rank[p] == current_rank:
                    for q in dominated[p]:
                        n[q] -= 1
                        if n[q] == 0 and rank[q] == 0:
                            rank[q] = current_rank + 1
                            next_rank.append(q)
            if not next_rank:
                break
            current_rank += 1

        return rank

    def crowding_distance(self, pop_obj, rank):
        """计算拥挤度距离"""
        pop_size = len(pop_obj)
        distance = [0.0] * pop_size
        max_rank = max(rank) if rank else 0

        # 对每个等级的个体计算拥挤度
        for r in range(max_rank + 1):
            # 获取当前等级的个体索引
            current_indices = [i for i in range(pop_size) if rank[i] == r]
            if len(current_indices) <= 1:
                continue

            # 对每个目标函数进行排序
            for m in range(self.obj_num):
                # 按目标m的值排序
                sorted_indices = sorted(current_indices, key=lambda x: pop_obj[x][m])
                min_val = pop_obj[sorted_indices[0]][m]
                max_val = pop_obj[sorted_indices[-1]][m]

                # 边界个体的拥挤度设为无穷大
                distance[sorted_indices[0]] = float('inf')
                distance[sorted_indices[-1]] = float('inf')

                # 计算中间个体的拥挤度
                for i in range(1, len(sorted_indices) - 1):
                    if max_val - min_val == 0:
                        continue
                    distance[sorted_indices[i]] += (pop_obj[sorted_indices[i + 1]][m] - pop_obj[sorted_indices[i - 1]][m]) / (max_val - min_val)

        return distance

    def selection(self, pop, pop_obj):
        """选择操作：基于非支配排序和拥挤度的锦标赛选择，惩罚拥挤度为0的个体"""
        pop_size = len(pop)
        rank = self.fast_non_dominated_sort(pop_obj)
        distance = self.crowding_distance(pop_obj, rank)
        selected = []

        for _ in range(pop_size):
            # 随机选择两个个体进行锦标赛
            i = random.randint(0, pop_size - 1)
            j = random.randint(0, pop_size - 1)

            # 优先选择等级低（更优）的个体
            if rank[i] < rank[j]:
                selected.append(pop[i])
            elif rank[i] > rank[j]:
                selected.append(pop[j])
            # 等级相同则选择拥挤度大的个体（惩罚拥挤度为0的重复解）
            else:
                # 对拥挤度为0的个体进行惩罚
                if distance[i] == 0 and distance[j] > 0:
                    selected.append(pop[j])
                elif distance[j] == 0 and distance[i] > 0:
                    selected.append(pop[i])
                elif distance[i] >= distance[j]:
                    selected.append(pop[i])
                else:
                    selected.append(pop[j])

        return selected