import copy
import math

import numpy as np

from opt.utils import axis_moving_time
from core.common import *
from data.type import Point


class PathPlanOpt:
    def __init__(self, config, part_data, step_data):
        self.part_data = part_data
        self.step_data = step_data
        self.config = config

    def dynamic_programming_cycle_path(self, cycle_point, cycle_slot):
        head_sequence = []
        num_pos = sum([placement != -1 for placement in cycle_point]) + 1
        intv_ratio = self.config.head_intv // self.config.slot_intv

        pos, head_set = [], []
        feeder_set = set()
        for head, slot in enumerate(cycle_slot):
            if slot == -1:
                continue

            head_set.append(head)
            placement = cycle_point[head]
            pos.append([self.step_data.loc[placement]['x'] - head * self.config.head_intv + self.config.stopper_pos.x,
                        self.step_data.loc[placement]['y'] + self.config.stopper_pos.y,
                        self.step_data.loc[placement]['r'], head])

            feeder_set.add(slot - head * intv_ratio)

        pos.insert(0, [self.config.slotf1_pos.x + ((min(list(feeder_set)) + max(list(feeder_set))) / 2 - 1) *
                       self.config.slot_intv, self.config.slotf1_pos.y, None, 0])

        def get_distance(pos_1, pos_2):
            # 拾取起始与终止位置 或 非同轴
            if pos_1[2] is None or pos_2[2] is None or pos_1[3] + (1 if pos_1[3] % 2 == 0 else -1) != pos_2[3]:
                return max(axis_moving_time(pos_1[0] - pos_2[0], 0), axis_moving_time(pos_1[1] - pos_2[1], 1))
            else:
                return max(axis_moving_time(pos_1[0] - pos_2[0], 0), axis_moving_time(pos_1[1] - pos_2[1], 1),
                           axis_moving_time(pos_1[2] - pos_2[2], 2))

        # 各节点之间的距离
        dist = [[get_distance(pos_1, pos_2) for pos_2 in pos] for pos_1 in pos]

        min_dist = [[np.inf for _ in range(num_pos)] for s in range(1 << num_pos)]
        min_path = [[[] for _ in range(num_pos)] for s in range(1 << num_pos)]

        # 状压dp搜索
        for s in range(1, 1 << num_pos, 2):
            # 考虑节点集合s必须包括节点0
            if not (s & 1):
                continue
            for j in range(1, num_pos):
                # 终点j需在当前考虑节点集合s内
                if not (s & (1 << j)):
                    continue
                if s == int((1 << j) | 1):
                    # 若考虑节点集合s仅含节点0和节点j，dp边界，赋予初值
                    # print('j:', j)
                    min_path[s][j] = [j]
                    min_dist[s][j] = dist[0][j]

                # 枚举下一个节点i，更新
                for i in range(1, num_pos):
                    # 下一个节点i需在考虑节点集合s外
                    if s & (1 << i):
                        continue
                    if min_dist[s][j] + dist[j][i] < min_dist[s | (1 << i)][i]:
                        min_path[s | (1 << i)][i] = min_path[s][j] + [i]
                        min_dist[s | (1 << i)][i] = min_dist[s][j] + dist[j][i]

        ans_dist = float('inf')
        ans_path = []
        # 求最终最短哈密顿回路
        for i in range(1, num_pos):
            if min_dist[(1 << num_pos) - 1][i] + dist[i][0] < ans_dist:
                # 更新，回路化
                ans_path = min_path[s][i]
                ans_dist = min_dist[(1 << num_pos) - 1][i] + dist[i][0]

        for parent in ans_path:
            head_sequence.append(head_set[parent - 1])

        start_head, end_head = head_sequence[0], head_sequence[-1]
        if self.step_data.loc[cycle_point[start_head]]['x'] - start_head * self.config.head_intv > \
                self.step_data.loc[cycle_point[end_head]]['x'] - end_head * self.config.head_intv:
            head_sequence = list(reversed(head_sequence))
        return ans_dist, head_sequence

    def scan_based(self, part_result, cycle_result, slot_result, hinter=True):
        point_result, sequence_result = [], []

        class Mount:
            def __init__(self):
                self.pos = []

                self.part = []
                self.step = []

            def pop(self, index):
                self.pos.pop(index)

                self.part.pop(index)
                self.step.pop(index)

        all_points = Mount()
        for step_index, data in self.step_data.iterrows():
            part_index = self.part_data[self.part_data.part == data.part].index.tolist()[0]

            # 记录贴装点序号索引和对应的位置坐标
            all_points.pos.append(Point(data.x + self.config.stopper_pos.x, data.y + self.config.stopper_pos.y, data.r))
            all_points.part.append(part_index)
            all_points.step.append(step_index)

        point_list = list(range(0, self.step_data.shape[0]))

        head_num = self.config.head_num
        left_boundary, right_boundary = min(all_points.pos, key=lambda p: p.x).x, \
                                        max(all_points.pos, key=lambda p: p.x).x
        search_step = max((right_boundary - left_boundary) / head_num / 2, 0)

        search_points, head_range = defaultdict(lambda: np.array([])), defaultdict(list)

        search_points[0] = np.arange(left_boundary, (left_boundary + right_boundary) / 2, search_step)
        head_range[0] = list(range(head_num))

        search_points[1] = np.arange(right_boundary + 1e-3, (left_boundary + right_boundary) / 2, -search_step)
        head_range[1] = list(range(head_num - 1, -1, -1))

        search_points[2] = np.arange(left_boundary, right_boundary, search_step)
        head_range[2], head_index = [], (head_num - 1) // 2
        while head_index >= 0:
            if 2 * head_index != head_num - 1:
                head_range[2].append(head_num - 1 - head_index)
            head_range[2].append(head_index)
            head_index -= 1

        pbar = tqdm(total=sum(cycle_result), desc='scan-based path plan process') if hinter else None

        ref_pos_y = min(all_points.pos, key=lambda p: p.y).y
        for cycle_index, component_cycle in enumerate(part_result):
            for _ in range(cycle_result[cycle_index]):
                min_dist = np.inf
                best_head_point, best_head_seq, best_point_list = [], [], []
                for dir in range(3):     # 不同的搜索方向，贴装头和起始点的选取方法各不相同
                    for start_points in search_points[dir]:
                        cur_point_list = point_list.copy()

                        head_point = [-1] * head_num
                        assigned_point = [Point(0, 0)] * head_num
                        head_counter, point_index = 0, -1
                        for head_index in head_range[dir]:
                            if head_counter == 0:
                                if part_result[cycle_index][head_index] == -1:
                                    continue

                                min_horizontal_distance = np.inf
                                for index, point in enumerate(cur_point_list):
                                    if all_points.part[point] != part_result[cycle_index][head_index]:
                                        continue

                                    horizontal_distance = abs(all_points.pos[point].x - start_points) + 0 * abs(
                                        all_points.pos[point].y - ref_pos_y)

                                    if horizontal_distance < min_horizontal_distance:
                                        min_horizontal_distance = horizontal_distance
                                        point_index = index
                            else:
                                point_index = -1
                                min_cheby_distance = np.inf

                                for index, point in enumerate(cur_point_list):
                                    if all_points.part[point] != part_result[cycle_index][head_index]:
                                        continue
                                    point_pos = [Point(all_points.pos[point].x - head_index * self.config.head_intv,
                                                       all_points.pos[point].y)]

                                    cheby_distance, euler_distance = 0, 0
                                    for next_head in range(head_num):
                                        if head_point[next_head] == -1:
                                            continue
                                        point_pos.append(Point(assigned_point[next_head].x - next_head * head_num,
                                                               assigned_point[next_head].y))

                                    point_pos = sorted(point_pos, key=lambda p: p.x)
                                    for mount_seq in range(len(point_pos) - 1):
                                        delta_x = axis_moving_time(
                                            point_pos[mount_seq].x - point_pos[mount_seq + 1].x, 0)
                                        delta_y = axis_moving_time(
                                            point_pos[mount_seq].y - point_pos[mount_seq + 1].y, 1)
                                        cheby_distance += max(delta_x, delta_y)
                                        euler_distance += math.sqrt(delta_x ** 2 + delta_y ** 2)

                                    # cheby_distance += 0.01 * euler_distance
                                    if cheby_distance < min_cheby_distance:
                                        min_cheby_distance, min_euler_distance = cheby_distance, euler_distance
                                        point_index = index

                            if point_index == -1:
                                continue

                            head_counter += 1

                            step_index = cur_point_list[point_index]
                            head_point[head_index] = all_points.step[step_index]
                            assigned_point[head_index] = all_points.pos[step_index]

                            cur_point_list.pop(point_index)

                        dist, head_seq = self.dynamic_programming_cycle_path(head_point, slot_result[cycle_index])

                        if dist < min_dist:
                            best_head_point, best_head_seq, best_point_list = head_point, head_seq, cur_point_list
                            min_dist = dist
                if pbar:
                    pbar.update(1)
                point_list = best_point_list
                point_result.append(best_head_point)
                sequence_result.append(best_head_seq)

        return point_result, sequence_result

    def greedy_cluster(self, part_result, cycle_result, slot_result):
        point_result, sequence_result = [], []

        # === assign CT group to feeder slot ===
        component_point_pos = defaultdict(list)

        for idx, data in self.step_data.iterrows():
            component_point_pos[data.part].append([data.x + self.config.stopper_pos.x,
                                                   data.y + self.config.stopper_pos.y, idx])

        for pos_list in component_point_pos.values():
            pos_list.sort(key=lambda x: (x[0], x[1]))

        component_point_index = defaultdict(int)
        for cycle_set in range(len(cycle_result)):
            for cycle in range(cycle_result[cycle_set]):
                point_result.append([-1 for _ in range(self.config.head_num)])
                for head in range(self.config.head_num):
                    part_index = part_result[cycle_set][head]
                    if part_index == -1:
                        continue

                    part = self.part_data.loc[part_index]['part']
                    point_info = component_point_pos[part][component_point_index[part]]

                    point_result[-1][head] = point_info[2]
                    # mount_point[head] = point_info[0:2]

                    component_point_index[part] += 1
                sequence_result.append(
                    self.dynamic_programming_cycle_path(point_result[-1], slot_result[cycle_set])[1])
        return point_result, sequence_result

    def greedy_level_placing(self, part_result, cycle_result, slot_result):
        point_result, sequence_result = [], []
        part_indices = defaultdict(int)
        for part_idx, data in self.part_data.iterrows():
            part_indices[data.part] = part_idx

        mount_point_pos = defaultdict(list)
        for pcb_idx, data in self.step_data.iterrows():
            mount_point_pos[part_indices[data.part]].append([data.x, data.y, pcb_idx])

        for index_ in mount_point_pos.keys():
            mount_point_pos[index_].sort(key=lambda x: (x[1], x[0]))

        for cycle_idx, _ in enumerate(cycle_result):
            for _ in range(cycle_result[cycle_idx]):
                point_result.append([-1 for _ in range(self.config.head_num)])
                for head in range(self.config.head_num):
                    if part_result[cycle_idx][head] == -1:
                        continue
                    index_ = part_result[cycle_idx][head]
                    point_result[-1][head] = mount_point_pos[index_][-1][2]
                    mount_point_pos[index_].pop()
                sequence_result.append(self.dynamic_programming_cycle_path(point_result[-1], slot_result[cycle_idx])[1])
        return point_result, sequence_result

