import numpy as np
from data_structures import OrderData, RiskEnterpriseData, SupplierData, Config
from chromosome_utils import ChromosomeUtils


class ObjectiveCalculator:
    """目标函数计算器：变更成本C + 交付延期T"""

    def __init__(self, order_data: OrderData, risk_data: RiskEnterpriseData, supplier_data: SupplierData,
                 utils: ChromosomeUtils, config: Config):
        self.order = order_data
        self.risk = risk_data
        self.supplier = supplier_data
        self.utils = utils
        self.config = config
        self.split_points = np.cumsum(utils.material_enterprise_count)

    def calculate_objectives(self, chromosome: np.ndarray) -> tuple[float, float]:
        """计算双目标值：(变更成本C, 交付延期T)"""
        enterprise_layer, capacity_layer, quantity_layer = self.utils._split_chromosome(chromosome)
        # 修复：传入capacity_layer参数
        C = self._calculate_change_cost(enterprise_layer, capacity_layer, quantity_layer)
        T = self._calculate_tardiness(enterprise_layer, capacity_layer, quantity_layer)
        return C, T

    # 修复：添加capacity_layer参数
    def _calculate_change_cost(self, enterprise_layer: np.ndarray, capacity_layer: np.ndarray, quantity_layer: np.ndarray) -> float:
        """计算变更成本C = C1 + C2 + C3 + C4"""
        C1 = 0.0  # 变更惩罚成本
        C2 = 0.0  # 采购变更成本
        C3 = 0.0  # 运输变更成本
        C4 = 0.0  # 提前交付惩罚成本

        # 原始成本（全部由风险企业生产）
        original_purchase_cost = sum(self.order.Q[i] * self.order.P0[i] for i in range(self.order.I))
        original_transport_cost = sum(self.order.Q[i] * self.order.T0[i] for i in range(self.order.I))

        # 变更后成本
        new_purchase_cost = 0.0
        new_transport_cost = 0.0
        risk_production = np.zeros(self.order.I)  # 风险企业生产数量
        supplier_production = np.zeros(self.order.I)  # 供应商生产数量

        start = 0
        for i in range(self.order.I):
            end = self.split_points[i]
            ents = self.utils.material_optional_enterprises[i]
            e_segment = enterprise_layer[start:end]
            q_segment = quantity_layer[start:end]

            for idx, ent in enumerate(ents):
                if e_segment[idx] == 1:
                    q = q_segment[idx]
                    if ent == 0:
                        # 风险企业
                        risk_production[i] += q
                        new_purchase_cost += q * self.order.P0[i]
                        new_transport_cost += q * self.order.T0[i]
                    else:
                        # 供应商
                        supplier_id = ent - 1
                        supplier_production[i] += q
                        new_purchase_cost += q * self.supplier.P_ij[supplier_id][i]
                        new_transport_cost += q * self.supplier.T_ij[supplier_id][i]
            start = end

        # 计算C1：变更惩罚成本
        for i in range(self.order.I):
            C1 += self.config.delta * supplier_production[i] * (self.order.P0[i] + self.order.T0[i])

        # 计算C2：采购变更成本
        C2 = new_purchase_cost - original_purchase_cost

        # 计算C3：运输变更成本
        C3 = new_transport_cost - original_transport_cost

        # 计算C4：提前交付惩罚成本
        # 修复：传入capacity_layer参数
        actual_delivery_time = self._calculate_actual_delivery_time(enterprise_layer, capacity_layer, quantity_layer)
        if actual_delivery_time < self.order.Dd:
            # 计算风险企业和各供应商的交货时间
            risk_delivery = []
            supplier_deliveries = {}
            start = 0
            for i in range(self.order.I):
                end = self.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]

                for idx, ent in enumerate(ents):
                    if e_segment[idx] == 1:
                        q = q_segment[idx]
                        c = c_segment[idx]
                        if c == 0:
                            production_time = 0
                        else:
                            production_time = q / c

                        if ent == 0:
                            transport_time = self.risk.distance / self.order.transport_speed
                            risk_delivery.append(production_time + transport_time)
                        else:
                            supplier_id = ent - 1
                            transport_time = self.supplier.distance[supplier_id] / self.order.transport_speed
                            if supplier_id not in supplier_deliveries:
                                supplier_deliveries[supplier_id] = []
                            supplier_deliveries[supplier_id].append(production_time + transport_time)
                start = end

            D0 = max(risk_delivery) if risk_delivery else 0
            Dj_sum = sum(max(times) for times in supplier_deliveries.values()) if supplier_deliveries else 0
            C4 = self.config.gamma * ((self.order.Dd - D0) + Dj_sum)

        return C1 + C2 + C3 + C4

    def _calculate_actual_delivery_time(self, enterprise_layer: np.ndarray, capacity_layer: np.ndarray,
                                        quantity_layer: np.ndarray) -> float:
        """计算实际交货期：所有企业中最长的生产+运输时间"""
        max_time = 0.0
        start = 0
        for i in range(self.order.I):
            end = self.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]

            for idx, ent in enumerate(ents):
                if e_segment[idx] == 1:
                    q = q_segment[idx]
                    c = c_segment[idx]
                    if c == 0:
                        production_time = 0
                    else:
                        production_time = q / c

                    # 计算运输时间
                    if ent == 0:
                        transport_time = self.risk.distance / self.order.transport_speed
                    else:
                        supplier_id = ent - 1
                        transport_time = self.supplier.distance[supplier_id] / self.order.transport_speed

                    total_time = production_time + transport_time
                    if total_time > max_time:
                        max_time = total_time
            start = end
        return max_time

    def _calculate_tardiness(self, enterprise_layer: np.ndarray, capacity_layer: np.ndarray,
                             quantity_layer: np.ndarray) -> float:
        """计算交付延期T = max(0, 实际交货期 - 需求交货期)"""
        actual_delivery = self._calculate_actual_delivery_time(enterprise_layer, capacity_layer, quantity_layer)
        return max(0.0, actual_delivery - self.order.Dd)