import copy

import pandas as pd

from base_optimizer.optimizer_common import *
from base_optimizer.result_analysis 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', 600)

    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)

    average_pos = 0
    for _, data in pcb_data.iterrows():
        average_pos += data.x
    slot_start = int(round(average_pos / len(pcb_data) + stopper_pos[0] - slotf1_pos[0]) / slot_interval) + 1

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

    H = max_head_index
    I = len(component_data)
    S = min(len(component_data) * ratio, 60)
    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 - 1), list_range(H), list_range(M), vtype=GRB.CONTINUOUS)
    # u = mdl.addVars(list_range(K), list_range(M), vtype=GRB.INTEGER)

    d_plus = mdl.addVars(list_range(J), list_range(H), list_range(K - 1), list_range(M), vtype=GRB.CONTINUOUS)
    d_minus = mdl.addVars(list_range(J), list_range(H), list_range(K - 1), list_range(M), vtype=GRB.CONTINUOUS)

    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='')
    x = mdl.addVars(list_range(I), list_range(S), list_range(K), list_range(H), list_range(M), vtype=GRB.BINARY)
    n = mdl.addVars(list_range(H), list_range(M), 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(x[i, s, k, h, m] for i in range(I) for s in range(S)) <= 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(CompOfNozzle[i][j] * x[i, s, k, h, m] for i in range(I) for s in range(S) for j in range(J)) <= 1 for k
        in range(K) for h in range(H) for m in range(M))

    # nozzle available number constraint
    mdl.addConstrs(
        quicksum(CompOfNozzle[i][j] * x[i, s, k, h, m] for i in range(I) for s in range(S) for h in range(H)) <= H for k
        in range(K) for j in range(J) for m in range(M))

    # work completion
    mdl.addConstrs(
        quicksum(x[i, s, k, h, m] for s in range(S) 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(quicksum(CompOfNozzle[i][j] * x[i, s, k, h, m] for i in range(I) for s in range(S)) - quicksum(
        CompOfNozzle[i][j] * x[i, s, k + 1, h, m] for i in range(I) for s in range(S)) == d_plus[j, h, k, m] - d_minus[
                       j, h, k, m] for k in range(K - 1) for j in range(J) for h in range(H) for m in range(M))

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

    mdl.addConstrs(n[h, m] == quicksum(d[k, h, m] for k in range(K - 1)) - 0.5 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(x[i, 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(x[i, s + h * ratio, k, h, m] for h in rng for i in range(I)) >= e[s, k, m] for m in range(M))

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

    # feeder related
    mdl.addConstrs(quicksum(f[s, i, m] for s in range(S) for m in range(M)) <= 1 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(x[i, 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(x[i, 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))
    mdl.addConstrs(
        quicksum(f[s, i, m] for i in range(I)) >= quicksum(f[s + 1, i, m] for i in range(I)) for s in range(S - 1) for m
        in range(M))

    # objective
    T_cy, T_nz, T_pu, T_pl = 2, 3, 1, 1
    mdl.addConstrs(obj >= T_cy * quicksum(g[k, m] for k in range(K)) + T_nz * quicksum(
        d[k, h, m] for h in range(H) for k in range(K - 1)) + T_pl * quicksum(
        e[s, k, m] for s in range(-(H - 1) * ratio, S) for k in range(K)) + T_pl * quicksum(
        x[i, s, k, h, m] for i in range(I) for s in range(S) for k in range(K) for h in range(H)) for m in range(M))

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

    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
        part_index = defaultdict(int)
        for idx, data in component_data.iterrows():
            part_index[data.part] = idx

        component_result, cycle_result, feeder_slot_result = [], [], []
        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):
                    for s in range(S):
                        if abs(x[i, s, k, h, m].x) < 1e-3:
                            continue
                        if component_result[-1][h] != -1:
                            assert 1
                        component_result[-1][h] = i
                        feeder_slot_result[-1][h] = slot_start + s * 2

                        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])
                        pcb_part_indices[component_data.iloc[i].part].pop(0)
                        partial_component_data.loc[i, 'points'] += 1
        print(component_result)
        print(cycle_result)
        print(feeder_slot_result)
        placement_result, head_sequence = greedy_placement_route_generation(partial_component_data, partial_pcb_data,
                                                                            component_result, cycle_result,
                                                                            feeder_slot_result, hinter=False)
        print('----- Placement machine ' + str(m + 1) + ' ----- ')
        info = placement_info_evaluation(partial_component_data, partial_pcb_data, component_result, cycle_result,
                                         feeder_slot_result, placement_result, head_sequence, hinter=False)
        optimization_assign_result(partial_component_data, partial_pcb_data, component_result, cycle_result,
                                   feeder_slot_result, nozzle_hinter=True, component_hinter=True, feeder_hinter=True)
        info.print()
        assembly_info.append(info)
        print('------------------------------ ')
    return assembly_info