import random
import numpy as np
from data_structures import Config
from chromosome_utils import ChromosomeUtils


class GeneticOperator:
    """遗传操作：交叉、变异"""
    def __init__(self, config: Config, utils: ChromosomeUtils):
        self.config = config
        self.utils = utils

    def two_point_crossover(self, parent1: np.ndarray, parent2: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """两点交叉"""
        length = self.utils.chromosome_length
        if length < 2:
            return parent1.copy(), parent2.copy()
        point1 = random.randint(0, length//2)
        point2 = random.randint(point1+1, length-1)
        child1 = parent1.copy()
        child2 = parent2.copy()
        child1[point1:point2] = parent2[point1:point2]
        child2[point1:point2] = parent1[point1:point2]
        child1 = self.utils.repair_chromosome(child1)
        child2 = self.utils.repair_chromosome(child2)
        return child1, child2

    def uniform_mutation(self, chromosome: np.ndarray) -> np.ndarray:
        """均匀变异：遍历三层每个基因"""
        mutated = chromosome.copy()
        enterprise_layer, capacity_layer, quantity_layer = self.utils._split_chromosome(mutated)
        split_points = np.cumsum(self.utils.material_enterprise_count)
        # 企业层变异（0/1翻转）
        start = 0
        for i in range(self.utils.I):
            end = split_points[i]
            for idx in range(start, end):
                if random.random() < self.config.mutation_prob:
                    enterprise_layer[idx] = 1 - enterprise_layer[idx]
            start = end
        enterprise_layer = self.utils.repair_enterprise_layer(enterprise_layer)
        # 能力层变异
        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]
            for idx in range(start, end):
                if random.random() < self.config.mutation_prob and e_segment[idx-start] == 1:
                    ent = ents[idx-start]
                    if ent == 0:
                        max_cap = self.utils.risk.C0_i_std[i]
                    else:
                        supplier_id = ent - 1
                        max_cap = self.utils.supplier.Cj_i_std[supplier_id][i]
                    capacity_layer[idx] = random.uniform(1, max_cap)
            start = end
        capacity_layer = self.utils.repair_capacity_layer(enterprise_layer, capacity_layer)
        # 数量层变异
        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]
            for idx in range(start, end):
                if random.random() < self.config.mutation_prob and e_segment[idx-start] == 1:
                    ent = ents[idx-start]
                    qi = self.utils.order.Q[i]
                    if ent == 0:
                        max_q = qi
                    else:
                        supplier_id = ent - 1
                        max_q = self.supplier.MaxOrder[supplier_id][i]
                    quantity_layer[idx] = random.uniform(1, max_q)
            start = end
        quantity_layer = self.utils.repair_quantity_layer(enterprise_layer, quantity_layer)
        # 合并并修复
        mutated = self.utils._merge_chromosome(enterprise_layer, capacity_layer, quantity_layer)
        return mutated