from base_optimizer.optimizer_common import *
from estimator import *


def random_component_assignment(pcb_data, component_data, machine_number, estimator=None):
    # == the set of feasible component type for each nozzle type
    nozzle_part_list = defaultdict(list)
    component_points = []
    for idx, data in component_data.iterrows():
        component_points.append(data.points)
        nozzle_part_list[data.nz].append(idx)

    component_number = len(component_data)
    assignment_result = [[0 for _ in range(component_number)] for _ in range(machine_number)]

    # === ensure every nozzle types ===
    selected_part = []
    for part_list in nozzle_part_list.values():
        part = random.sample(part_list, 1)[0]
        machine_index = random.randint(0, machine_number - 1)

        assignment_result[machine_index][part] += 1
        component_points[part] -= 1
        selected_part.append(part)

    # === assign one placement which has not been selected ===
    for part in range(component_number):
        if part in selected_part:
            continue

        assignment_result[random.randint(0, machine_number - 1)][part] += 1
        component_points[part] -= 1

    machine_assign = list(range(machine_number))
    random.shuffle(machine_assign)
    finished_assign_counter = 0
    while finished_assign_counter < component_number:
        for machine_index in machine_assign:
            part = random.randint(0, component_number - 1)
            feeder_counter = 0
            for idx in range(machine_number):
                if assignment_result[idx][part] > 0 or idx == machine_index:
                    feeder_counter += 1

            if component_points[part] == 0 or feeder_counter > component_data.iloc[part]['feeder-limit']:
                continue

            # feeder limit restriction
            points = random.randint(1, component_points[part])
            assignment_result[machine_index][part] += points
            component_points[part] -= points
            if component_points[part] == 0:
                finished_assign_counter += 1

    assert sum(component_points) == 0
    val = 0
    if estimator:
        cp_items = estimator.convert(pcb_data, component_data, assignment_result)
        for machine_index in range(machine_number):
            cp_points, cp_nozzle, cp_width, cp_height, board_width, board_height = cp_items[machine_index]
            # objective_value.append(
            #     estimator.neural_network(cp_points, cp_nozzle, cp_width, cp_height, board_width, board_height))
            val = max(val, estimator.heuristic(cp_points, cp_nozzle))

    return val, assignment_result


def greedy_component_assignment(component_points, component_nozzle, component_feeders, task_block_weight):
    pass    # 不清楚原文想说什么


def local_search_component_assignment(pcb_data, component_data, machine_number, estimator):
    # maximum number of iterations : 5000
    # maximum number of unsuccessful iterations: 50
    component_number = len(component_data)
    iteration_counter, unsuccessful_iteration_counter = 5000, 50
    optimal_val, optimal_assignment = random_component_assignment(pcb_data, component_data, machine_number, estimator)
    for _ in range(iteration_counter):
        machine_idx = random.randint(0, machine_number - 1)
        if sum(optimal_assignment[machine_idx]) == 0:
            continue

        part_set = []
        for part_idx in range(component_number):
            if optimal_assignment[machine_idx][part_idx] != 0:
                part_set.append(part_idx)
        part_idx = random.sample(part_set, 1)[0]
        r = random.randint(1, optimal_assignment[machine_idx][part_idx])

        assignment = copy.deepcopy(optimal_assignment)
        cyclic_counter = 0
        swap_machine_idx = None
        while cyclic_counter <= 2 * machine_idx:
            cyclic_counter += 1
            swap_machine_idx = random.randint(0, machine_number - 1)
            feeder_available = 0
            for machine in range(machine_number):
                if optimal_assignment[machine][part_idx] or machine == swap_machine_idx:
                    feeder_available += 1

            if feeder_available <= component_data.iloc[part_idx]['feeder-limit'] and swap_machine_idx != machine_idx:
                break
        assert swap_machine_idx is not None
        assignment[machine_idx][part_idx] -= r
        assignment[swap_machine_idx][part_idx] += r

        val = 0
        cp_items = estimator.convert(pcb_data, component_data, assignment)
        for machine_index in range(machine_number):
            cp_points, cp_nozzle, _, _, _, _ = cp_items[machine_index]
            val = max(val, estimator.heuristic(cp_points, cp_nozzle))

        if val < optimal_val:
            optimal_assignment, optimal_val = assignment, val
            unsuccessful_iteration_counter = 50
        else:
            unsuccessful_iteration_counter -= 1
            if unsuccessful_iteration_counter <= 0:
                break

    return optimal_val, optimal_assignment


