# implementation of <<An integrated allocation method for the PCB assembly line balancing problem with nozzle changes>>
from base_optimizer.optimizer_common import *
from lineopt_hyperheuristic import *


def selective_initialization(component_points, component_feeders, population_size, machine_number):
    population = []    # population initialization
    for _ in range(population_size):
        individual = []
        for part_index, points in component_points:
            if points == 0:
                continue
            # 可用机器数
            avl_machine_num = random.randint(1, min(machine_number, component_feeders[part_index], points))

            selective_possibility = []
            for p in range(1, avl_machine_num + 1):
                selective_possibility.append(pow(2, avl_machine_num - p + 1))

            sel_machine_num = random_selective([p + 1 for p in range(avl_machine_num)], selective_possibility)  # 选择的机器数
            sel_machine_set = random.sample([p for p in range(machine_number)], sel_machine_num)

            sel_machine_points = [1 for _ in range(sel_machine_num)]
            for p in range(sel_machine_num - 1):
                if points == sum(sel_machine_points):
                    break
                assign_points = random.randint(1, points - sum(sel_machine_points))
                sel_machine_points[p] += assign_points

            if sum(sel_machine_points) < points:
                sel_machine_points[-1] += (points - sum(sel_machine_points))

            # code component allocation into chromosome
            for p in range(machine_number):
                if p in sel_machine_set:
                    individual += [0 for _ in range(sel_machine_points[0])]
                    sel_machine_points.pop(0)
                individual.append(1)
            individual.pop(-1)

        population.append(individual)
    return population


def selective_crossover(component_points, component_feeders, mother, father, machine_number, non_decelerating=True):
    assert len(mother) == len(father)

    offspring1, offspring2 = mother.copy(), father.copy()
    one_counter, feasible_cut_line = 0, []

    idx = 0
    for part_index, points in component_points.items():
        one_counter = 0

        idx_, mother_cut_line, father_cut_line = 0, [-1], [-1]
        for idx_, gene in enumerate(mother[idx: idx + points + machine_number - 1]):
            if gene:
                mother_cut_line.append(idx_)
        mother_cut_line.append(idx_ + 1)

        for idx_, gene in enumerate(father[idx: idx + points + machine_number - 1]):
            if gene:
                father_cut_line.append(idx_)
        father_cut_line.append(idx_ + 1)

        for offset in range(points + machine_number - 1):
            if mother[idx + offset] == 1:
                one_counter += 1
            if father[idx + offset] == 1:
                one_counter -= 1

            # first constraint: the total number of '1's (the number of partitions) in the chromosome is unchanged
            if one_counter != 0 or offset == 0 or offset == points + machine_number - 2:
                continue

            # the selected cut-line should guarantee there are the same or a larger number unassigned machine
            # for each component type
            n_bro, n_new = 0, 0
            if mother[idx + offset] and mother[idx + offset + 1]:
                n_bro += 1
            if father[idx + offset] and father[idx + offset + 1]:
                n_bro += 1
            if mother[idx + offset] and father[idx + offset + 1]:
                n_new += 1
            if father[idx + offset] and mother[idx + offset + 1]:
                n_new += 1

            # second constraint: non_decelerating or accelerating crossover
            # non_decelerating or accelerating means that the number of machine without workload is increased
            if n_new < n_bro or (n_new == n_bro and not non_decelerating):
                continue

            # third constraint (customized constraint):
            # no more than the maximum number of available machine for each component type
            new_mother_cut_line, new_father_cut_line = [], []
            for idx_ in range(machine_number + 1):
                if mother_cut_line[idx_] <= offset:
                    new_mother_cut_line.append(mother_cut_line[idx_])
                else:
                    new_father_cut_line.append(mother_cut_line[idx_])

                if father_cut_line[idx_] <= offset:
                    new_father_cut_line.append(father_cut_line[idx_])
                else:
                    new_mother_cut_line.append(father_cut_line[idx_])

            sorted(new_mother_cut_line, reverse=False)
            sorted(new_father_cut_line, reverse=False)
            n_mother_machine, n_father_machine = 0, 0

            for idx_ in range(machine_number):
                if new_mother_cut_line[idx_ + 1] - new_mother_cut_line[idx_] > 1:
                    n_mother_machine += 1

                if new_father_cut_line[idx_ + 1] - new_father_cut_line[idx_] > 1:
                    n_father_machine += 1

            if n_mother_machine > component_feeders[part_index] or n_father_machine > component_feeders[part_index]:
                continue

            feasible_cut_line.append(idx + offset)

        idx += (points + machine_number - 1)

    if len(feasible_cut_line) == 0:
        return offspring1, offspring2

    cut_line_idx = feasible_cut_line[random.randint(0, len(feasible_cut_line) - 1)]
    offspring1, offspring2 = mother[:cut_line_idx + 1] + father[cut_line_idx + 1:], father[:cut_line_idx + 1] + mother[
                                                                                                     cut_line_idx + 1:]
    return offspring1, offspring2


