import matplotlib.pyplot as plt
import numpy as np
from chromosome_utils import ChromosomeUtils  # 染色体工具类


class ResultVisualizer:
    """结果可视化工具类：绘制帕累托前沿、收敛曲线，打印最优解详情"""

    def __init__(self, utils: ChromosomeUtils):
        """
        初始化可视化工具
        :param utils: 染色体工具类实例（用于解析染色体）
        """
        self.utils = utils
        self.figsize = (12, 8)  # 图表大小
        # 企业名称列表（风险企业+供应商）
        self.supplier_names = ["RiskEnterprise"] + self.utils.supplier.names

    def plot_pareto_front(self, all_objectives: list[tuple[float, float]],
                          non_dominated_objectives: list[tuple[float, float]]):
        """
        绘制帕累托前沿（所有解 vs 帕累托最优解）
        :param all_objectives: 所有解的目标值
        :param non_dominated_objectives: 帕累托最优解的目标值
        """
        plt.figure(figsize=self.figsize)
        # 绘制所有解（蓝色点）
        if all_objectives:
            c_all = [obj[0] for obj in all_objectives]  # 成本
            t_all = [obj[1] for obj in all_objectives]  # 延期
            plt.scatter(c_all, t_all, color='blue', alpha=0.3, label='All Solutions', zorder=1)
        # 绘制帕累托前沿（红色点）
        if non_dominated_objectives:
            c_nd = [obj[0] for obj in non_dominated_objectives]
            t_nd = [obj[1] for obj in non_dominated_objectives]
            plt.scatter(c_nd, t_nd, color='red', s=50, label='Pareto Front', zorder=5)
        # 图表设置
        plt.title('Pareto Front: Change Cost vs Tardiness')
        plt.xlabel('Change Cost')
        plt.ylabel('Tardiness')
        plt.legend()
        plt.grid(True, alpha=0.5)
        plt.savefig('pareto_front.png', dpi=300, bbox_inches='tight')  # 保存图片
        plt.show()

    def plot_convergence(self, convergence_history: list[tuple[float, float]]):
        """
        绘制收敛趋势图（每代最优前沿的平均成本和延期）
        :param convergence_history: 收敛历史数据（(平均成本, 平均延期)）
        """
        if len(convergence_history) == 0:  # 无数据时不绘制
            return
        plt.figure(figsize=self.figsize)
        generations = list(range(len(convergence_history)))  # 代数
        costs = [h[0] for h in convergence_history]  # 平均成本
        tardiness = [h[1] for h in convergence_history]  # 平均延期

        # 子图1：平均成本趋势
        plt.subplot(2, 1, 1)
        plt.plot(generations, costs, 'b-')
        plt.title('Convergence History')
        plt.ylabel('Average Change Cost')
        plt.grid(True, alpha=0.5)

        # 子图2：平均延期趋势
        plt.subplot(2, 1, 2)
        plt.plot(generations, tardiness, 'r-')
        plt.xlabel('Generation')
        plt.ylabel('Average Tardiness')
        plt.grid(True, alpha=0.5)

        plt.tight_layout()  # 调整布局
        plt.savefig('convergence_history.png', dpi=300, bbox_inches='tight')  # 保存图片
        plt.show()

    def print_solution_details(self, solution: np.ndarray, objectives: tuple[float, float]):
        """
        打印单个解的详细信息（企业选择、产能、数量等）
        :param solution: 染色体（解）
        :param objectives: 该解的目标值（成本, 延期）
        """
        # 拆分染色体为三层
        enterprise_layer, capacity_layer, quantity_layer = self.utils._split_chromosome(solution)
        split_points = np.cumsum(self.utils.material_enterprise_count)  # 分割点

        print(f"\n变更成本: {objectives[0]:.2f}, 交付延期: {objectives[1]:.2f}")
        print("=" * 80)
        total_q_check = []  # 检查各物料的总数量是否满足需求

        start = 0
        for i in range(self.utils.I):  # 遍历每种物料
            end = split_points[i]
            ents = self.utils.material_optional_enterprises[i]  # 可选企业
            e_segment = enterprise_layer[start:end]  # 企业选择状态
            c_segment = capacity_layer[start:end]  # 产能
            q_segment = quantity_layer[start:end]  # 数量

            demand_q = self.utils.order.Q[i]  # 需求数量
            allocated_q = np.sum(q_segment[e_segment == 1])  # 分配的总数量
            print(f"物料 {i} - 需求数量: {demand_q}, 分配总量: {allocated_q:.2f}")
            total_q_check.append(abs(allocated_q - demand_q) < 1e-2)  # 检查是否满足需求

            print(f"  选择的企业及其分配:")
            for idx, ent in enumerate(ents):
                if e_segment[idx] == 1:  # 仅打印被选中的企业
                    ent_name = self.supplier_names[ent]  # 企业名称
                    cap = c_segment[idx]  # 产能
                    qty = q_segment[idx]  # 数量
                    print(f"    企业: {ent_name}, 产能: {cap:.2f}, 分配数量: {qty:.2f}")

            start = end
            print("-" * 80)

        # 验证数量约束是否满足
        if all(total_q_check):
            print("✅ 所有物料数量满足需求约束")
        else:
            print("❌ 部分物料数量未满足需求约束")

    def print_pareto_solutions(self, population: np.ndarray, objectives: list[tuple[float, float]]):
        """
        打印帕累托前沿解的详细信息（最多打印前5个）
        :param population: 帕累托前沿解的种群
        :param objectives: 对应的目标值列表
        """
        print("\n" + "=" * 100)
        print("帕累托前沿解详细方案")
        print("=" * 100)
        # 打印前5个解（或所有解，取较少者）
        for i in range(min(5, len(population))):
            print(f"\n===== 最优解 {i + 1} =====")
            self.print_solution_details(population[i], objectives[i])