from opt.smm.basis import *
from opt.utils import *
from opt.smm.solver import *


class Aggregation(BaseOpt):
    def __init__(self, config, part_data, step_data, feeder_data=pd.DataFrame(columns=['slot', 'part'])):
        super().__init__(config, part_data, step_data, feeder_data)
        self.feeder_assigner = FeederAssignOpt(config, part_data, step_data)

    def optimize(self, hinter=True):
        # === phase 0: data preparation ===
        M = 1000  # a sufficient large number
        a, b = 1, 6  # coefficient

        part_list, nozzle_list = defaultdict(int), defaultdict(int)
        cpidx_2_part, nzidx_2_nozzle = {}, {}
        for _, data in self.step_data.iterrows():
            part = data.part
            if part not in cpidx_2_part.values():
                cpidx_2_part[len(cpidx_2_part)] = part

            part_list[part] += 1

            idx = self.part_data[self.part_data['part'] == part].index.tolist()[0]
            nozzle = self.part_data.loc[idx]['nz']
            if nozzle not in nzidx_2_nozzle.values():
                nzidx_2_nozzle[len(nzidx_2_nozzle)] = nozzle
            nozzle_list[nozzle] += 1

        I, J = len(part_list.keys()), len(nozzle_list.keys())  # the maximum number of part types and nozzle types
        L = I + 1  # the maximum number of batch level
        K = self.config.head_num  # the maximum number of heads
        HC = [[M for _ in range(J)] for _ in range(I)]  # represent the nozzle-part compatibility

        for i in range(I):
            for _, item in enumerate(cpidx_2_part.items()):
                index, part = item
                cp_idx = self.part_data[self.part_data['part'] == part].index.tolist()[0]
                nozzle = self.part_data.loc[cp_idx]['nz']

                for j in range(J):
                    if nzidx_2_nozzle[j] == nozzle:
                        HC[index][j] = 0

        # === phase 1: mathematical model solver ===
        mdl = Model('SMT')
        # mdl.setParam('OutputFlag', hinter)

        # === Decision Variables ===
        # the largest workload of all placement heads
        WL = mdl.addVar(vtype=GRB.INTEGER, lb=0, ub=len(self.step_data), name='WL')

        # the number of parts of type i that are placed by nozzle type j on placement head k
        X = mdl.addVars(I, J, K, vtype=GRB.INTEGER, ub=max(part_list.values()), name='X')

        # the total number of nozzle changes on placement head k
        N = mdl.addVars(K, vtype=GRB.INTEGER, name='N')

        # whether batch Xijk is placed on level l
        Z = mdl.addVars(I, J, L, K, vtype=GRB.BINARY, name='Z')

        # Dlk := 2 if a change of nozzles in the level l + 1 on placement head k
        # Dlk := 1 if there are no batches placed on levels higher than l
        # Dlk := 0 otherwise
        D = mdl.addVars(L, K, vtype=GRB.BINARY, name='D')
        D_plus = mdl.addVars(L, J, K, vtype=GRB.INTEGER, name='D_plus')
        D_minus = mdl.addVars(L, J, K, vtype=GRB.INTEGER, name='D_minus')

        # == Objective function ===
        mdl.setObjective(a * WL + b * quicksum(N[k] for k in range(K)), GRB.MINIMIZE)

        # === Constraint ===
        mdl.addConstrs(
            quicksum(X[i, j, k] for j in range(J) for k in range(K)) == part_list[cpidx_2_part[i]] for i in range(I))

        mdl.addConstrs(quicksum(X[i, j, k] for i in range(I) for j in range(J)) <= WL for k in range(K))

        mdl.addConstrs(
            X[i, j, k] <= M * quicksum(Z[i, j, l, k] for l in range(L)) for i in range(I) for j in range(J) for k in
            range(K))

        mdl.addConstrs(
            quicksum(Z[i, j, l, k] for l in range(L)) <= 1 for i in range(I) for j in range(J) for k in range(K))
        mdl.addConstrs(
            quicksum(Z[i, j, l, k] for l in range(L)) <= X[i, j, k] for i in range(I) for j in range(J) for k in
            range(K))

        mdl.addConstrs(quicksum(Z[i, j, l, k] for j in range(J) for i in range(I)) >= quicksum(
            Z[i, j, l + 1, k] for j in range(J) for i in range(I)) for k in range(K) for l in range(L - 1))

        mdl.addConstrs(
            quicksum(Z[i, j, l, k] for i in range(I) for j in range(J)) <= 1 for k in range(K) for l in range(L))
        mdl.addConstrs(D_plus[l, j, k] - D_minus[l, j, k] == quicksum(Z[i, j, l, k] for i in range(I)) - quicksum(
            Z[i, j, l + 1, k] for i in range(I)) for l in range(L - 1) for j in range(J) for k in range(K))

        mdl.addConstrs(
            D[l, k] == quicksum((D_plus[l, j, k] + D_minus[l, j, k]) for j in range(J)) for k in range(K) for l in
            range(L))

        # mdl.addConstrs(2 * N[k] == quicksum(D[l, k] for l in range(L)) - 1 for k in range(K))
        # mdl.addConstrs(
        #     0 >= quicksum(HC[i][j] * Z[i, j, l, k] for i in range(I) for j in range(J)) for l in range(L) for k in
        #     range(K))

        # === Main Process ===
        mdl.TimeLimit = 100
        mdl.optimize()
        if mdl.Status == GRB.OPTIMAL or mdl.Status == GRB.TIME_LIMIT:
            print('total cost = {}'.format(mdl.objval))

            # convert cp model solution to standard output
            model_cycle_result, model_part_result = [], []
            for l in range(L):
                model_part_result.append([None for _ in range(K)])
                model_cycle_result.append([0 for _ in range(K)])
                for k in range(K):
                    for i in range(I):
                        for j in range(J):
                            if abs(Z[i, j, l, k].x - 1) <= 1e-3:
                                model_part_result[-1][k] = cpidx_2_part[i]
                                model_cycle_result[-1][k] = round(X[i, j, k].x)

                # remove redundant term
                if sum(model_cycle_result[-1]) == 0:
                    model_part_result.pop()
                    model_cycle_result.pop()

            head_part_index = [0 for _ in range(self.config.head_num)]
            while True:
                head_cycle = []
                for head, index in enumerate(head_part_index):
                    head_cycle.append(model_cycle_result[index][head])

                if len([cycle for cycle in head_cycle if cycle > 0]) == 0:
                    break

                self.result.part.append([None for _ in range(self.config.head_num)])
                min_cycle = min([cycle for cycle in head_cycle if cycle > 0])
                for head, index in enumerate(head_part_index):
                    if model_cycle_result[index][head] != 0:
                        self.result.part[-1][head] = model_part_result[index][head]
                    else:
                        continue

                    model_cycle_result[index][head] -= min_cycle
                    if model_cycle_result[index][head] == 0 and index + 1 < len(model_cycle_result):
                        head_part_index[head] += 1

                self.result.cycle.append(min_cycle)

            part_2_index = {}
            for index, data in self.part_data.iterrows():
                part_2_index[data['part']] = index

            for cycle in range(len(self.result.part)):
                for head in range(self.config.head_num):
                    part = self.result.part[cycle][head]
                    self.result.part[cycle][head] = -1 if part is None else part_2_index[part]

            self.result.slot = self.feeder_assigner.do(self.result.part, self.result.cycle)
            # === phase 2: heuristic method ===
            self.result.point, self.result.sequence = self.path_planner.greedy_level_placing(self.result.part,
                                                                                             self.result.cycle,
                                                                                             self.result.slot)
        else:
            warnings.warn('No solution found!', UserWarning)





