from core.common import *
from data.type import *


head_rotary_velocity = 8e-5  # 贴装头R轴旋转时间
x_max_velocity, y_max_velocity = 1.4, 1.2
x_max_acceleration, y_max_acceleration = x_max_velocity / 0.079, y_max_velocity / 0.079


def axis_moving_time(distance, axis=0):
    assert 0 <= axis <= 2
    distance = abs(distance) * 1e-3
    Lamax = x_max_velocity ** 2 / x_max_acceleration if axis == 0 else y_max_velocity ** 2 / y_max_acceleration
    Tmax = x_max_velocity / x_max_acceleration if axis == 0 else y_max_velocity / y_max_acceleration
    if axis == 0:
        return 2 * math.sqrt(distance / x_max_acceleration) if distance < Lamax else 2 * Tmax + (
                    distance - Lamax) / x_max_velocity
    elif axis == 1:
        return 2 * math.sqrt(distance / y_max_acceleration) if distance < Lamax else 2 * Tmax + (
                    distance - Lamax) / y_max_velocity
    elif axis == 2:
        return abs((distance + 180) % 360 - 180) * head_rotary_velocity
    return .0


def optimizer_result_hinter(config, part_data, opt_res, nozzle_hinter=False, part_hinter=False,
                            slot_hinter=False, place_hinter=False):
    if nozzle_hinter:
        columns = ['H{}'.format(i + 1) for i in range(config.head_num)] + ['cycle']

        nozzle_assign = pd.DataFrame(columns=columns)
        for cycle, components in enumerate(opt_res.part):
            nozzle_assign_row = len(nozzle_assign)
            nozzle_assign.loc[nozzle_assign_row, 'cycle'] = opt_res.cycle[cycle]

            for head in range(config.head_num):
                index = opt_res.part[cycle][head]
                if index == -1:
                    nozzle_assign.loc[nozzle_assign_row, 'H{}'.format(head + 1)] = ''
                else:
                    nozzle = part_data.loc[index]['nz']
                    nozzle_assign.loc[nozzle_assign_row, 'H{}'.format(head + 1)] = nozzle

            for head in range(config.head_num):
                if nozzle_assign_row == 0 or nozzle_assign.loc[nozzle_assign_row - 1, 'H{}'.format(head + 1)] != \
                        nozzle_assign.loc[nozzle_assign_row, 'H{}'.format(head + 1)]:
                    break
            else:
                nozzle_assign.loc[nozzle_assign_row - 1, 'cycle'] += nozzle_assign.loc[nozzle_assign_row, 'cycle']
                nozzle_assign.drop([len(nozzle_assign) - 1], inplace=True)

        print(nozzle_assign)
        print('')

    if part_hinter:
        columns = ['H{}'.format(i + 1) for i in range(config.head_num)] + ['cycle']

        part_assign = pd.DataFrame(columns=columns)
        for cycle, components in enumerate(opt_res.part):
            part_assign.loc[cycle, 'cycle'] = opt_res.cycle[cycle]
            for head in range(config.head_num):
                index = opt_res.part[cycle][head]
                if index == -1:
                    part_assign.loc[cycle, 'H{}'.format(head + 1)] = ''
                else:
                    # component_assign.loc[cycle, 'H{}'.format(head + 1)] = component_data.loc[index]['part']
                    part_assign.loc[cycle, 'H{}'.format(head + 1)] = 'C' + str(index)

        print(part_assign)
        print('')

    if slot_hinter:
        columns = ['H{}'.format(i + 1) for i in range(config.head_num)] + ['cycle']

        slot_assign = pd.DataFrame(columns=columns)
        for cycle, components in enumerate(opt_res.slot):
            slot_assign.loc[cycle, 'cycle'] = opt_res.cycle[cycle]
            for head in range(config.head_num):
                slot = opt_res.slot[cycle][head]
                if slot == -1:
                    slot_assign.loc[cycle, 'H{}'.format(head + 1)] = 'A'
                else:
                    try:
                        slot_assign.loc[cycle, 'H{}'.format(head + 1)] = 'F{}'.format(
                            slot) if slot <= config.slot_num // 2 else 'R{}'.format(slot - config.head_num)
                    except:
                        print('')

        print(slot_assign)
        print('')

    if place_hinter:
        columns = ['H{}'.format(i + 1) for i in range(config.head_num)] + ['cycle']

        place_assign = pd.DataFrame(columns=columns)
        for cycle, _ in enumerate(opt_res.point):
            place_assign.loc[cycle, 'cycle'] = 1
            for head in range(config.head_num):
                point = opt_res.point[cycle][head]
                if point != -1:
                    place_assign.loc[cycle, 'H{}'.format(head + 1)] = 'P{}'.format(point)
                else:
                    place_assign.loc[cycle, 'H{}'.format(head + 1)] = ''

        headseq_assign = pd.DataFrame(columns=columns)
        for cycle, headseq in enumerate(opt_res.sequence):
            headseq_assign.loc[cycle, 'cycle'] = 1
            for head in range(len(headseq)):
                headseq_assign.loc[cycle, 'H{}'.format(head + 1)] = 'H{}'.format(headseq[head])

        print(place_assign)
        print(headseq_assign)
        print('')


