from core.common import *
from opt.smm.basis import BaseOpt

import torch


class Predictor:
    def __init__(self):
        pass

    @staticmethod
    def training(self, params):
        pass

    @staticmethod
    def testing(self, params):
        pass

    @staticmethod
    def predict(self, cp_points, cp_nozzle, board_width=None, board_height=None):
        pass


class Net(torch.nn.Module):
    def __init__(self, input_size, hidden_size=1000, output_size=1):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(input_size, hidden_size)
        self.relu = torch.nn.ReLU()  # 激活函数
        self.fc2 = torch.nn.Linear(hidden_size, hidden_size)
        # self.relu1 = torch.nn.ReLU()  # 激活函数
        self.fc3 = torch.nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        # x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x


class NeuralPredictor(Predictor, BaseOpt):
    def __init__(self):
        super().__init__()
        self.min_placement_points = 10
        self.max_placement_points = 1000

        self.max_component_types = 30
        self.default_feeder_limit = 1
        self.max_nozzle_types = 4

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.net = Net(input_size=self.get_feature(), output_size=1).to(self.device)
        self.net_file = 'opt/param.pth'
        try:
            self.net.load_state_dict(torch.load(self.net_file, map_location=self.device))
        except:
            warnings.warn('the parameters of neural net model load failed', UserWarning)

    def init_weights(self):
        for m in self.net.modules():
            if isinstance(m, torch.nn.Linear):
                torch.nn.init.xavier_uniform_(m.weight)
                torch.nn.init.zeros_(m.bias)

    def subobjective(self, cp_points, cp_nozzle, config):
        if len(cp_points.keys()) or sum(cp_points.values()) == 0:
            return 0, 0, 0, 0
        nozzle_heads, nozzle_points = defaultdict(int), defaultdict(int)
        for idx, points in cp_points.items():
            if points == 0:
                continue
            nozzle = cp_nozzle[idx]
            nozzle_points[nozzle] += points
            nozzle_heads[nozzle] = 1

        anc_round_counter = 0
        while sum(nozzle_heads.values()) != config.head_num:
            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(config.head_num)]
            for idx, nozzle_item in enumerate(nozzle_group):
                if idx < config.head_num:
                    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 = config.cycle_time * max(head_total_points)
            for head in range(config.head_num):
                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 += self.nozzle_change_weight
                    else:
                        if tmp_head_nozzle_assignment[head][cycle][0] != tmp_head_nozzle_assignment[head][cycle + 1][0]:
                            cost += self.nozzle_change_weight

            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(config.head_num):
                    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(config.head_num):
                    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 self.cycle_weight * (max(head_total_points) - max_cycle) < self.nozzle_change_weight * nozzle_change_counter:
                    break

                cost -= self.cycle_weight * (max(head_total_points) - max_cycle) - self.nozzle_change_weight * 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(config.head_num):
            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 cp_points.items():
            if points == 0:
                continue
            feeder_limit = 1  # todo: 暂时仅考虑一种吸嘴的情形
            reminder_points = points % feeder_limit
            for _ in range(feeder_limit):
                cp_info.append([idx, points // feeder_limit + (1 if reminder_points > 0 else 0), cp_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())

        return cycle_counter, nozzle_change_counter, anc_round_counter, pickup_counter

    def encode(self, cp_points: defaultdict[str], cp_nozzle: defaultdict[str], board_width, board_height, config):
        assert len(cp_points.keys()) == len(cp_nozzle.keys())

        # === general info ===
        total_points = sum(points for points in cp_points.values())
        total_component_types, total_nozzle_types = len(cp_points.keys()), len(set(cp_nozzle.values()))

        data = [total_points, total_component_types, total_nozzle_types]
        data.extend([board_width, board_height])

        # === heuristic info ===
        cycle, nozzle_change, anc_move, pickup = self.subobjective(cp_points, cp_nozzle, config)
        data.extend([cycle, nozzle_change, anc_move, pickup])

        # === nozzle info ===
        nozzle_points = defaultdict(int)
        for cp_idx, nozzle in cp_nozzle.items():
            nozzle_points[cp_nozzle[cp_idx]] += cp_points[cp_idx]       # points for different nozzle type
        nozzle_items = [[nozzle, points] for nozzle, points in nozzle_points.items()]
        nozzle_items = sorted(nozzle_items, key=lambda x: x[1], reverse=True)

        nz2idx = defaultdict(int)
        nozzle_slice = [0 for _ in range(self.max_nozzle_types)]
        for idx, [nozzle, points] in enumerate(nozzle_items):
            nz2idx[nozzle] = idx if idx < self.max_nozzle_types else self.max_nozzle_types - 1
            nozzle_slice[idx if idx < self.max_nozzle_types else -1] += points
        data.extend(nozzle_slice)

        # === part info ===
        part_data_slice = defaultdict(list)
        for idx in range(self.max_nozzle_types):
            part_data_slice[idx] = []

        cp_items = [[component, points] for component, points in cp_points.items()]
        cp_items = sorted(cp_items, key=lambda x: (x[1], nz2idx[cp_nozzle[x[0]]] * 0.1 + x[1]), reverse=True)
        for component, points in cp_items:
            nozzle = cp_nozzle[component]
            part_data_slice[nz2idx[nozzle]].append(points)

        data_slice = [0 for _ in range(self.max_nozzle_types)]
        for idx, part_list in part_data_slice.items():
            data_slice[idx] = len(part_list)
        data.extend(data_slice)

        for idx in range(self.max_nozzle_types):
            if len(part_data_slice[idx]) <= self.max_component_types:
                part_data_slice[idx].extend([0 for _ in range(self.max_component_types - len(part_data_slice[idx]))])
            else:
                part_data_slice[idx] = part_data_slice[idx][:self.max_component_types]
            data.extend(part_data_slice[idx])

        return data

    def get_feature(self):
        return (self.max_component_types + 2) * self.max_nozzle_types + 5 + 4

    def eval(self, cp_points, cp_nozzle, board_width, board_height, config):
        encoding = np.array(self.encode(cp_points, cp_nozzle, board_width, board_height, config))
        encoding = torch.from_numpy(encoding.reshape((-1, np.shape(encoding)[0]))).float().to(self.device)
        return self.net(encoding)[0, 0].item()
