import copy
import math
import random

import numpy as np

from base_optimizer.optimizer_common import *
from base_optimizer.optimizer_feederpriority import *
from base_optimizer.result_analysis import *


# TODO: nozzle tool available restriction
# TODO: consider with the PCB placement topology
def assembly_time_estimator(assignment_points, component_feeders, component_nozzle):
    nozzle_heads, nozzle_points = defaultdict(int), defaultdict(int)
    for idx, points in enumerate(assignment_points):
        if points == 0:
            continue
        nozzle_points[component_nozzle[idx]] += points
        nozzle_heads[component_nozzle[idx]] = 1

    while sum(nozzle_heads.values()) != max_head_index:
        max_cycle_nozzle = None

        for nozzle, head_num in nozzle_heads.items():
            if max_cycle_nozzle is None or nozzle_points[nozzle] / head_num > nozzle_points[max_cycle_nozzle] / \
                    nozzle_heads[max_cycle_nozzle]:
                max_cycle_nozzle = nozzle

        assert max_cycle_nozzle is not None
        nozzle_heads[max_cycle_nozzle] += 1

    head_nozzle_assignment, min_cost = None, None

    # generate initial nozzle group
    nozzle_group = []
    # averagely assign for the same type of nozzles, and generate nozzle group
    nozzle_points_cpy = copy.deepcopy(nozzle_points)
    for nozzle, heads in nozzle_heads.items():
        points = nozzle_points_cpy[nozzle] // heads
        for _ in range(heads):
            nozzle_group.append([nozzle, points])
        nozzle_points_cpy[nozzle] -= heads * points

    for idx, [nozzle, _] in enumerate(nozzle_group):
        if nozzle_points_cpy[nozzle]:
            nozzle_group[idx][1] += 1
            nozzle_points_cpy[nozzle] -= 1

    while True:
        # assign nozzle group to each head
        nozzle_group.sort(key=lambda x: -x[1])

        tmp_head_nozzle_assignment = []
        head_total_points = [0 for _ in range(max_head_index)]
        for idx, nozzle_item in enumerate(nozzle_group):
            if idx < max_head_index:
                tmp_head_nozzle_assignment.append([nozzle_item.copy()])
                head_total_points[idx] += nozzle_item[1]
            else:
                min_head = np.argmin(head_total_points)
                tmp_head_nozzle_assignment[min_head].append(nozzle_item.copy())
                head_total_points[min_head] += nozzle_item[1]

        cost = t_cycle * max(head_total_points)
        for head in range(max_head_index):
            for cycle in range(len(tmp_head_nozzle_assignment[head])):
                if cycle + 1 == len(tmp_head_nozzle_assignment[head]):
                    if tmp_head_nozzle_assignment[head][cycle][0] != tmp_head_nozzle_assignment[head][-1][0]:
                        cost += t_nozzle_change
                else:
                    if tmp_head_nozzle_assignment[head][cycle][0] != tmp_head_nozzle_assignment[head][cycle + 1][0]:
                        cost += t_nozzle_change

        while True:
            min_head, max_head = np.argmin(head_total_points), np.argmax(head_total_points)
            min_head_nozzle, max_head_nozzle = tmp_head_nozzle_assignment[min_head][-1][0], \
                                               tmp_head_nozzle_assignment[max_head][-1][0]
            if min_head_nozzle == max_head_nozzle:
                break

            min_head_list, max_head_list = [min_head], [max_head]
            minmax_head_points = 0
            for head in range(max_head_index):
                if head in min_head_list or head in max_head_list:
                    minmax_head_points += head_total_points[head]
                    continue

                # the max/min heads with the sum nozzle type
                if tmp_head_nozzle_assignment[head][-1][0] == tmp_head_nozzle_assignment[min_head][-1][0]:
                    min_head_list.append(head)
                    minmax_head_points += head_total_points[head]

                if tmp_head_nozzle_assignment[head][-1][0] == tmp_head_nozzle_assignment[max_head][-1][0]:
                    max_head_list.append(head)
                    minmax_head_points += head_total_points[head]

            # todo: restriction of available nozzle
            # the reduction of cycles is not offset the cost of nozzle change
            average_points = minmax_head_points // (len(min_head_list) + len(max_head_list))
            reminder_points = minmax_head_points % (len(min_head_list) + len(max_head_list))
            max_cycle = average_points + (1 if reminder_points > 0 else 0)
            for head in range(max_head_index):
                if head in min_head_list or head in max_head_list:
                    continue
                max_cycle = max(max_cycle, head_total_points[head])

            nozzle_change_counter = 0
            for head in min_head_list:
                if tmp_head_nozzle_assignment[head][0] == tmp_head_nozzle_assignment[head][-1]:
                    nozzle_change_counter += 2
                else:
                    nozzle_change_counter += 1

            if t_cycle * (max(head_total_points) - max_cycle) < t_nozzle_change * nozzle_change_counter:
                break

            cost -= t_cycle * (max(head_total_points) - max_cycle) - t_nozzle_change * nozzle_change_counter

            required_points = 0  # 待均摊的贴装点数较多的吸嘴类型
            for head in min_head_list:
                points = average_points - head_total_points[head]
                tmp_head_nozzle_assignment[head].append([max_head_nozzle, points])
                head_total_points[head] = average_points
                required_points += points

            for head in max_head_list:
                tmp_head_nozzle_assignment[head][-1][1] -= required_points // len(max_head_list)
                head_total_points[head] -= required_points // len(max_head_list)

            required_points -= (required_points // len(max_head_list)) * len(max_head_list)

            for head in max_head_list:
                if required_points <= 0:
                    break
                tmp_head_nozzle_assignment[head][-1][1] -= 1
                head_total_points[head] -= 1
                required_points -= 1

        if min_cost is None or cost < min_cost:
            min_cost = cost
            head_nozzle_assignment = copy.deepcopy(tmp_head_nozzle_assignment)
        else:
            break

        # 在吸嘴组中增加一个吸嘴
        idx, nozzle = 0, nozzle_group[0][0]
        for idx, [nozzle_, _] in enumerate(nozzle_group):
            if nozzle_ != nozzle:
                break

        average_points, remainder_points = nozzle_points[nozzle] // (idx + 1), nozzle_points[nozzle] % (idx + 1)
        nozzle_group.append([nozzle, 0])
        for idx, [nozzle_, _] in enumerate(nozzle_group):
            if nozzle_ == nozzle:
                nozzle_group[idx][1] = average_points + (1 if remainder_points > 0 else 0)
                remainder_points -= 1

    cycle_counter, nozzle_change_counter = 0, 0
    for head in range(max_head_index):
        head_cycle_counter = 0
        for cycle in range(len(head_nozzle_assignment[head])):
            if cycle + 1 == len(head_nozzle_assignment[head]):
                if head_nozzle_assignment[head][0][0] != head_nozzle_assignment[head][-1][0]:
                    nozzle_change_counter += 1
            else:
                if head_nozzle_assignment[head][cycle][0] != head_nozzle_assignment[head][cycle + 1][0]:
                    nozzle_change_counter += 1
            head_cycle_counter += head_nozzle_assignment[head][cycle][1]
        cycle_counter = max(cycle_counter, head_cycle_counter)

    # === 元件拾取次数预估 ===
    cp_info = []
    for idx, points in enumerate(assignment_points):
        if points == 0:
            continue
        reminder_points = points % component_feeders[idx]
        for _ in range(component_feeders[idx]):
            cp_info.append(
                [idx, points // component_feeders[idx] + (1 if reminder_points > 0 else 0), component_nozzle[idx]])
            reminder_points -= 1

    cp_info.sort(key=lambda x: -x[1])

    nozzle_level, nozzle_counter = defaultdict(int), defaultdict(int)
    level_points = defaultdict(int)
    for info in cp_info:
        nozzle = info[2]
        if nozzle_counter[nozzle] and nozzle_counter[nozzle] % nozzle_heads[nozzle] == 0:
            nozzle_level[nozzle] += 1
        level = nozzle_level[nozzle]
        level_points[level] = max(level_points[level], info[1])
        nozzle_counter[nozzle] += 1

    pickup_counter = sum(points for points in level_points.values())
    placement_counter = sum(assignment_points)

    pickup_movement = 0
    for points in assignment_points:
        if points:
            pickup_movement += 1
    # 返回加权预估时间
    return t_cycle * cycle_counter + t_nozzle_change * nozzle_change_counter + t_pick * pickup_counter + \
        t_place * placement_counter + 0.1 * pickup_movement


def assemblyline_optimizer_heuristic(pcb_data, component_data, machine_number):
    # the number of placement points, the number of available feeders, and nozzle type of component respectively
    component_number = len(component_data)

    component_points = [0 for _ in range(component_number)]
    component_feeders = [0 for _ in range(component_number)]
    component_nozzle = [0 for _ in range(component_number)]
    component_part = [0 for _ in range(component_number)]

    nozzle_points = defaultdict(int)        # the number of placements of nozzle

    for _, data in pcb_data.iterrows():
        part_index = component_data[component_data['part'] == data['part']].index.tolist()[0]
        nozzle = component_data.loc[part_index]['nz']

        component_points[part_index] += 1
        component_feeders[part_index] = component_data.loc[part_index]['feeder-limit']
        # component_feeders[part_index] = math.ceil(component_data.loc[part_index]['feeder-limit'] / max_feeder_limit)
        component_nozzle[part_index] = nozzle
        component_part[part_index] = data['part']

        nozzle_points[nozzle] += 1

    # first step: generate the initial solution with equalized workload
    assignment_result = [[0 for _ in range(len(component_points))] for _ in range(machine_number)]
    assignment_points = [0 for _ in range(machine_number)]
    average_points = len(pcb_data) // machine_number

    weighted_points = list(
        map(lambda x: x[1] + 1e-5 * nozzle_points[component_nozzle[x[0]]], enumerate(component_points)))

    # for part_index in np.argsort(weighted_points)[::-1]:
    for part_index in np.argsort(weighted_points)[::-1]:
        if (total_points := component_points[part_index]) == 0:        # total placements for each component type
            continue
        machine_set = []

        # define the machine that assigning placement points (considering the feeder limitation)
        for machine_index in np.argsort(assignment_points):
            if len(machine_set) >= component_points[part_index] or len(machine_set) >= component_feeders[part_index]:
                break

            machine_set.append(machine_index)
            if weighted_points[part_index] + assignment_points[machine_index] < average_points:
                break

        # Allocation of mounting points to available machines according to the principle of equality
        while total_points:
            assign_machine = list(filter(lambda x: assignment_points[x] == min(assignment_points), machine_set))

            if len(assign_machine) == len(machine_set):
                # averagely assign point to all available machines
                points = total_points // len(assign_machine)
                for machine_index in machine_set:
                    assignment_points[machine_index] += points
                    assignment_result[machine_index][part_index] += points

                total_points -= points * len(assign_machine)
                for machine_index in machine_set:
                    if total_points == 0:
                        break
                    assignment_points[machine_index] += 1
                    assignment_result[machine_index][part_index] += 1
                    total_points -= 1
            else:
                # assigning placements to make up for the gap between the least and the second least
                second_least_machine, second_least_machine_points = -1, max(assignment_points) + 1
                for idx in machine_set:
                    if assignment_points[idx] < second_least_machine_points and assignment_points[idx] != min(
                            assignment_points):
                        second_least_machine_points = assignment_points[idx]
                        second_least_machine = idx

                assert second_least_machine != -1

                if len(assign_machine) * (second_least_machine_points - min(assignment_points)) < total_points:
                    min_points = min(assignment_points)
                    total_points -= len(assign_machine) * (second_least_machine_points - min_points)
                    for machine_index in assign_machine:
                        assignment_points[machine_index] += (second_least_machine_points - min_points)
                        assignment_result[machine_index][part_index] += (
                                    second_least_machine_points - min_points)
                else:
                    points = total_points // len(assign_machine)
                    for machine_index in assign_machine:
                        assignment_points[machine_index] += points
                        assignment_result[machine_index][part_index] += points

                    total_points -= points * len(assign_machine)
                    for machine_index in assign_machine:
                        if total_points == 0:
                            break
                        assignment_points[machine_index] += 1
                        assignment_result[machine_index][part_index] += 1
                        total_points -= 1

    prev_max_assembly_time, prev_assignment_result = None, None
    while True:
        # second step: estimate the assembly time for each machine
        arranged_feeders = defaultdict(list)
        for machine_index in range(machine_number):
            arranged_feeders[machine_index] = [0 for _ in range(len(component_data))]

        for part_index in range(len(component_data)):
            feeder_limit = component_feeders[part_index]    # 总体可用数
            for machine_index in range(machine_number):
                if assignment_result[machine_index][part_index] == 0:
                    continue
                feeder_limit -= 1
                # 已分配元件的机器至少安装1把供料器
                arranged_feeders[machine_index][part_index] = 1
                assert feeder_limit >= 0

        for part_index in range(len(component_data)):
            total_feeder_limit = component_feeders[part_index] - sum(
                [arranged_feeders[machine_index][part_index] for machine_index in range(machine_number)])
            while total_feeder_limit > 0:
                max_ratio, max_ratio_machine = None, -1
                for machine_index in range(machine_number):
                    if assignment_result[machine_index][part_index] == 0:
                        continue
                    ratio = assignment_result[machine_index][part_index] / arranged_feeders[machine_index][part_index]
                    if max_ratio is None or ratio > max_ratio:
                        max_ratio, max_ratio_machine = ratio, machine_index
                assert max_ratio_machine is not None
                arranged_feeders[max_ratio_machine][part_index] += 1
                total_feeder_limit -= 1

        assembly_time, chip_per_hour = [], []
        for machine_index in range(machine_number):
            assembly_time.append(
                assembly_time_estimator(assignment_result[machine_index], arranged_feeders[machine_index],
                                        component_nozzle))
            chip_per_hour.append(sum(assignment_result[machine_index]) / (assembly_time[-1] + 1e-10))

        max_assembly_time = max(assembly_time)
        if prev_max_assembly_time and (prev_max_assembly_time < max_assembly_time or abs(
                max_assembly_time - prev_max_assembly_time) < 1e-10):
            if prev_max_assembly_time < max_assembly_time:
                assignment_result = copy.deepcopy(prev_assignment_result)
            break
        else:
            prev_max_assembly_time = max_assembly_time
            prev_assignment_result = copy.deepcopy(assignment_result)

        # third step: adjust the assignment results to reduce maximal assembly time among all machines
        # ideal averagely assigned points
        total_points = len(pcb_data)
        average_assign_points = [round(total_points * chip_per_hour[mi] / sum(chip_per_hour)) for mi in
                                 range(machine_number)]

        machine_index = 0
        while total_points != sum(average_assign_points):
            if total_points > sum(average_assign_points):
                average_assign_points[machine_index] += 1
            else:
                average_assign_points[machine_index] -= 1
            machine_index += 1
            if machine_index >= machine_number:
                machine_index = 0

        # the placement points that need to be re-allocated
        machine_reallocate_points = [sum(assignment_result[mi]) - average_assign_points[mi] for mi in
                                     range(machine_number)]

        # workload balance
        # 1. balance the number of placements of the same type between different machines.
        for demand_mi in range(machine_number):
            if machine_reallocate_points[demand_mi] >= 0:
                continue

            supply_machine_list = [mi for mi in range(machine_number) if machine_reallocate_points[mi] > 0]
            supply_machine_list.sort(key=lambda mi: -machine_reallocate_points[mi])

            for supply_mi in supply_machine_list:
                for part_index in range(len(component_data)):
                    if assignment_result[supply_mi][part_index] <= 0:
                        continue

                    reallocate_points = min(assignment_result[supply_mi][part_index],
                                            -machine_reallocate_points[demand_mi])

                    # upper available feeder restrictions
                    tmp_reallocate_result = [assignment_result[mi][part_index] for mi in range(machine_number)]
                    tmp_reallocate_result[supply_mi] -= reallocate_points
                    tmp_reallocate_result[demand_mi] += reallocate_points

                    if sum(1 for pt in tmp_reallocate_result if pt > 0) > component_feeders[part_index]:
                        continue

                    assignment_result[supply_mi][part_index] -= reallocate_points
                    machine_reallocate_points[supply_mi] -= reallocate_points
                    assignment_result[demand_mi][part_index] += reallocate_points
                    machine_reallocate_points[demand_mi] += reallocate_points

                    if machine_reallocate_points[demand_mi] <= 0:
                        break

        # 2. balance the number of placements of the different type between different machines.
        cp_info = []
        for part_index in range(len(component_data)):
            for machine_index in range(machine_number):
                if assignment_result[machine_index][part_index] == 0:
                    continue
                cp_info.append([machine_index, part_index, assignment_result[machine_index][part_index]])

        for machine_index in range(machine_number):
            if machine_reallocate_points[machine_index] >= 0:
                continue

            filter_cp_info = [info for info in cp_info if
                              info[0] != machine_index and machine_reallocate_points[info[0]] > 0]
            while True:
                if len(filter_cp_info) == 0 or machine_reallocate_points[machine_index] >= 0:
                    break
                # todo: 对同时拾取数的影响
                filter_cp_info.sort(key=lambda x: x[2] + machine_reallocate_points[machine_index])

                info = filter_cp_info[0]
                filter_cp_info.remove(info)
                if abs(machine_reallocate_points[machine_index]) + abs(machine_reallocate_points[info[0]]) < abs(
                        machine_reallocate_points[machine_index] + info[2]) + abs(
                        machine_reallocate_points[info[0]] - info[2]):
                    continue

                cp_info.remove(info)

                assignment_result[info[0]][info[1]] = 0
                assignment_result[machine_index][info[1]] += info[2]

                machine_reallocate_points[info[0]] -= info[2]
                machine_reallocate_points[machine_index] += info[2]

    return assignment_result