from opt.smm.basis import *
from opt.utils import *


class CellDivisionOpt(BaseOpt):
    def __init__(self, config, part_data, step_data, feeder_data=None):
        super().__init__(config, part_data, step_data, feeder_data)

        self.feeder_assigner = FeederAssignOpt(config, part_data, step_data)

        self.e_gang_pick = 0.6
        self.e_nz_change = 4

    def evaluate(self, part_result, cycle_result, slot_result) -> float:
        nozzle_change_counter = 0
        for head in range(self.config.head_num):
            nozzle = ''
            for cycle in range(len(part_result)):
                component_index = part_result[cycle][head]
                if component_index == -1:
                    continue

                if cycle != 0 and nozzle != self.part_data.loc[component_index, 'nz']:
                    nozzle_change_counter += 1
                nozzle = self.part_data.loc[component_index, 'nz']

        gang_pick_counter = 0
        for cycle, feeder_slot in enumerate(slot_result):
            pick_slot = defaultdict(int)
            for head, slot in enumerate(feeder_slot):
                if slot == -1:
                    continue
                pick_slot[slot - head * (self.config.head_intv / self.config.slot_intv)] += 1
            for _ in pick_slot.values():
                gang_pick_counter += cycle_result[cycle]

        return sum(cycle_result) + self.e_nz_change * nozzle_change_counter + self.e_gang_pick * gang_pick_counter

    def convertor(self, part_cell, population):
        assert part_cell['points'].sum() == len(self.step_data)
        head_num = self.config.head_num
        head_assignment = [[] for _ in range(head_num)]

        wl = [0 for _ in range(head_num)]  # workload

        e1, e2, e3 = 1, 2, 1. / 6

        part_result, cycle_result, feeder_slot_result = [], [], []
        for index in population:
            if part_cell.loc[index]['points'] == 0:
                continue
            # 元胞对应的元件类型和贴装点数
            part_type, part_points = int(part_cell.loc[index, 'index']), int(part_cell.loc[index, 'points'])

            nozzle_change, maxwl = [0 for _ in range(head_num)], [0 for _ in range(head_num)]
            for head in range(head_num):
                if head_assignment[head]:
                    assigned_part = head_assignment[head][-1][0]
                    if self.part_data.loc[assigned_part]['nz'] != self.part_data.loc[part_type]['nz']:
                        nozzle_change[head] = 1
                wl1 = wl.copy()
                wl1[head] += part_points
                maxwl[head] = max(wl1) + e1 * nozzle_change[head]

            awl, wl2 = min(maxwl), wl.copy()
            for idx, val in enumerate(maxwl):
                if val > awl:
                    wl2[idx] += e3
            head_ = wl2.index(min(wl2))
            wl[head_] += part_points
            head_assignment[head_].append([part_type, part_points])

        head_assignment_counter = [0 for _ in range(head_num)]
        while True:
            assigned_part, assigned_cycle = [-1 for _ in range(head_num)], [0 for _ in range(head_num)]
            for head in range(head_num):
                counter = head_assignment_counter[head]

                if head_assignment[head] and head_assignment[head][counter][1] > 0:
                    assigned_part[head] = head_assignment[head][counter][0]
                    assigned_cycle[head] = head_assignment[head][counter][1]

            nonzero_cycle = [cycle for cycle in assigned_cycle if cycle > 0]
            if not nonzero_cycle:
                break

            cycle_result.append(min(nonzero_cycle))
            part_result.append(assigned_part)

            for head in range(head_num):
                counter = head_assignment_counter[head]

                if head_assignment[head] and head_assignment[head][counter][1] > 0:
                    head_assignment[head][counter][1] -= cycle_result[-1]
                    if head_assignment[head][counter][1] == 0 and counter < len(head_assignment[head]) - 1:
                        head_assignment_counter[head] += 1

        slot_result = self.feeder_assigner.do(part_result, cycle_result)
        return part_result, cycle_result, slot_result

    def optimize(self, hinter=True):
        # Crossover method: Two-point crossover
        # Mutation method: Swap
        # Parent selection method: Roulette wheel
        # Termination condition: 20 successive non-improvement iterations
        population_size = 40  # 种群规模
        crossover_rate, mutation_rate = .6, .02
        golden_section = 0.618

        # 获取元件元胞
        part_points = defaultdict(int)
        for _, data in self.step_data.iterrows():
            part_points[data.part] += 1
        feeder_num = sum(self.part_data['fdn'])
        part_cell = pd.DataFrame({'index': np.arange(feeder_num), 'points': np.zeros(feeder_num, dtype=int)})
        cell_index = 0
        for part_index, data in self.part_data.iterrows():
            total_points, div_points = part_points[data.part], math.ceil(part_points[data.part] / data.fdn)
            for _ in range(data.fdn):
                part_cell.loc[cell_index, 'index'] = part_index
                part_cell.loc[cell_index, 'points'] = min(div_points, total_points)
                total_points -= div_points
                cell_index += 1

        part_cell = part_cell[~part_cell['points'].isin([0])]

        # part_cell.sort_values(by = "points" , inplace = True, ascending = False)
        best_population, best_part_cell = [], []
        min_pop_val = float('inf')  # 最优种群价值
        Div, Imp = 0, 0
        while True:
            # randomly generate permutations
            generation_ = np.array(part_cell.index)
            pop_generation = []
            for _ in range(population_size):
                np.random.shuffle(generation_)
                pop_generation.append(generation_.tolist())

            pop_val = []
            for pop in range(population_size):
                part_result, cycle_result, slot_result = self.convertor(part_cell, pop_generation[pop])
                pop_val.append(self.evaluate(part_result, cycle_result, slot_result))

            # 初始化随机生成种群
            Upit = int(1.5 * np.sqrt(len(part_cell)))

            while Div < Upit:
                if hinter:
                    print('----- current div :   ' + str(Div) + ' , total div :   ' + str(Upit) + '  -----')

                # 选择
                new_pop_generation, new_pop_val = [], []
                top_k_index = GenOpe.get_top_kth(pop_val, int(population_size * 0.3))
                for index in top_k_index:
                    new_pop_generation.append(pop_generation[index])
                    new_pop_val.append(pop_val[index])
                index = [i for i in range(population_size)]

                select_index = random.choices(index, weights=pop_val, k=population_size - int(population_size * 0.3))
                for index in select_index:
                    new_pop_generation.append(pop_generation[index])
                    new_pop_val.append(pop_val[index])
                pop_generation, pop_val = new_pop_generation, new_pop_val

                # 交叉
                for pop in range(population_size):
                    if pop % 2 == 0 and np.random.random() < crossover_rate:
                        index1, index2 = GenOpe.roulette_wheel_selection(pop_val), -1
                        while True:
                            index2 = GenOpe.roulette_wheel_selection(pop_val)
                            if index1 != index2:
                                break
                        # 两点交叉算子
                        pop_generation[index1], pop_generation[index2] = GenOpe.partially_mapped_crossover(pop_generation[index1],
                                                                                                           pop_generation[index2])

                    if np.random.random() < mutation_rate:
                        index_ = GenOpe.roulette_wheel_selection(pop_val)
                        GenOpe.swap_mutation(pop_generation[index_])

                # 将元件元胞分配到各个吸杆上，计算价值函数
                for pop in range(population_size):
                    part_result, cycle_result, slot_result = self.convertor(part_cell, pop_generation[pop])
                    pop_val[pop] = self.evaluate(part_result, cycle_result, slot_result)

                    assert pop_val[pop] > 0

                if min(pop_val) < min_pop_val:
                    min_pop_val = min(pop_val)
                    best_population = copy.deepcopy(pop_generation[np.argmin(pop_val)])
                    best_part_cell = copy.deepcopy(part_cell)
                    Div, Imp = 0, 1
                else:
                    Div += 1

            if Imp == 1:
                Div, Imp = 0, 0
                # Section: cell division operation
                if hinter:
                    print(' -------------  cell division operation  ------------- ')
                div_part_cell = pd.DataFrame()
                for idx, rows in part_cell.iterrows():
                    if part_cell.loc[idx, 'points'] <= 1:
                        div_part_cell = pd.concat([div_part_cell, pd.DataFrame([rows])], ignore_index=True)
                    else:
                        div_part_cell = pd.concat([div_part_cell, pd.DataFrame([rows] * 2)], ignore_index=True)

                        rows_counter = len(div_part_cell)
                        div_points = int(max(np.ceil(div_part_cell.loc[rows_counter - 2, 'points'] * golden_section), 1))

                        # 避免出现空元胞的情形
                        if div_points == 0 or div_points == div_part_cell.loc[rows_counter - 2, 'points']:
                            div_part_cell.loc[rows_counter - 2, 'points'] = 1
                        else:
                            div_part_cell.loc[rows_counter - 2, 'points'] = div_points

                        div_part_cell.loc[rows_counter - 1, 'points'] -= div_part_cell.loc[rows_counter - 2, 'points']

                        if div_part_cell.loc[rows_counter - 2, 'points'] == 0 or \
                                div_part_cell.loc[rows_counter - 1, 'points'] == 0:
                            raise ValueError

                part_cell = div_part_cell

                # 完成分裂后重新生成染色体组
                generation_ = np.array(range(len(part_cell)))
                pop_generation = []
                for _ in range(population_size):
                    np.random.shuffle(generation_)
                    pop_generation.append(generation_.tolist())
            else:
                break

        assert len(best_part_cell) == len(best_population)
        self.result.part, self.result.cycle, self.result.slot = self.convertor(best_part_cell, best_population)
        self.result.point, self.result.sequence = self.path_planner.scan_based(self.result.part,
                                                                               self.result.cycle, self.result.slot)
