import copy

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


class TwoPhaseOpt(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)
        self.reduction = True
        self.partition = True
        self.initial = False
        
    def optimize(self, hinter=True):
        # data preparation: convert data to index
        part_list, nozzle_list = defaultdict(int), defaultdict(int)
        part_feeder = defaultdict(int)
        cpidx_2_part, nzidx_2_nozzle, cpidx_2_nzidx = {}, {}, {}
        arg_slot_rng = None if len(self.feeder_data) == 0 else [self.feeder_data.iloc[0].slot, self.feeder_data.iloc[-1].slot]
        for idx, data in self.part_data.iterrows():
            part, nozzle = data.part, data.nz

            cpidx_2_part[idx] = part
            nz_key = [key for key, val in nzidx_2_nozzle.items() if val == nozzle]

            nz_idx = len(nzidx_2_nozzle) if len(nz_key) == 0 else nz_key[0]
            nzidx_2_nozzle[nz_idx] = nozzle

            part_list[part] = 0
            part_feeder[part] = data.fdn
            cpidx_2_nzidx[idx] = nz_idx

        for _, data in self.step_data.iterrows():
            idx = self.part_data[self.part_data.part == data.part].index.tolist()[0]
            nozzle = self.part_data.loc[idx].nz

            nozzle_list[nozzle] += 1
            part_list[data.part] += 1

        part_feederbase = defaultdict(int)
        if self.feeder_data is not None:
            for _, data in self.feeder_data.iterrows():
                idx = -1
                for idx, part_ in cpidx_2_part.items():
                    if data.part == part_:
                        break
                assert idx != -1
                part_feederbase[idx] = data.slot  # part index - slot

        ratio = 1 if self.reduction else 2
        I, J = len(cpidx_2_part.keys()), len(nzidx_2_nozzle.keys())
        # === determine the hyper-parameter of L ===
        # first phase: calculate the number of heads for each type of nozzle
        nozzle_heads = defaultdict(int)
        for nozzle in nozzle_list.keys():
            nozzle_heads[nozzle] = 1
            
        head_num = self.config.head_num
        while sum(nozzle_heads.values()) != head_num:
            max_cycle_nozzle = None

            for nozzle, head_num in nozzle_heads.items():
                if max_cycle_nozzle is None or nozzle_list[nozzle] / head_num > nozzle_list[max_cycle_nozzle] / \
                        nozzle_heads[max_cycle_nozzle]:
                    max_cycle_nozzle = nozzle

            assert max_cycle_nozzle is not None
            nozzle_heads[max_cycle_nozzle] += 1

        nozzle_comp_points = defaultdict(list)
        for part, points in part_list.items():
            idx = self.part_data[self.part_data.part == part].index.tolist()[0]
            nozzle = self.part_data.loc[idx].nz
            nozzle_comp_points[nozzle].append([part, points])

        level = 1 if len(part_list) == 1 or len(part_list) % head_num == 0 else 2
        part_assignment, cycle_assignment = [], []

        def aux_func(info):
            return max(map(lambda points: max([p[1] for p in points]), info))

        pre_objbst, pre_changetime = None, None

        def terminate_condition(mdl, where):
            if where == GRB.Callback.MIP:
                objbst, objbnd = mdl.cbGet(GRB.Callback.MIP_OBJBST), mdl.cbGet(GRB.Callback.MIP_OBJBND)
                changetime = mdl.cbGet(GRB.Callback.RUNTIME)
                nonlocal pre_objbst, pre_changetime
                #            condition: value change
                if abs(objbst - 1e+100) > 1:  # 避免未找到可行解提前退出
                    if pre_objbst and abs(pre_objbst - objbst) < 1e-3:
                        if pre_changetime and changetime - pre_changetime > 90 * (1 - objbnd / objbst):
                            mdl.terminate()
                    else:
                        pre_changetime = changetime

                pre_objbst = objbst

        def recursive_assign(assign_points, nozzle_compo_points, cur_level, total_level) -> int:
            def func(points):
                return map(lambda points: max([p[1] for p in points]), points)

            if cur_level > total_level and sum(func(nozzle_compo_points.values())) == 0:
                return 0
            elif assign_points <= 0 and cur_level == 1:
                return -1  # backtrack
            elif assign_points <= 0 or cur_level > total_level:
                return 1  # fail

            nozzle_compo_points_cpy = copy.deepcopy(nozzle_compo_points)
            prev_assign = 0
            for part in part_assignment[cur_level - 1]:
                if part != -1:
                    prev_assign += 1

            head_idx = 0
            for nozzle, head in nozzle_heads.items():
                while head:
                    min_idx = -1
                    for idx, (part, points) in enumerate(nozzle_compo_points_cpy[nozzle]):
                        if points >= assign_points and (
                                min_idx == -1 or points < nozzle_compo_points_cpy[nozzle][min_idx][1]):
                            min_idx = idx
                    part_assignment[cur_level - 1][head_idx] = -1 if min_idx == -1 else \
                        nozzle_compo_points_cpy[nozzle][min_idx][0]
                    if min_idx != -1:
                        nozzle_compo_points_cpy[nozzle][min_idx][1] -= assign_points
                    head -= 1
                    head_idx += 1

            cycle_assignment[cur_level - 1] = assign_points
            for part in part_assignment[cur_level - 1]:
                if part != -1:
                    prev_assign -= 1

            if prev_assign == 0:
                res = 1
            else:
                points = min(len(self.step_data) // head_num + 1, aux_func(nozzle_compo_points_cpy.values()))
                res = recursive_assign(points, nozzle_compo_points_cpy, cur_level + 1, total_level)
            if res == 0:
                return 0
            elif res == 1:
                # All cycles have been completed, but there are still points left to be allocated
                return recursive_assign(assign_points - 1, nozzle_compo_points, cur_level, total_level)

        # second phase: (greedy) recursive search to assign points for each cycle set and obtain an initial solution
        while True:
            part_assignment = [[-1 for _ in range(head_num)] for _ in range(level)]
            cycle_assignment = [-1 for _ in range(level)]
            points = min(len(self.step_data) // head_num + 1, max(part_list.values()))
            if recursive_assign(points, nozzle_comp_points, 1, level) == 0:
                break
            level += 1

        L = len(cycle_assignment) if self.partition else len(self.step_data)
        S = ratio * sum(part_feeder.values()) if len(self.feeder_data) == 0 else arg_slot_rng[-1] - arg_slot_rng[
            0] + 1  # the available feeder num
        M = len(self.step_data)  # a sufficiently large number (number of placement points)
        HC = [[0 for _ in range(J)] for _ in range(I)]
        for i in range(I):
            for j in range(J):
                HC[i][j] = 1 if cpidx_2_nzidx[i] == j else 0

        mdl = Model('SMT')
        mdl.setParam('Seed', 0)
        mdl.setParam('OutputFlag', hinter)  # set whether output the debug information
        mdl.setParam('TimeLimit', 3600)
        mdl.setParam('PoolSearchMode', 2)
        mdl.setParam('PoolSolutions', 100)
        mdl.setParam('PoolGap', 1e-4)
        mdl.setParam("Heuristics", 0.5)

        # Use only if other methods, including exploring the tree with the default settings, do not yield a viable solution
        # mdl.setParam("ZeroObjNodes", 100)

        # === Decision Variables ===
        H = head_num
        x = mdl.addVars(I, S, H, L, vtype=GRB.BINARY, name='x')
        y = mdl.addVars(I, H, L, vtype=GRB.BINARY, name='y')
        v = mdl.addVars(S, H, L, vtype=GRB.BINARY, name='v')
        c = mdl.addVars(I, H, L, vtype=GRB.INTEGER, name='c')

        mdl.addConstrs(c[i, h, l] <= part_list[cpidx_2_part[i]] for i in range(I) for h in range(H) for l in range(L))

        f = {}
        for i in range(I):
            if i not in part_feederbase.keys():
                for s in range(S):
                    f[s, i] = mdl.addVar(vtype=GRB.BINARY, name='f_' + str(s) + '_' + str(i))
            else:
                for s in range(S):
                    f[s, i] = 1 if part_feederbase[i] == s + arg_slot_rng[0] else 0

        p = mdl.addVars(S + (H - 1) * ratio, L, vtype=GRB.BINARY, name='p')
        z = mdl.addVars(J, H, L, vtype=GRB.BINARY)

        d = mdl.addVars(L, H, vtype=GRB.INTEGER, name='d')
        d_plus = mdl.addVars(J, H, L, vtype=GRB.INTEGER, name='d_plus')
        d_minus = mdl.addVars(J, H, L, vtype=GRB.INTEGER, name='d_minus')

        max_cycle = math.ceil(len(self.step_data) / H)
        PU = mdl.addVars(S + (H - 1) * ratio, L, vtype=GRB.INTEGER, name='PU')
        WL = mdl.addVars(L, vtype=GRB.INTEGER, ub=max_cycle, name='WL')
        NC = mdl.addVars(H, vtype=GRB.INTEGER, name='NC')

        part_2_cpidx = defaultdict(int)
        for idx, part in cpidx_2_part.items():
            part_2_cpidx[part] = idx

        if self.initial:
            # initial some variables to speed up the search process
            # ensure the priority of the workload assignment
            cycle_index = sorted(range(len(cycle_assignment)), key=lambda k: cycle_assignment[k], reverse=True)
            part_list = []

            for cycle in cycle_index:
                cycle_part = part_assignment[cycle]
                for part in cycle_part:
                    if part != -1 and part not in part_list:
                        part_list.append(part)
            slot = 0
            for part in part_list:
                if self.feeder_data is not None:
                    while slot in self.feeder_data.keys():
                        slot += 1  # skip assigned feeder slot

                if part_2_cpidx[part] in part_feederbase.keys():
                    continue

                part_feederbase[part_2_cpidx[part]] = slot
                f[slot, part_2_cpidx[part]].Start = 1
                slot += 1

            for idx, cycle in enumerate(cycle_index):
                WL[idx].Start = cycle_assignment[cycle]
                for h in range(H):
                    part = part_assignment[cycle][h]
                    if part == -1:
                        continue
                    i = part_2_cpidx[part]
                    y[i, h, idx].Start = 1
                    v[part_feederbase[i], h, idx].Start = 1

        # === Objective ===
        mdl.setObjective(self.cycle_weight * quicksum(WL[l] for l in range(L)) +
                         2 * self.nozzle_change_weight * quicksum(NC[h] for h in range(H)) +
                         self.pickup_weight * quicksum(PU[s, l] for s in range(-(H - 1) * ratio, S) for l in range(L)))

        # === Constraint ===
        if not self.partition:
            mdl.addConstrs(WL[l] <= 1 for l in range(L))

        # work completion
        mdl.addConstrs(c[i, h, l] == WL[l] * y[i, h, l] for i in range(I) for h in range(H) for l in range(L))
        mdl.addConstrs(c[i, h, l] <= max_cycle * y[i, h, l] for i in range(I) for h in range(H) for l in range(L))
        mdl.addConstrs(c[i, h, l] <= WL[l] for i in range(I) for h in range(H) for l in range(L))
        mdl.addConstrs(c[i, h, l] >= WL[l] - max_cycle * (1 - y[i, h, l]) for i in range(I) for h in range(H)
                       for l in range(L))

        mdl.addConstrs(quicksum(c[i, h, l] for h in range(H) for l in range(L)) == part_list[cpidx_2_part[i]]
                       for i in range(I))

        # variable constraint
        mdl.addConstrs(quicksum(y[i, h, l] for i in range(I)) <= 1 for h in range(H) for l in range(L))

        # simultaneous pick
        for s in range(S + (H - 1) * ratio):
            rng = list(range(max(0, -math.floor(s / ratio)), min(H, math.ceil((S - s) / ratio))))
            for l in range(L):
                mdl.addConstr(quicksum(v[s + h * ratio, h, l] for h in rng) <= H * p[s, l])
                mdl.addConstr(quicksum(v[s + h * ratio, h, l] for h in rng) >= p[s, l])

        mdl.addConstrs(PU[s, l] == p[s, l] * WL[l] for s in range(S + (H - 1) * ratio) for l in range(L))
        # mdl.addConstrs(PU[s, l] <= max_cycle * p[s, l] for s in range(S + (H - 1) * ratio) for l in range(L))
        # mdl.addConstrs(PU[s, l] <= WL[l] for s in range(-(H - 1) * ratio, S) for l in range(L))
        # mdl.addConstrs(PU[s, l] >= WL[l] - max_cycle * (1 - p[s, l]) for s in range(-(H - 1) * ratio, S) for l in
        #                range(L))

        # nozzle change
        mdl.addConstrs(
            z[j, h, l] - z[j, h, l + 1] == d_plus[j, h, l] - d_minus[j, h, l] for l in range(L - 1) for j in range(J)
            for h in range(H))

        mdl.addConstrs(z[j, h, 0] - z[j, h, L - 1] == d_plus[j, h, L - 1] - d_minus[j, h, L - 1] for j in range(J)
                       for h in range(H))

        mdl.addConstrs(
            2 * d[l, h] == quicksum(d_plus[j, h, l] for j in range(J)) + quicksum(d_minus[j, h, l] for j in range(J))
            for l in range(L) for h in range(H))

        mdl.addConstrs(NC[h] == quicksum(d[l, h] for l in range(L)) for h in range(H))
        mdl.addConstrs(quicksum(y[i, h, l] for i in range(I) for h in range(H)) * M >= WL[l] for l in range(L))

        # nozzle-component compatibility
        mdl.addConstrs(y[i, h, l] <= quicksum(HC[i][j] * z[j, h, l] for j in range(J)) for i in range(I) for h in
                       range(H) for l in range(L))

        # available number of feeder
        mdl.addConstrs(quicksum(f[s, i] for s in range(S)) <= part_feeder[cpidx_2_part[i]] for i in range(I))

        # available number of nozzle
        mdl.addConstrs(quicksum(z[j, h, l] for h in range(H)) <= H for j in range(J) for l in range(L))

        # upper limit for occupation for feeder slot
        mdl.addConstrs(quicksum(f[s, i] for i in range(I)) <= 1 for s in range(S))
        mdl.addConstrs(quicksum(v[s, h, l] for s in range(S)) >= quicksum(y[i, h, l] for i in range(I)) for h in
                       range(H) for l in range(L))

        # others
        mdl.addConstrs(quicksum(z[j, h, l] for j in range(J)) <= 1 for h in range(H) for l in range(L))
        mdl.addConstrs(quicksum(x[i, s, h, l] for h in range(H) for l in range(L)) >= f[s, i] for i in range(I)
                       for s in range(S))
        mdl.addConstrs(quicksum(x[i, s, h, l] for h in range(H) for l in range(L)) <= M * f[s, i] for i in
                       range(I) for s in range(S))

        # mdl.addConstrs(f[s, i] >= x[i, s, h, l] for s in range(S) for i in range(I) for h in range(H) for l in range(L))
        #
        # mdl.addConstrs(quicksum(x[i, s, h, l] for h in range(H) for l in range(L)) >= f[s, i] for s in
        #                range(S) for i in range(I))

        # the constraints to speed up the search process
        mdl.addConstrs(quicksum(x[i, s, h, l] for i in range(I) for s in range(S)) <= 1 for h in range(H) for l
                       in range(L))

        if self.reduction:
            mdl.addConstrs(WL[l] >= WL[l + 1] for l in range(L - 1))
            mdl.addConstr(quicksum(WL[l] for l in range(L)) <= sum(cycle_assignment))
            mdl.addConstr(quicksum(WL[l] for l in range(L)) >= math.ceil(len(self.step_data) / H))
            mdl.addConstrs(quicksum(z[j, h, l] for j in range(J) for h in range(H)) >= quicksum(
                z[j, h, l + 1] for j in range(J) for h in range(H)) for l in range(L - 1))

        mdl.addConstrs(y[i, h, l] <= WL[l] for i in range(I) for h in range(H) for l in range(L))
        mdl.addConstrs(v[s, h, l] <= WL[l] for s in range(S) for h in range(H) for l in range(L))

        mdl.addConstrs(x[i, s, h, l] >= y[i, h, l] + v[s, h, l] - 1 for i in range(I) for s in range(S) for h in range(H)
                       for l in range(L))
        mdl.addConstrs(x[i, s, h, l] <= y[i, h, l] for i in range(I) for s in range(S) for h in range(H)
                       for l in range(L))

        mdl.addConstrs(x[i, s, h, l] <= v[s, h, l] for i in range(I) for s in range(S) for h in range(H)
                       for l in range(L))

        # === search process ===
        # mdl.update()
        # mdl.write('mdl.lp')

        mdl.optimize(terminate_condition)
        # mdl.optimize()

        # === result generation ===
        opt_res_list = defaultdict(OptResult)
        if mdl.Status == GRB.OPTIMAL or mdl.Status == GRB.INTERRUPTED or mdl.Status == GRB.TIME_LIMIT:
            # === selection from solution pool ===
            component_pos = defaultdict(list[Point])
            for _, data in self.step_data.iterrows():
                component_index = self.part_data[self.part_data.part == data.part].index.tolist()[0]
                component_pos[component_index].append(Point(data.x, data.y))

            for part in component_pos.keys():
                component_pos[part] = sorted(component_pos[part], key=lambda pos: (pos.x, pos.y))

            for sol_counter in range(mdl.SolCount):
                mdl.Params.SolutionNumber = sol_counter
                opt_res = OptResult()
                # == 转换标准的贴装头分配的解 ===
                for l in range(L):
                    if abs(WL[l].Xn) <= 1e-4:
                        continue
                    opt_res.cycle.append(round(WL[l].Xn))
                    opt_res.part.append([-1] * head_num)
                    opt_res.slot.append([-1] * head_num)

                    for h in range(head_num):
                        for i in range(I):
                            if abs(y[i, h, l].Xn) <= 1e-4:
                                continue
                            opt_res.part[-1][h] = i

                        for s in range(S):
                            if abs(v[s, h, l].Xn - 1) < 1e-4 and opt_res.part[-1][h] != -1:
                                opt_res.slot[-1][h] = s

                # 根据贴装头位置，转换供料器槽位
                cp_avg_head, cp_sum_cycle = defaultdict(float), defaultdict(int)
                for cycle, component_assign in enumerate(opt_res.part):
                    for head, part in enumerate(component_assign):
                        if part == -1:
                            continue
                        cp_avg_head[part] += opt_res.cycle[cycle] * head
                        cp_sum_cycle[part] += opt_res.cycle[cycle]

                for part, head in cp_avg_head.items():
                    cp_avg_head[part] = head / cp_sum_cycle[part]

                avg_position = sum([data.x - cp_avg_head[part_2_cpidx[data.part]] * self.config.head_intv for _, data in
                                    self.step_data.iterrows()]) / len(self.step_data)
                avg_slot = 0
                D_PU, D_PL, D_BW, D_FW = 0, 0, 0, 0
                for cycle, slots in enumerate(opt_res.slot):
                    min_slot, max_slot = self.config.slot_num, 0
                    for head, slot in enumerate(slots):
                        if slot == -1:
                            continue
                        min_slot = min(min_slot, slot - head * ratio)
                        max_slot = max(max_slot, slot - head * ratio)
                    avg_slot += (max_slot - min_slot) * opt_res.cycle[cycle]
                    D_PU += (max_slot - min_slot) * self.config.slot_intv * opt_res.cycle[cycle]  # 拾取路径

                avg_slot /= sum(opt_res.cycle)
                start_slot = round((avg_position + self.config.stopper_pos.x - self.config.slotf1_pos.x)
                                   / self.config.slot_intv + avg_slot / 2) + 1

                for cycle in range(len(opt_res.slot)):
                    for head in range(head_num):
                        if (slot := opt_res.slot[cycle][head]) == -1:
                            continue
                        opt_res.slot[cycle][head] = start_slot + slot * (2 if ratio == 1 else 1)

                component_pos_counter = defaultdict(int)
                cycle_place_pos = defaultdict(list[Point])
                for head in range(head_num):
                    for cycle in range(len(opt_res.cycle)):
                        if (part := opt_res.part[cycle][head]) == -1:
                            continue

                        avg_place_pos = Point(0, 0, _h=head)
                        for counter in range(round(opt_res.cycle[cycle])):
                            avg_place_pos.x = (1 - 1.0 / (counter + 1)) * avg_place_pos.x + (
                                    component_pos[part][component_pos_counter[part]].x - head * self.config.head_intv) / \
                                              (counter + 1)
                            avg_place_pos.y = (1 - 1.0 / (counter + 1)) * avg_place_pos.y + component_pos[part][
                                component_pos_counter[part]].y / (counter + 1)
                            component_pos_counter[part] += 1
                        avg_place_pos.x += self.config.stopper_pos.x
                        avg_place_pos.y += self.config.stopper_pos.y
                        cycle_place_pos[cycle].append(avg_place_pos)

                intv_ratio = self.config.head_intv // self.config.slot_intv
                for cycle in range(len(opt_res.cycle)):
                    min_slot, max_slot = self.config.slot_num, 0
                    for head in range(head_num):
                        if (slot := opt_res.slot[cycle][head]) == -1:
                            continue
                        min_slot = min(min_slot, slot - head * intv_ratio)
                        max_slot = max(max_slot, slot - head * intv_ratio)
                    # cycle_place_pos[cycle] = sorted(cycle_place_pos[cycle], key=lambda pt: pt.x)

                    pick_pos = copy.deepcopy(self.config.slotf1_pos)
                    pick_pos.x += (min_slot + max_slot) / 2 * self.config.slot_intv
                    _, seq = self.path_planner.dynamic_programming_cycle_path(cycle_place_pos[cycle], pick_pos)
                    head_position = [Point(0, 0) for _ in range(head_num)]
                    for point in cycle_place_pos[cycle]:
                        head_position[point.h] = point

                    for idx in range(len(seq) - 1):
                        h1, h2 = seq[idx], seq[idx + 1]
                        D_PL += max(abs(head_position[h1].x - head_position[h2].x),
                                    abs(head_position[h1].y - head_position[h2].y)) * opt_res.cycle[cycle]

                opt_res_list[sol_counter] = opt_res

        solution_number = 0
        # mdl.Params.SolutionNumber = 0
        if hinter:
            print('total cost = {}'.format(mdl.objval))
            print('cycle = {}, nozzle change = {}, pick up = {}'.format(quicksum(WL[l].Xn for l in range(L)), quicksum(
                NC[h].Xn for h in range(head_num)), quicksum(
                PU[s, l].Xn for s in range(-(head_num - 1) * ratio, S) for l in range(L))))

            print('workload: ')
            for l in range(L):
                print(WL[l].Xn, end=', ')

            print('')
            print('result')
            print('component assignment: ', opt_res_list[solution_number].part)
            print('feeder assignment: ', opt_res_list[solution_number].slot)
            print('cycle assignment: ', opt_res_list[solution_number].cycle)

        self.result = opt_res_list[solution_number]

