from base_optimizer.optimizer_common import *


def load_data(filename: str, default_feeder_limit=1, load_cp_data=True, load_feeder_data=True, cp_auto_register=False):
    # 读取PCB数据
    filename = 'data/' + filename
    part_content, step_content = False, False
    part_start_line, step_start_line, part_end_line, step_end_line = -1, -1, -1, -1
    line_counter = 0
    with open(filename, 'r') as file:
        line = file.readline()
        while line:
            if line == '[Part]\n':
                part_content = True
                part_start_line = line_counter
            elif line == '[Step]\n':
                step_content = True
                step_start_line = line_counter
            elif line == '\n':
                if part_content:
                    part_content = False
                    part_end_line = line_counter
                elif step_content:
                    step_content = False
                    step_end_line = line_counter
            line_counter += 1
            line = file.readline()

    if part_content:
        part_end_line = line_counter
    elif step_content:
        step_end_line = line_counter

    pcb_data = pd.DataFrame(
        pd.read_csv(filepath_or_buffer=filename, skiprows=step_start_line + 1, nrows=step_end_line - step_start_line + 1,
                    sep='\t', header=None))

    if len(pcb_data.columns) <= 17:
        step_col = ["ref", "x", "y", "z", "r", "part", "desc", "fdr", "nz", "hd", "cs", "cy", "sk", "bl", "ar",
                    "pl", "lv"]
    elif len(pcb_data.columns) <= 18:
        step_col = ["ref", "x", "y", "z", "r", "part", "desc", "fdr", "nz", "hd", "cs", "cy", "sk", "bl", "ar", "fid",
                    "pl", "lv"]
    else:
        step_col = ["ref", "x", "y", "z", "r", "part", "desc", "fdr", "nz", "hd", "cs", "cy", "sk", "bl", "ar", "fid",
                    "", "pl", "lv"]

    pcb_data.columns = step_col
    pcb_data = pcb_data.dropna(axis=1)

    # 坐标系处理
    # pcb_data = pcb_data.sort_values(by = ['x', 'y'], ascending = True)
    # pcb_data["x"] = pcb_data["x"].apply(lambda x: -100+x)

    # 注册元件检查
    part_feeder_assign = defaultdict(set)
    part_col = ["part", "fdr", "nz", 'feeder-limit']

    try:
        if part_start_line != -1:
            component_data = pd.DataFrame(
                pd.read_csv(filepath_or_buffer=filename, sep='\t', header=None, skiprows=part_start_line + 1,
                            nrows=part_end_line - part_start_line - 1))
            component_data.columns = part_col
        else:
            component_data = pd.DataFrame(columns=part_col)
    except:
        component_data = pd.DataFrame(columns=part_col)

    component_data['points'] = 0
    part_col.append('points')
    for _, data in pcb_data.iterrows():
        part, nozzle = data.part, data.nz.split(' ')[1]
        slot = data['fdr'].split(' ')[0]
        if part not in component_data['part'].values:
            if not cp_auto_register:
                raise Exception("unregistered component:  " + component_data['part'].values)
            else:
                component_data = pd.concat([component_data, pd.DataFrame(
                    [part, '', 'SM8', nozzle, default_feeder_limit, 0], index=part_col).T],
                                           ignore_index=True)
                warning_info = 'register component ' + part + ' with default feeder type'
                warnings.warn(warning_info, UserWarning)
        part_index = component_data[component_data['part'] == part].index.tolist()[0]
        part_feeder_assign[part].add(slot)
        component_data.loc[part_index, 'points'] += 1

        if nozzle != 'A' and component_data.loc[part_index, 'nz'] != nozzle:
            warning_info = 'the nozzle type of component ' + part + ' is not consistent with the pcb data'
            warnings.warn(warning_info, UserWarning)
    # 清除点数为0的数据
    component_data = component_data[component_data['points'] != 0].reset_index(drop=True)
    for idx, data in component_data.iterrows():
        if data['fdr'][0:3] == 'SME':       # 电动供料器和气动供料器参数一致
            component_data.at[idx, 'fdr'] = data['fdr'][0:2] + data['fdr'][3:]

    # 读取供料器基座数据
    feeder_data = pd.DataFrame(columns=['slot', 'part', 'arg'])      # arg表示是否为预分配，不表示分配数目
    if load_feeder_data:
        for _, data in pcb_data.iterrows():
            slot, part = data['fdr'].split(' ')
            if slot[0] != 'F' and slot[0] != 'R':
                continue
            slot = int(slot[1:]) if slot[0] == 'F' else int(slot[1:]) + max_slot_index // 2
            feeder_data = pd.concat([feeder_data, pd.DataFrame([slot, part, 1]).T])

        feeder_data.drop_duplicates(subset='slot', inplace=True, ignore_index=True)
        # 随机移除部分已安装的供料器
        if load_feeder_data == 2:
            drop_index = random.sample(list(range(len(feeder_data))), len(feeder_data) // 2)
            feeder_data.drop(index=drop_index, inplace=True)

        feeder_data.sort_values(by='slot', ascending=True, inplace=True, ignore_index=True)

    pcb_data = pcb_data.sort_values(by="x", ascending=False)
    return pcb_data, component_data, feeder_data