def reconfig_crossover_operation(component_data, parent1, parent2, machine_number):
    offspring1, offspring2 = copy.deepcopy(parent1), copy.deepcopy(parent2)
    component_number = len(component_data)

    # === crossover ===
    mask_bit = []
    for _ in range(machine_number):
        mask_bit.append(random.randint(0, 1))
    if sum(mask_bit) == 0 or sum(mask_bit) == machine_number:
        return offspring1, offspring2

    for machine_index in range(machine_number):
        if mask_bit:
            offspring1[machine_index] = copy.deepcopy(parent1[machine_index])
            offspring2[machine_index] = copy.deepcopy(parent2[machine_index])
        else:
            offspring1[machine_index] = copy.deepcopy(parent2[machine_index])
            offspring2[machine_index] = copy.deepcopy(parent1[machine_index])

    # === balancing ===
    # equally to reach the correct number
    for part_index in range(component_number):
        for offspring in [offspring1, offspring2]:
            additional_points = sum([offspring[mt][part_index] for mt in range(machine_number)]) - \
                                component_data.iloc[part_index]['points']
            if additional_points > 0:
                # if a component type has more placements, decrease the assigned values on every head equally keeping
                # the proportion of the number of placement among the heads
                points_list = []
                for machine_index in range(machine_number):
                    points = math.floor(
                        additional_points * offspring[machine_index][part_index] / component_data[part_index]['points'])
                    points_list.append(points)
                    offspring[machine_index][part_index] -= points
                additional_points -= sum(points_list)

                for machine_index in range(machine_number):
                    if additional_points == 0:
                        break
                    if offspring[machine_index][part_index] == 0:
                        continue
                    offspring[machine_index][part_index] -= 1
                    additional_points += 1
            elif additional_points < 0:
                # otherwise, increase the assigned nonzero values equally
                machine_set = []
                for machine_index in range(machine_number):
                    if offspring[machine_index][part_index] == 0:
                        continue
                    machine_set.append(machine_index)

                points = -math.ceil(additional_points / len(machine_set))
                for machine_index in machine_set:
                    offspring[machine_index][part_index] += points
                    additional_points += points

                for machine_index in machine_set:
                    if additional_points == 0:
                        break
                    offspring[machine_index][part_index] += 1
                    additional_points -= 1

    # === 结果校验 ===
    for offspring in [offspring1, offspring2]:
        for part in range(component_number):
            pt = sum(offspring[mt][part] for mt in range(machine_number))
            assert pt == component_data.iloc[part]['points']

    return offspring1, offspring2


def reconfig_mutation_operation(component_data, parent, machine_number):
    offspring = copy.deepcopy(parent)

    swap_direction = random.randint(0, 1)
    if swap_direction:
        swap_machine1, swap_machine2 = random.sample(list(range(machine_number)), 2)
    else:
        swap_machine2, swap_machine1 = random.sample(list(range(machine_number)), 2)

    component_list = []
    for component_index, points in enumerate(offspring[swap_machine1]):
        if points:
            component_list.append(component_index)
    swap_component_index = random.sample(component_list, 1)[0]
    swap_points = random.randint(1, offspring[swap_machine1][swap_component_index])

    feeder_counter = 0
    for machine_index in range(machine_number):
        if offspring[swap_machine1][swap_component_index] < swap_points or machine_index == swap_machine2:
            feeder_counter += 1
    if feeder_counter > component_data.iloc[swap_component_index]['feeder-limit']:
        return offspring

    offspring[swap_machine1][swap_component_index] -= swap_points
    offspring[swap_machine2][swap_component_index] += swap_points
    return offspring