def cal_individual_val(component_points, component_nozzle, machine_number, individual, estimator):
    idx, objective_val = 0, []
    machine_component_points = [[] for _ in range(machine_number)]

    # decode the component allocation
    for part_index, points in component_points.items():
        component_gene = individual[idx: idx + points + machine_number - 1]
        machine_idx, component_counter = 0, 0
        for gene in component_gene:
            if gene:
                machine_component_points[machine_idx].append(component_counter)
                machine_idx += 1
                component_counter = 0
            else:
                component_counter += 1
        machine_component_points[-1].append(component_counter)
        idx += (points + machine_number - 1)

    objective_val = 0
    for machine_idx in range(machine_number):
        machine_points = sum(machine_component_points[machine_idx])         # num of placement points
        if machine_points == 0:
            continue

        cp_points, cp_nozzle = defaultdict(int), defaultdict(str)
        for part_index, points in enumerate(machine_component_points[machine_idx]):
            if points == 0:
                continue
            cp_points[part_index], cp_nozzle[part_index] = points, component_nozzle[part_index]
        objective_val = max(objective_val, estimator.predict(cp_points, cp_nozzle))
    return objective_val, machine_component_points


def individual_convert(component_points, individual):
    machine_number = len(individual)
    machine_component_points = [[] for _ in range(machine_number)]
    idx = 0
    # decode the component allocation
    for comp_idx, points in component_points:
        component_gene = individual[idx: idx + points + machine_number - 1]
        machine_idx, component_counter = 0, 0
        for gene in component_gene:
            if gene:
                machine_component_points[machine_idx].append(component_counter)
                machine_idx += 1
                component_counter = 0
            else:
                component_counter += 1
        machine_component_points[-1].append(component_counter)
        idx += (points + machine_number - 1)

    return machine_component_points


def line_optimizer_genetic(component_data, machine_number):
    # basic parameter
    # crossover rate & mutation rate: 80% & 10%
    # population size: 200
    # the number of generation: 500
    crossover_rate, mutation_rate = 0.8, 0.1
    population_size, n_generations = 200, 500

    estimator = HeuristicEstimator()
    # the number of placement points, the number of available feeders, and nozzle type of component respectively
    cp_points, cp_feeders, cp_nozzle = defaultdict(int), defaultdict(int), defaultdict(int)
    for part_index, data in component_data.iterrows():
        cp_points[part_index] += data.points
        cp_feeders[part_index], cp_nozzle[part_index] = data.fdn, data.nz

    # population initialization
    population = selective_initialization(sorted(cp_points.items(), key=lambda x: x[0]), cp_feeders, population_size,
                                          machine_number)
    with tqdm(total=n_generations) as pbar:
        pbar.set_description('genetic algorithm process for PCB assembly line balance')

        new_population = []
        for _ in range(n_generations):
            # calculate fitness value
            pop_val = []
            for individual in population:
                val, assigned_points = cal_individual_val(cp_points, cp_nozzle, machine_number, individual, estimator)
                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:
                val, _ = cal_individual_val(cp_points, cp_nozzle, machine_number, individual, estimator)
                pop_val.append(val)

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

            # 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)
                    while True:
                        index2 = roulette_wheel_selection(pop_val)
                        if index1 != index2:
                            break

                    offspring1, offspring2 = selective_crossover(cp_points, cp_feeders,
                                                                 population[index1], population[index2], machine_number)

                    if np.random.random() < mutation_rate:
                        offspring1 = constraint_swap_mutation(cp_points, offspring1, machine_number)

                    if np.random.random() < mutation_rate:
                        offspring2 = constraint_swap_mutation(cp_points, offspring2, machine_number)

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

            pbar.update(1)

    best_individual = population[np.argmax(pop_val)]
    val, assignment_result = cal_individual_val(cp_points, cp_nozzle, machine_number, best_individual, estimator)

    print('final value: ', val)
    # available feeder check
    for part_index, data in component_data.iterrows():
        feeder_limit = data.fdn
        for machine_index in range(machine_number):
            if assignment_result[machine_index][part_index]:
                feeder_limit -= 1
        assert feeder_limit >= 0

    return assignment_result