import random
import matplotlib.pyplot as plt
import numpy as np
from Decode import Decode
from Encode import Encode
from GA import GA
from Instance import *
from NSGA2 import NSGA2

# 绘制甘特图
def Gantt(Machines):
    M = ['red', 'blue', 'yellow', 'orange', 'green', 'palegoldenrod', 'purple', 'pink', 'Thistle', 'Magenta',
         'SlateBlue', 'RoyalBlue', 'Cyan', 'Aqua', 'floralwhite', 'ghostwhite', 'goldenrod', 'mediumslateblue',
         'navajowhite', 'nawy', 'sandybrown', 'moccasin']
    for i in range(len(Machines)):
        Machine = Machines[i]
        Start_time = Machine.O_start
        End_time = Machine.O_end
        for i_1 in range(len(End_time)):
            plt.barh(i, width=End_time[i_1] - Start_time[i_1], height=0.8, left=Start_time[i_1],
                     color=M[Machine.assigned_task[i_1][0] - 1], edgecolor='black')
            plt.text(x=Start_time[i_1] + (End_time[i_1] - Start_time[i_1]) / 2 - 0.5, y=i,
                     s=Machine.assigned_task[i_1][0])
    plt.yticks(np.arange(len(Machines) + 1), np.arange(1, len(Machines) + 2))
    plt.title('Scheduling Gantt chart')
    plt.ylabel('Machines')
    plt.xlabel('Time(min)')
    plt.savefig('优化后排程方案的甘特图.png')
    plt.show()

if __name__ == '__main__':
    # 初始化参数
    g = GA()
    e = Encode(Processing_time, g.Pop_size, J, J_num, M_num)
    CHS1 = e.Global_initial()
    CHS2 = e.Random_initial()
    CHS3 = e.Local_initial()
    C = np.vstack((CHS1, CHS2, CHS3))

    # 双目标优化相关
    nsga2 = NSGA2(g.Pop_size, 2)  # 2个优化目标
    Optimal_solutions = []  # 存储非支配解
    Optimal_fit_values = []  # 存储非支配解的目标值

    # 新增：用于存储所有代的所有解
    all_solutions = []  # 存储所有个体（染色体）
    all_fitnesses = []  # 存储所有个体的目标值

    for i in range(g.Max_Itertions):
        print(f"iter_{i} start!")
        Fit = g.fitness(C, J, Processing_time, M_num, O_num)

        # 记录当前代的所有解
        all_solutions.extend(C)
        all_fitnesses.extend(Fit)

        # 非支配排序
        rank = nsga2.fast_non_dominated_sort(Fit)
        current_non_dominated = [C[j] for j in range(len(C)) if rank[j] == 0]
        current_non_dominated_fit = [Fit[j] for j in range(len(C)) if rank[j] == 0]

        # 更新全局非支配解
        Optimal_solutions.extend(current_non_dominated)
        Optimal_fit_values.extend(current_non_dominated_fit)

        # 对全局解重新筛选非支配解
        if Optimal_solutions:
            rank_all = nsga2.fast_non_dominated_sort(Optimal_fit_values)
            # 保留等级为0的解
            Optimal_solutions = [Optimal_solutions[j] for j in range(len(Optimal_solutions)) if rank_all[j] == 0]
            Optimal_fit_values = [Optimal_fit_values[j] for j in range(len(Optimal_fit_values)) if rank_all[j] == 0]
            # 控制解的数量，过多时保留拥挤度高的解
            if len(Optimal_solutions) > g.Pop_size:
                distance = nsga2.crowding_distance(Optimal_fit_values, rank_all)
                # 按拥挤度排序并保留前Pop_size个解
                sorted_indices = sorted(range(len(distance)), key=lambda k: distance[k], reverse=True)
                Optimal_solutions = [Optimal_solutions[j] for j in sorted_indices[:g.Pop_size]]
                Optimal_fit_values = [Optimal_fit_values[j] for j in sorted_indices[:g.Pop_size]]

        # 选择操作（基于NSGA-II）
        selected = nsga2.selection(C, Fit)

        # 交叉变异操作
        new_pop = []
        for j in range(len(selected)):
            if random.random() < g.Pc:
                # 选择另一个个体进行交叉
                mate_idx = random.randint(0, len(selected) - 1)
                if random.random() < g.Pv:
                    offspring1, offspring2 = g.machine_cross(selected[j], selected[mate_idx], O_num)
                else:
                    offspring1, offspring2 = g.operation_cross(selected[j], selected[mate_idx], O_num, J_num)
                new_pop.append(offspring1)
                new_pop.append(offspring2)
            else:
                new_pop.append(selected[j])

            # 变异操作
            if random.random() < g.Pm:
                if random.random() < g.Pw:
                    mutated = g.machine_variation(selected[j], Processing_time, O_num, J)
                else:
                    mutated = g.operation_variation(selected[j], O_num, J_num, J, Processing_time, M_num)
                new_pop.append(mutated)

        # 保持种群规模
        if len(new_pop) > g.Pop_size:
            # 对新种群进行筛选
            new_fit = g.fitness(new_pop, J, Processing_time, M_num, O_num)
            new_rank = nsga2.fast_non_dominated_sort(new_fit)
            new_distance = nsga2.crowding_distance(new_fit, new_rank)
            # 按等级和拥挤度排序
            sorted_indices = sorted(range(len(new_pop)), key=lambda k: (new_rank[k], -new_distance[k]))
            C = [new_pop[j] for j in sorted_indices[:g.Pop_size]]
        else:
            C = new_pop

    # 输出结果
    print("\n=== 优化结果 ===")
    print(f"非支配解数量: {len(Optimal_solutions)}")
    print("非支配解目标值 (最大完工时间, 负载标准差):")
    for fit in Optimal_fit_values:
        print(f"({fit[0]}, {fit[1]:.2f})")

    # 选择一个折中解绘制甘特图（例如Cmax最小的解）
    if Optimal_solutions:
        # 找到Cmax最小的解
        cmax_values = [fit[0] for fit in Optimal_fit_values]
        best_idx = cmax_values.index(min(cmax_values))
        Optimal_CHS = Optimal_solutions[best_idx]
        # 解码并绘图
        d = Decode(J, Processing_time, M_num)
        final_fitness = d.decode(Optimal_CHS, O_num)
        print(f"\n选中解的目标值: (最大完工时间: {final_fitness[0]}, 负载标准差: {final_fitness[1]:.2f})")
        Gantt(d.Machines)

    # 绘制帕累托前沿（非支配解）
    if Optimal_fit_values:
        plt.figure(figsize=(10, 6))
        cmax_nd = [fit[0] for fit in Optimal_fit_values]
        load_std_nd = [fit[1] for fit in Optimal_fit_values]
        plt.scatter(cmax_nd, load_std_nd, color='red', label='Non-dominated solutions', zorder=5)

    # 绘制所有解的点图（所有代的所有个体）
    if all_fitnesses:
        cmax_all = [fit[0] for fit in all_fitnesses]
        load_std_all = [fit[1] for fit in all_fitnesses]
        plt.scatter(cmax_all, load_std_all, color='blue', alpha=0.5, label='All solutions', zorder=1)

        plt.title('All Solutions and Pareto Front')
        plt.xlabel('Max Completion Time (Cmax)')
        plt.ylabel('Load Standard Deviation')
        plt.legend()
        plt.grid(True)
        plt.savefig('all_solutions_pareto.png')
        plt.show()