import random

import numpy as np

from dataloader import *
from optimizer_genetic import line_optimizer_genetic
from optimizer_heuristic import line_optimizer_heuristic
from optimizer_reconfiguration import line_optimizer_reconfiguration
from optimizer_hyperheuristic import line_optimizer_hyperheuristic

from base_optimizer.optimizer_interface import *


def optimizer(pcb_data, component_data, line_optimizer, machine_optimizer, machine_number):
    if machine_number > 1:
        if line_optimizer == 'hyper-heuristic':
            assignment_result = line_optimizer_hyperheuristic(component_data, pcb_data, machine_number)
        elif line_optimizer == "heuristic":
            assignment_result = line_optimizer_heuristic(component_data, machine_number)
        elif line_optimizer == "genetic":
            assignment_result = line_optimizer_genetic(component_data, machine_number)
        elif line_optimizer == "reconfiguration":
            assignment_result = line_optimizer_reconfiguration(component_data, pcb_data, machine_number)
        else:
            raise 'line optimizer method is not existed'
    else:
        assignment_result = [[]]
        for _, data in component_data.iterrows():
            assignment_result[-1].append(data.points)

    partial_pcb_data, partial_component_data = convert_line_assigment(pcb_data, component_data, assignment_result)
    assembly_info = []
    for machine_index in range(machine_number):
        assembly_info.append(
            base_optimizer(machine_index + 1, partial_pcb_data[machine_index], partial_component_data[machine_index],
                           feeder_data=pd.DataFrame(columns=['slot', 'part', 'arg']), method=machine_optimizer,
                           hinter=True))

    for machine_index in range(machine_number):
        total_component_types = sum(1 if pt else 0 for pt in assignment_result[machine_index])
        total_placement_points = sum(assignment_result[machine_index])
        total_time = assembly_info[machine_index].total_time
        print(f'assembly time for machine {machine_index + 1: d}: {total_time: .3f} s, total placement: '
              f'{total_placement_points}, total component types {total_component_types: d}', end='')
        for part_index in range(len(assignment_result[machine_index])):
            if assignment_result[machine_index][part_index]:
                print(', ', part_index, end='')
        print('')

    print(f'finial assembly time: {max(info.total_time for info in assembly_info): .3f} s, '
          f'standard deviation: {np.std([info.total_time for info in assembly_info]): .3f}')


@timer_wrapper
def main():
    warnings.simplefilter(action='ignore', category=FutureWarning)
    # 参数解析
    parser = argparse.ArgumentParser(description='assembly line optimizer implementation')
    parser.add_argument('--filename', default='PCB.txt', type=str, help='load pcb data')
    parser.add_argument('--auto_register', default=1, type=int, help='register the component according the pcb data')
    parser.add_argument('--machine_number', default=3, type=int, help='the number of machine in the assembly line')
    parser.add_argument('--machine_optimizer', default='feeder-scan', type=str, help='optimizer for single machine')
    parser.add_argument('--line_optimizer', default='hyper-heuristic', type=str, help='optimizer for PCB assembly line')
    # parser.add_argument('--line_optimizer', default='genetic', type=str, help='optimizer for PCB assembly line')
    parser.add_argument('--feeder_limit', default=1, type=int, help='the upper feeder limit for each type of component')
    params = parser.parse_args()

    # 结果输出显示所有行和列
    pd.set_option('display.max_columns', None)
    pd.set_option('display.max_rows', None)

    # 加载PCB数据
    pcb_data, component_data, _ = load_data(params.filename, default_feeder_limit=params.feeder_limit,
                                            cp_auto_register=params.auto_register, load_feeder_data=False)  # 加载PCB数据

    optimizer(pcb_data, component_data, params.line_optimizer, params.machine_optimizer, params.machine_number)

    # index_list, part_list = [1,  4,  8,  9,  12,  13,  14,  18,  20,  22,  23,  25,  33,  35,  38,  39,  40], []
    # for idx in index_list:
    #     part_list.append(component_data.iloc[idx].part)
    # pcb_data = pcb_data[pcb_data['part'].isin(part_list)].reset_index(drop=True)
    # component_data = component_data.iloc[index_list].reset_index(drop=True)
    # optimizer(pcb_data, component_data, params.line_optimizer, params.machine_optimizer, 1)
    #
    # from optimizer_hyperheuristic import DataMgr, Net
    # data_mgr = DataMgr()

    # cp_points, cp_nozzle = defaultdict(int), defaultdict(str)
    # for _, data in component_data.iterrows():
    #     cp_points[data.part], cp_nozzle[data.part] = data.points, data.nz

    # idx = 1832
    # data = data_mgr.loader(file_name)
    # encoding = np.array(data[0][idx])

    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # net = Net(input_size=data_mgr.get_feature(), output_size=1).to(device)
    #
    # net.load_state_dict(torch.load('model/net_model.pth'))
    # board_width, board_height = pcb_data['x'].max() - pcb_data['x'].min(), pcb_data['y'].max() - pcb_data['y'].min()
    # encoding = np.array(data_mgr.encode(cp_points, cp_nozzle, board_width, board_height))
    # encoding = torch.from_numpy(encoding.reshape((-1, np.shape(encoding)[0]))).float().to("cuda")
    # print(f'net pred time: {net(encoding)[0, 0].item():.3f}')

    # with open('model/lr_model.pkl', 'rb') as f:
    #     lr = pickle.load(f)
    #
    # print('lr model train data: ', np.array(data[2:]).T[idx].reshape(1, -1))
    # print('lr model pred time: ', lr.predict(np.array(data[2:]).T[idx].reshape(1, -1)))
    # print('real time: ', data[-1][idx] * 3600 / data[1][idx])


if __name__ == '__main__':
    main()