from base_optimizer.optimizer_common import *
from base_optimizer.result_analysis import *
from base_optimizer.smtopt_route import *


def line_optimizer_model(component_data, pcb_data, machine_num, hinter=True):
    mdl = Model('pcb assembly line optimizer')
    mdl.setParam('Seed', 0)
    mdl.setParam('OutputFlag', hinter)  # set whether output the debug information
    mdl.setParam('TimeLimit', 1000)

    nozzle_type, component_type = [], []
    for _, data in component_data.iterrows():
        if not data.nz in nozzle_type:
            nozzle_type.append(data.nz)
        component_type.append(data.part)

    ratio = 1
    J = len(nozzle_type)
    N = 10000
    M = machine_num

    H = max_head_index
    I = len(component_data)
    S = sum([data.fdn * ratio for _, data in component_data.iterrows()])
    K = math.ceil(len(pcb_data) * 1.0 / M) + 1

    # K = len(pcb_data)
    CompOfNozzle = [[0 for _ in range(J)] for _ in range(I)]  # Compatibility

    component_point = [0 for _ in range(I)]
    for idx, data in component_data.iterrows():
        nozzle = component_data.iloc[idx].nz
        CompOfNozzle[idx][nozzle_type.index(nozzle)] = 1
        component_point[idx] = data.points

    # objective related
    g = mdl.addVars(list_range(K), list_range(M), vtype=GRB.BINARY)
    d = mdl.addVars(list_range(K), list_range(H), list_range(M), vtype=GRB.INTEGER)
    u = mdl.addVars(list_range(I), list_range(K), list_range(H), list_range(M), vtype=GRB.BINARY)
    v = mdl.addVars(list_range(S), list_range(K), list_range(H), list_range(M), vtype=GRB.BINARY)
    w = mdl.addVars(list_range(J), list_range(K), list_range(H), list_range(M), vtype=GRB.BINARY)
    d_plus = mdl.addVars(list_range(J), list_range(K), list_range(H), list_range(M), vtype=GRB.CONTINUOUS)
    d_minus = mdl.addVars(list_range(J), list_range(K), list_range(H), list_range(M), vtype=GRB.CONTINUOUS)
    z = mdl.addVars(list_range(K), list_range(M), vtype=GRB.INTEGER)

    e = mdl.addVars(list_range(-(H - 1) * ratio, S), list_range(K), list_range(M), vtype=GRB.BINARY)
    f = mdl.addVars(list_range(S), list_range(I), list_range(M), vtype=GRB.BINARY, name='')
    t = mdl.addVars(list_range(M), lb=0, ub=N, vtype=GRB.CONTINUOUS)
    obj = mdl.addVar(lb=0, ub=N, vtype=GRB.CONTINUOUS)

    mdl.addConstrs(g[k, m] >= g[k + 1, m] for k in range(K - 1) for m in range(M))

    mdl.addConstrs(
        quicksum(u[i, k, h, m] for i in range(I)) <= g[k, m] for k in range(K) for h in range(H) for m in range(M))

    # nozzle no more than 1 for head h and cycle k
    mdl.addConstrs(quicksum(w[j, k, h, m] for j in range(J)) <= 1 for k in range(K) for h in range(H) for m in range(M))

    # work completion
    mdl.addConstrs(
        quicksum(u[i, k, h, m] for k in range(K) for h in range(H) for m in range(M)) == component_point[i] for i in
        range(I))

    # nozzle change
    mdl.addConstrs(
        u[i, k, h, m] <= quicksum(CompOfNozzle[i][j] * w[j, k, h, m] for j in range(J)) for i in range(I) for k in
        range(K) for h in range(H) for m in range(M))

    mdl.addConstrs(w[j, k, h, m] - w[j, k + 1, h, m] == d_plus[j, k, h, m] - d_minus[j, k, h, m] for k in
                   range(K - 1) for j in range(J) for h in range(H) for m in range(M))

    mdl.addConstrs(w[j, 0, h, m] - w[j, K - 1, h, m] == d_plus[j, K - 1, h, m] - d_minus[j, K - 1, h, m]
                   for j in range(J) for h in range(H) for m in range(M))

    mdl.addConstrs(d[k, h, m] == quicksum(d_plus[j, k, h, m] for j in range(J)) + quicksum(
        d_minus[j, k, h, m] for j in range(J)) for k in range(K) for h in range(H) for m in range(M))

    # simultaneous pick
    for s in range(-(H - 1) * ratio, S):
        rng = list(range(max(0, -math.floor(s / ratio)), min(H, math.ceil((S - s) / ratio))))
        for k in range(K):
            mdl.addConstrs(
                quicksum(u[i, k, h, m] * v[s + h * ratio, k, h, m] for h in rng for i in range(I)) <= N * e[s, k, m] for
                m in range(M))
            mdl.addConstrs(
                quicksum(u[i, k, h, m] * v[s + h * ratio, k, h, m] for h in rng for i in range(I)) >= e[s, k, m] for m
                in range(M))

    # head - feeder slot relationship
    mdl.addConstrs(
        quicksum(v[s, k, h, m] for s in range(S)) == quicksum(u[i, k, h, m] for i in range(I)) for h in range(H) for k
        in range(K) for m in range(M))

    # feeder related
    mdl.addConstrs(
        quicksum(f[s, i, m] for s in range(S) for m in range(M)) <= component_data.iloc[i].fdn for i in range(I))
    mdl.addConstrs(quicksum(f[s, i, m] for i in range(I)) <= 1 for s in range(S) for m in range(M))
    mdl.addConstrs(
        quicksum(u[i, k, h, m] * v[s, k, h, m] for h in range(H) for k in range(K)) >= f[s, i, m] for i in range(I) for
        s in range(S) for m in range(M))
    mdl.addConstrs(
        quicksum(u[i, k, h, m] * v[s, k, h, m] for h in range(H) for k in range(K)) <= N * f[s, i, m] for i in range(I)
        for s in range(S) for m in range(M))

    # pickup movement
    mdl.addConstrs(z[k, m] >= s1 * e[s1, k, m] - s2 * e[s2, k, m] + N * (e[s1, k, m] + e[s2, k, m] - 2) for s1 in
                   range(-(H - 1) * ratio, S) for s2 in range(-(H - 1) * ratio, S) for k in range(K) for m in range(M))

    # objective
    mdl.addConstrs(t[m] == Fit_cy * quicksum(g[k, m] for k in range(K)) + Fit_nz * quicksum(
                        d[k, h, m] for h in range(H) for k in range(K)) + Fit_pu * quicksum(
                        e[s, k, m] for s in range(-(H - 1) * ratio, S) for k in range(K)) + Fit_pl * quicksum(
                        u[i, k, h, m] for i in range(I) for k in range(K) for h in range(H)) + Fit_mv * quicksum(
                        z[k, m] for k in range(K)) for m in range(M))
    for m in range(M - 1):
        mdl.addConstr(t[m] >= t[m + 1])

    mdl.addConstrs(obj >= t[m] for m in range(M))

    mdl.setObjective(obj, GRB.MINIMIZE)
    mdl.optimize()

    for m in range(M):
        print(f'machine {m} : cycle : {sum(g[k, m].x for k in range(K))}, '
              f'nozzle change : {sum(d[k, h, m].x for h in range(H) for k in range(K))}, '
              f'pick up : {sum(e[s, k, m].x for s in range(-(H - 1) * ratio, S) for k in range(K))}, '
              f'placement : {sum(u[i, k, h, m].x for i in range(I) for k in range(K) for h in range(H))}, '
              f'pick movement : {sum(z[k, m].x for k in range(K))}')

    pcb_part_indices = defaultdict(list)
    for idx, data in pcb_data.iterrows():
        pcb_part_indices[data.part].append(idx)

    assembly_info = []
    for m in range(M):
        partial_component_data, partial_pcb_data = copy.deepcopy(component_data), pd.DataFrame(columns=pcb_data.columns)
        partial_component_data['points'] = 0

        component_result, cycle_result, feeder_slot_result = [], [], []
        head_place_pos = []
        for k in range(K):
            if abs(g[k, m].x) < 1e-3:
                continue
            component_result.append([-1 for _ in range(H)])
            cycle_result.append(1)
            feeder_slot_result.append([-1 for _ in range(H)])

            for h in range(H):
                for i in range(I):
                    if abs(u[i, k, h, m].x) < 1e-3:
                        continue

                    component_result[-1][h] = i

                    idx = pcb_part_indices[component_data.iloc[i].part][0]
                    partial_pcb_data = pd.concat([partial_pcb_data, pd.DataFrame(pcb_data.iloc[idx]).T])

                    head_place_pos.append(pcb_data.iloc[idx].x - h * head_interval)
                    pcb_part_indices[component_data.iloc[i].part].pop(0)
                    partial_component_data.loc[i, 'points'] += 1

                for s in range(S):
                    if abs(v[s, k, h, m].x) < 1e-3:
                        continue
                    feeder_slot_result[-1][h] = s
            if sum(component_result[-1]) == -max_head_index:
                component_result.pop(-1)
                cycle_result.pop(-1)
                feeder_slot_result.pop(-1)

        average_pos = round(
            (sum(head_place_pos) / len(head_place_pos) + stopper_pos[0] - slotf1_pos[0] + 1) / slot_interval)

        for k in range(len(feeder_slot_result)):
            for h in range(H):
                if feeder_slot_result[k][h] == -1:
                    continue
                feeder_slot_result[k][h] = feeder_slot_result[k][h] * 2 + average_pos

        placement_result, head_sequence = place_allocate_sequence_route_generation(partial_component_data,
                                                                                   partial_pcb_data,
                                                                                   component_result, cycle_result,
                                                                                   feeder_slot_result, hinter=False)

        opt_res = OptResult(component_result, cycle_result, feeder_slot_result, placement_result, head_sequence)
        info = placement_info_evaluation(partial_component_data, partial_pcb_data, opt_res, hinter=hinter)
        if hinter:
            print('----- Placement machine ' + str(m + 1) + ' ----- ')
            optimization_assign_result(partial_component_data, partial_pcb_data, opt_res, nozzle_hinter=True,
                                       component_hinter=True, feeder_hinter=True)
            info.print()
            assembly_info.append(info)
            print('------------------------------ ')
    return assembly_info