from dataloader import *
from lineopt_genetic import line_optimizer_genetic
from lineopt_heuristic import line_optimizer_heuristic
from lineopt_reconfiguration import line_optimizer_reconfiguration
from lineopt_hyperheuristic import line_optimizer_hyperheuristic
from lineopt_model import line_optimizer_model

from base_optimizer.optimizer_interface import *


def optimizer(pcb_data, component_data, feeder_data, params, hinter=True):
    if params.machine_number == 1:
        assembly_info = [base_optimizer(1, pcb_data, component_data, feeder_data, params, hinter=hinter)]
        return assembly_info

    if params.line_optimizer == 'hyper-heuristic' or params.line_optimizer == 'heuristic' or params.line_optimizer \
            == 'genetic' or params.line_optimizer == 'reconfiguration':
        if params.line_optimizer == 'hyper-heuristic':
            assignment_result = line_optimizer_hyperheuristic(component_data, pcb_data, params.machine_number)
        elif params.line_optimizer == "heuristic":
            assignment_result = line_optimizer_heuristic(component_data, params.machine_number)
        elif params.line_optimizer == "genetic":
            assignment_result = line_optimizer_genetic(component_data, params.machine_number)
        else:
            assignment_result = line_optimizer_reconfiguration(component_data, pcb_data, params.machine_number)

        partial_pcb_data, partial_component_data = convert_line_assigment(pcb_data, component_data, assignment_result)
        assembly_info = []
        for machine_index in range(params.machine_number):
            assembly_info.append(base_optimizer(machine_index + 1, partial_pcb_data[machine_index],
                                                partial_component_data[machine_index],
                                                pd.DataFrame(columns=['slot', 'part']), params, hinter=hinter))
    elif params.line_optimizer == 'mip-model':
        assembly_info = line_optimizer_model(component_data, pcb_data, params.machine_number)
    else:
        raise 'line optimizer method is not existed'

    return assembly_info