def evaluation(config: MachineConfig, part_data, step_data, opt_res: OptResult, hinter=False):
    # === 优化结果参数 ===
    info = OptInfo()
    # === 校验 ===
    info.total_points = 0
    for cycle, part in enumerate(opt_res.part):
        for head, component in enumerate(part):
            if component == -1:
                continue
            info.total_points += opt_res.cycle[cycle]

    if info.total_points != len(step_data):
        warning_info = 'the number of ' + str(info.total_points) + ' placement point(s) is not match with the PCB data. '
        warnings.warn(warning_info, UserWarning)

    interval_ratio = config.head_intv / config.slot_intv
    if opt_res.point:
        total_points = info.total_points
        for placements in opt_res.point:
            for placement in placements:
                if placement == -1:
                    continue
                total_points -= 1

        if total_points != 0:
            warnings.warn(
                'the optimization result of component assignment result and placement result are not consistent. ',
                UserWarning)
            return OptInfo()

    feeder_arrangement = defaultdict(set)
    for cycle, feeder_slots in enumerate(opt_res.slot):
        for head, slot in enumerate(feeder_slots):
            if slot == -1:
                continue
            feeder_arrangement[opt_res.part[cycle][head]].add(slot)

    info.total_components = len(feeder_arrangement.keys())
    for part, data in part_data.iterrows():
        if part in feeder_arrangement.keys() and data.fdn < len(feeder_arrangement[part]):
            info = 'the number of arranged feeder of [' + data['part'] + '] exceeds the quantity limit'
            warnings.warn(info, UserWarning)

    cur_pos, next_pos = config.anc_pos, Point(0, 0)                         # 贴装头当前位置

    # 初始化首个周期的吸嘴装配信息
    nozzle_assigned = ['Empty' for _ in range(config.head_num)]
    for head in range(config.head_num):
        for cycle in range(len(opt_res.part)):
            idx = opt_res.part[cycle][head]
            if idx == -1:
                continue
            else:
                nozzle_assigned[head] = part_data.loc[idx]['nz']

    for cycle_set, _ in enumerate(opt_res.part):
        floor_cycle, ceil_cycle = sum(opt_res.cycle[:cycle_set]), sum(opt_res.cycle[:(cycle_set + 1)])
        for cycle in range(floor_cycle, ceil_cycle):
            if sum(opt_res.part[cycle_set]) == -config.head_num:
                continue
            pick_slot, mount_pos, mount_angle = [], [], []
            nozzle_pick_counter, nozzle_put_counter = 0, 0  # 吸嘴更换次数统计（拾取/放置分别算一次）
            for head in range(config.head_num):
                if opt_res.slot[cycle_set][head] != -1:
                    pick_slot.append(opt_res.slot[cycle_set][head] - interval_ratio * head)
                if opt_res.part[cycle_set][head] == -1:
                    continue
                nozzle = part_data.loc[opt_res.part[cycle_set][head]]['nz']
                if nozzle != nozzle_assigned[head]:
                    if nozzle_assigned[head] != 'Empty':
                        nozzle_put_counter += 1
                    nozzle_pick_counter += 1
                    nozzle_assigned[head] = nozzle

            # ANC处进行吸嘴更换
            if nozzle_pick_counter + nozzle_put_counter > 0:
                next_pos = config.anc_pos
                move_time = max(axis_moving_time(cur_pos.x - next_pos.x, 0),
                                axis_moving_time(cur_pos.y - next_pos.y, 1))
                info.round_time += move_time
                info.anc_round_counter += 1
                info.total_distance += max(abs(cur_pos.x - next_pos.x), abs(cur_pos.y - next_pos.y))
                cur_pos = next_pos

            pick_slot = list(set(pick_slot))
            pick_slot = sorted(pick_slot, reverse=True)

            # 拾取路径(自右向左)
            for idx, slot in enumerate(pick_slot):
                if slot < config.slot_num // 2:
                    next_pos = Point(config.slotf1_pos.x + config.slot_intv * (slot - 1), config.slotf1_pos.y)
                else:
                    next_pos = Point(config.slotr1_pos.x - config.slot_intv * (config.slot_num - slot - 1), config.slotr1_pos.y)
                info.operation_time += config.pick_time
                info.pickup_counter += 1

                move_time = max(axis_moving_time(cur_pos.x - next_pos.x, 0),
                                axis_moving_time(cur_pos.y - next_pos.y, 1))
                if idx == 0:
                    info.round_time += move_time
                else:
                    info.pickup_time += move_time

                info.total_distance += max(abs(cur_pos.x - next_pos.x), abs(cur_pos.y - next_pos.y))
                if slot != pick_slot[0]:
                    info.pickup_distance += max(abs(cur_pos.x - next_pos.x), abs(cur_pos.y - next_pos.y))
                cur_pos = next_pos

            # 贴装路径
            if opt_res.point and opt_res.sequence:
                head_angle = [0 for _ in range(config.head_num)]
                for head in opt_res.sequence[cycle]:
                    index = opt_res.point[cycle][head]
                    if index == -1:
                        continue
                    mount_pos.append(Point(step_data.loc[index].x - head * config.head_intv + config.stopper_pos.x,
                                           step_data.loc[index].y + config.stopper_pos.y))
                    head_angle[head] = step_data.loc[index]['r']

                # 单独计算贴装路径
                for cntPoints in range(len(mount_pos) - 1):
                    info.place_distance += max(abs(mount_pos[cntPoints].x - mount_pos[cntPoints + 1].x),
                                               abs(mount_pos[cntPoints].y - mount_pos[cntPoints + 1].y))

                if mount_pos[0].x < mount_pos[-1].x:
                    mount_pos = reversed(mount_pos)

                # 考虑R轴预旋转，补偿同轴角度转动带来的额外贴装用时
                info.operation_time += config.nozzle_install_time * nozzle_put_counter + \
                                       config.nozzle_uninstall_time * nozzle_pick_counter
                for idx, pos in enumerate(mount_pos):
                    info.operation_time += config.place_time

                    if idx == 0:
                        move_time = max(axis_moving_time(cur_pos.x - pos.x, 0),
                                        axis_moving_time(cur_pos.y - pos.y, 1))
                        info.round_time += move_time
                    else:

                        cur_head = opt_res.sequence[cycle][idx]
                        side_head = cur_head - 1 if cur_head % 2 else cur_head + 1
                        if opt_res.sequence[cycle][idx - 1] != side_head:
                            move_time = max(axis_moving_time(cur_pos.x - pos.x, 0),
                                            axis_moving_time(cur_pos.y - pos.y, 1))
                        else:
                            move_time = max(axis_moving_time(cur_pos.x - pos.x, 0),
                                            axis_moving_time(cur_pos.y - pos.y, 1),
                                            axis_moving_time(head_angle[cur_head] - head_angle[side_head], 2))
                        info.place_time += move_time

                    info.total_distance += max(abs(cur_pos.x - pos.x), abs(cur_pos.y - pos.y))
                    cur_pos = pos

            info.nozzle_change_counter += nozzle_put_counter + nozzle_pick_counter

    info.total_time = info.pickup_time + info.round_time + info.place_time + info.operation_time
    info.cycle_counter = sum(opt_res.cycle)
    if hinter:
        optimizer_result_hinter(config, part_data, opt_res, part_hinter=True, nozzle_hinter=True, slot_hinter=True)
        info.print()

    return info