import pandas as pd

from core.interface import *
from data.type import *
from data.loader import *

from opt.hyper_heuristic import *


class TopBtmOpt(BaseOpt):
    def __init__(self, config, part_data, step_data, feeder_data=None):
        super().__init__(config, part_data, step_data, feeder_data)
        self.feeder_data = pd.DataFrame(columns=['slot', 'part'])
        self.optimizer = FeederPriorityOpt

    def optimize(self, hinter=True):
        rows = pd.Series()
        top_part, btm_part = set(), set()

        part_points = defaultdict(int)
        part_top_points = defaultdict(int)
        for r, data in self.step_data.iterrows():
            if data.layer == 1:
                top_part.add(data.part)
                part_top_points[data.part] += 1
            else:
                btm_part.add(data.part)
            part_points[data.part] += 1
        # part_points = {part: points for part, points in part_points.items() if part in top_part and part in btm_part}
        points_stat = pd.DataFrame(columns=['part', 'top', 'btm', 'total'])
        for part, points in part_points.items():
            points_stat.loc[len(points_stat)] = [part, part_top_points[part], points - part_top_points[part], points]

        for r, data in self.step_data.iterrows():
            rows[r] = True if data.part in top_part and data.part in btm_part and part_points[data.part] > 7 else False

        dblayer_step_data = self.step_data[rows]
        opt = self.optimizer(self.config, self.part_data, dblayer_step_data, self.feeder_data)
        opt.optimize()
        self.feeder_data = opt.feeder_data

        for _, layer_step in self.step_data.groupby('layer'):
            opt = self.optimizer(self.config, self.part_data, layer_step, self.feeder_data)
            opt.optimize(hinter=hinter)

            self.result = self.result + opt.result
            self.feeder_data = opt.feeder_data


@timer_wrapper
def main():
    # 参数解析
    parser = argparse.ArgumentParser(description='assembly line optimizer implementation')
    # parser.add_argument('--filename', default='PCB-1.txt', type=str, help='pcb data')
    parser.add_argument('--filename', default='US2.259.0600-US7.820.51964-A1-B3-64T.txt', type=str, help='pcb data')
    params = parser.parse_args()

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

    part_data, step_data = load_data(params.filename)
    config = MachineConfig()

    # # === 整线 ===
    # opt = HyperHeuristicOpt(2, part_data, step_data)
    # opt.optimize()
    # for machine_index, result in enumerate(opt.result):
    #     print('-------- Machine: ' + str(machine_index) + '--------')
    #     evaluation(config, part_data, step_data, result, hinter=True)     # 估算贴装用时

    # === 单机 ===
    opt = TopBtmOpt(config, part_data, step_data)
    opt.optimize()
    evaluation(config, part_data, step_data, opt.result, hinter=True)  # 估算贴装用时
    print('------------------------------ ')


if __name__ == '__main__':
    main()