@timer_wrapper
def main():
    # 参数解析
    parser = argparse.ArgumentParser(description='assembly line optimizer implementation')
    parser.add_argument('--mode', default=1, type=int, help='mode: 0 -directly load pcb data without optimization '
                                                            'for data analysis, 1 -optimize pcb data, 2 -batch test')
    parser.add_argument('--filename', default='data-1.txt', type=str, help='load pcb data')

    parser.add_argument('--comp_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-priority', 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('--feeder_limit', default=1, type=int, help='the upper feeder limit for each type of component')
    parser.add_argument('--save', default=0, type=int, help='save the optimization result')
    parser.add_argument('--save_suffix', default='(1)', type=str, help='load pcb data')
    params = parser.parse_args()

    # 结果输出显示所有行和列
    pd.set_option('display.max_columns', None)
    pd.set_option('display.max_rows', None)
    if params.mode == 0:
        partial_pcb_data, partial_component_data, _ = load_data(params.filename, load_feeder=False)
        assembly_info = []
        for machine_index in range(len(partial_pcb_data)):
            opt_res = convert_pcbdata_to_result(partial_pcb_data[machine_index], partial_component_data[machine_index])

            info = placement_info_evaluation(partial_component_data[machine_index], partial_pcb_data[machine_index],
                                             opt_res)
            assembly_info.append(info)
            optimization_assign_result(partial_component_data[machine_index], partial_pcb_data[machine_index], opt_res,
                                       nozzle_hinter=True, component_hinter=True, feeder_hinter=True)

            info.print()

            if params.save:
                output_optimize_result(f'result/{params.filename[:-4]}-T-Solution-M0{machine_index + 1}',
                                       partial_component_data[machine_index], partial_pcb_data[machine_index], opt_res)
            print('------------------------------ ')

        for machine_idx, info in enumerate(assembly_info):
            print(f'assembly time for machine {machine_idx + 1: d}: {info.total_time: .3f} s, total placement: '
                  f'{info.total_points}, total component types {info.total_components: d}')

        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}')

    elif params.mode == 1:

        # 加载PCB数据
        partial_pcb_data, partial_component_data, feeder_data = load_data(params.filename, load_feeder=True)
        pcb_data, component_data = merge_data(partial_pcb_data, partial_component_data)
        start_time = time.time()

        assembly_info = optimizer(pcb_data, component_data, feeder_data, params)

        # sys.stdout = sys.__stdout__
        print(f'optimizer running time: {time.time() - start_time: .3f}')
        for machine_idx, info in enumerate(assembly_info):
            print(f'assembly time for machine {machine_idx + 1: d}: {info.total_time: .3f} s, total placement: '
                  f'{info.total_points}, total component types {info.total_components: d}')

        print(f'assembly metric evaluation： {max(info.metric() for info in assembly_info)}')

        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}')

    elif params.mode == 2:

        # sys.stdout = open(f'record/dissertation-experiment.txt', 'w')
        machine_optimizer = ['two-phase', 'hybrid-genetic', 'cell-division', 'feeder-priority', 'aggregation']
        running_round = 10
        opt_columns = ['Cycle', 'Pick', 'Nozzle-Change', 'Running-Time']
        opt_result, opt_runtime = defaultdict(pd.DataFrame), defaultdict(pd.DataFrame)
        for opt in machine_optimizer:
            opt_result[opt] = pd.DataFrame(columns=opt_columns)
            opt_result[opt].index.name = 'file'

        for _, file in enumerate(os.listdir('data/')):
            if file[-3:] != 'txt':
                continue

            partial_pcb_data, partial_component_data, feeder_data = load_data(file)
            pcb_data, component_data = merge_data(partial_pcb_data, partial_component_data)

            for opt in machine_optimizer:
                for round_idx in range(running_round):
                    print(f'--- file ：  {file},  round  ：  {round_idx + 1},  optimizer ：  {opt} --- ')

                    params = parser.parse_args(['--machine_optimizer', opt, '--machine_number', str(1), '--filename', file])

                    start_time = time.time()
                    assembly_info = optimizer(pcb_data, component_data, feeder_data, params, hinter=False)

                    opt_result[opt].loc[file + str(round_idx + 1), 'Cycle'] = assembly_info[0].cycle_counter
                    opt_result[opt].loc[file + str(round_idx + 1), 'Pick'] = assembly_info[0].pickup_counter
                    opt_result[opt].loc[file + str(round_idx + 1), 'Nozzle-Change'] = assembly_info[0].nozzle_change_counter
                    opt_result[opt].loc[file + str(round_idx + 1), 'Running-Time'] = time.time() - start_time

        with pd.ExcelWriter('result/machine_optimizer.xlsx', engine='openpyxl') as writer:
            for opt, result in opt_result.items():
                result.to_excel(writer, sheet_name=opt, float_format='%.3f', na_rep='')
    else:
        line_optimizer = ['T-Solution', 'hyper-heuristic', 'genetic', 'reconfiguration']
        file_dirs = ['L01', 'L02', 'L03']

        running_round = 10
        line_opt_result, line_opt_runtime = defaultdict(pd.DataFrame), defaultdict(pd.DataFrame)

        opt_columns = []
        for line_opt in line_optimizer:
            if line_opt == 'T-Solution':
                opt_columns.append(line_opt)
            else:
                opt_columns.extend([line_opt + str(idx + 1) for idx in range(running_round)])

        for file_dir in file_dirs:
            line_opt_result[file_dir] = pd.DataFrame(columns=opt_columns)
            line_opt_runtime[file_dir] = pd.DataFrame(columns=opt_columns)
            line_opt_result[file_dir].index.name, line_opt_runtime[file_dir].index.name = 'file', 'file'

            for file_index, file in enumerate(os.listdir('data/' + file_dir)):
                sys.stdout = sys.__stdout__
                print(f'--- {file_dir} : ({file_index + 1}) file ：  {file} --- ')
                try:
                    partial_pcb_data, partial_component_data, _ = load_data(file_dir + '/' + file, load_feeder=False)
                except:
                    traceback.print_exc()
                    warning_info = f'file: {file_dir}/{file}: an unexpected error occurs for data loader'
                    warnings.warn(warning_info, SyntaxWarning)
                    continue

                machine_number = len(partial_pcb_data)
                if not os.path.exists(f'record/{file_dir}'):
                    os.makedirs(f'record/{file_dir}')
                merge_pcb_data, merge_component_data = merge_data(partial_pcb_data, partial_component_data)
                for line_opt in line_optimizer:
                    assembly_info = []
                    if line_opt == 'T-Solution':
                        sys.stdout = open(f'record/{file_dir}/{file[:-4]}-{line_opt}.txt', 'w')
                        for machine_index in range(machine_number):
                            opt_res = convert_pcbdata_to_result(partial_pcb_data[machine_index],
                                                                partial_component_data[machine_index])

                            print('----- Placement machine ' + str(machine_index + 1) + ' ----- ')
                            info = placement_info_evaluation(partial_component_data[machine_index],
                                                             partial_pcb_data[machine_index], opt_res, hinter=True)
                            print('------------------------------ ')
                            assembly_info.append(info)
                            if params.save:
                                output_optimize_result(f'result/{file_dir}/{file[:-4]}-T-Solution-M0{machine_index + 1}',
                                                       partial_component_data[machine_index],
                                                       partial_pcb_data[machine_index], opt_res)

                        line_opt_result[file_dir].loc[file, line_opt] = max(info.total_time for info in assembly_info)
                        for machine_idx, info in enumerate(assembly_info):
                            print(
                                f'assembly time for machine {machine_idx + 1: d}: {info.total_time: .3f} s, total placement: '
                                f'{info.total_points}, total component types {info.total_components: d}')

                        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}')
                    else:
                        for round_idx in range(running_round):
                            sys.stdout = open(f'record/{file_dir}/{file[:-4]}-{line_opt} ({round_idx + 1}).txt', 'w')

                            params = parser.parse_args(
                                ['--filename', file_dir + '/' + file, '--machine_number', str(machine_number),
                                 '--line_optimizer', line_opt, '--save_suffix', f'({round_idx + 1})'])
                            start_time = time.time()
                            assembly_info = optimizer(merge_pcb_data, merge_component_data, None, params)

                            line_opt_result[file_dir].loc[file, line_opt + str(round_idx + 1)] = max(
                                info.total_time for info in assembly_info)
                            line_opt_runtime[file_dir].loc[file, line_opt + str(round_idx + 1)] = time.time() - start_time

                            for machine_idx, info in enumerate(assembly_info):
                                print(
                                    f'assembly time for machine {machine_idx + 1: d}: {info.total_time: .3f} s, '
                                    f'total placement:  {info.total_points}, '
                                    f'total component types {info.total_components: d}')

                            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}')

        with pd.ExcelWriter('result/line_optimizer.xlsx', engine='openpyxl') as writer:
            for file_dir, result in line_opt_result.items():
                result.to_excel(writer, sheet_name='result-' + file_dir, float_format='%.3f', na_rep='')

            for file_dir, running_time in line_opt_runtime.items():
                running_time.to_excel(writer, sheet_name='running_time-' + file_dir, float_format='%.3f', na_rep='')


if __name__ == '__main__':
    main()