def evolutionary_component_assignment(pcb_data, component_data, machine_number, estimator):
    # population size: 10
    # probability of the mutation: 0.1
    # probability of the crossover: 0.8
    # number of generation: 100
    population_size = 10
    generation_number = 100
    mutation_rate, crossover_rate = 0.1, 0.8

    population = []
    for _ in range(population_size):
        population.append(random_component_assignment(pcb_data, component_data, machine_number, None)[1])

    with tqdm(total=generation_number) as pbar:
        pbar.set_description('evolutionary algorithm process for PCB assembly line balance')

        new_population = []
        for _ in range(generation_number):
            # calculate fitness value
            pop_val = []
            for individual in population:
                val = 0
                cp_items = estimator.convert(pcb_data, component_data, individual)
                for machine_index in range(machine_number):
                    cp_points, cp_nozzle, _, _, _, _ = cp_items[machine_index]
                    val = max(val, estimator.heuristic(cp_points, cp_nozzle))
                pop_val.append(val)

            select_index = get_top_k_value(pop_val, population_size - len(new_population), reverse=False)
            population = [population[idx] for idx in select_index]
            pop_val = [pop_val[idx] for idx in select_index]

            population += new_population
            for individual in new_population:
                cp_items = estimator.convert(pcb_data, component_data, individual)
                val = 0
                for machine_index in range(machine_index):
                    cp_points, cp_nozzle, _, _, _, _ = cp_items[machine_index]
                    val = max(val, estimator.heuristic(cp_points, cp_nozzle))
                pop_val.append(val)

            # min-max convert
            max_val = max(pop_val)
            pop_val_sel = list(map(lambda v: max_val - v, pop_val))
            sum_pop_val = sum(pop_val_sel) + 1e-10
            pop_val_sel = [v / sum_pop_val + 1e-3 for v in pop_val_sel]

            # crossover and mutation
            new_population = []
            for pop in range(population_size):
                if pop % 2 == 0 and np.random.random() < crossover_rate:
                    index1 = roulette_wheel_selection(pop_val_sel)
                    while True:
                        index2 = roulette_wheel_selection(pop_val_sel)
                        if index1 != index2:
                            break

                    offspring1, offspring2 = reconfig_crossover_operation(component_data, population[index1],
                                                                          population[index2], machine_number)

                    if np.random.random() < mutation_rate:
                        offspring1 = reconfig_mutation_operation(component_data, offspring1, machine_number)

                    if np.random.random() < mutation_rate:
                        offspring2 = reconfig_mutation_operation(component_data, offspring2, machine_number)

                    new_population.append(offspring1)
                    new_population.append(offspring2)

            pbar.update(1)

    return min(pop_val), population[np.argmin(pop_val)]


@timer_wrapper
def line_optimizer_reconfiguration(component_data, pcb_data, machine_number):
    # === assignment of heads to modules is omitted ===
    optimal_assignment, optimal_val = [], None
    estimator = Estimator(task_block_weight=5)    # element from list [0, 1, 2, 5, 10] task_block ~= cycle
    # === assignment of components to heads
    for i in range(5):
        if i == 0:
            # random
            val, assignment = random_component_assignment(pcb_data, component_data, machine_number, estimator)
        elif i == 1:
            # brute force
            # which is proved to be useless, since it only ran in reasonable time for the smaller test instances
            continue
        elif i == 2:
            # local search
            val, assignment = local_search_component_assignment(pcb_data, component_data, machine_number, estimator)
        elif i == 3:
            # evolutionary
            val, assignment = evolutionary_component_assignment(pcb_data, component_data, machine_number, estimator)
        else:
            # greedy: unclear description
            continue

        if optimal_val is None or val < optimal_val:
            optimal_val, optimal_assignment = val, assignment.copy()

    if optimal_val is None:
        raise Exception('no feasible solution! ')

    return optimal_assignment