import random
import numpy as np


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

    def fast_non_dominated_sort(self, objective_values):
        """标准快速非支配排序实现"""
        pop_size = len(objective_values)
        domination_count = [0] * pop_size
        dominated_solutions = [[] for _ in range(pop_size)]
        ranks = [0] * pop_size
        front = [[]]  # 第0层

        # 计算支配关系
        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)
                # q支配p？
                elif self._dominates(objective_values[q], objective_values[p]):
                    domination_count[p] += 1
            if domination_count[p] == 0:
                ranks[p] = 0
                front[0].append(p)

        # 处理后续层
        i = 0
        while len(front[i]) > 0:
            next_front = []
            for p in front[i]:
                for q in dominated_solutions[p]:
                    domination_count[q] -= 1
                    if domination_count[q] == 0:
                        ranks[q] = i + 1
                        next_front.append(q)
            i += 1
            front.append(next_front)

        return ranks, front[:-1]  # 返回rank数组和各层front列表

    def _dominates(self, a, b):
        """判断a是否支配b：a的所有目标≤b，且至少一个目标<b"""
        all_le = all(a[i] <= b[i] for i in range(self.objective_num))
        any_lt = any(a[i] < b[i] for i in range(self.objective_num))
        return all_le and any_lt

    def crowding_distance(self, objective_values, front):
        """计算拥挤度"""
        pop_size = len(front)
        distance = [0.0] * pop_size
        if pop_size <= 1:
            return distance
        # 按每个目标排序
        for m in range(self.objective_num):
            sorted_indices = sorted(range(pop_size), key=lambda i: objective_values[front[i]][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):
        """锦标赛选择（2选1）"""
        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):
        """计算所有个体的拥挤度"""
        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个个体"""
        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:
                # 计算当前front的拥挤度，选择剩余数量的个体
                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]