from numpy import put, sum, append, where, zeros
from numpy.random import choice, seed


def get_first_random_solution(n, m):
    seed(42)
    solution = zeros(shape=n, dtype=bool)
    random_indices = choice(n, size=m, replace=False)
    put(solution, ind=random_indices, v=True)
    return solution


def get_genotype(element):
    genotype = where(element == True)
    return genotype[0]


def evaluate_element(element, data):
    fitness = []
    genotype = get_genotype(element)
    distances = data.query(f"source in @genotype and destination in @genotype")
    for item in genotype[:-1]:
        element_df = distances.query(f"source == {item} or destination == {item}")
        max_distance = element_df["distance"].astype(float).max()
        fitness = append(arr=fitness, values=max_distance)
        distances = distances.query(f"source != {item} and destination != {item}")
    return sum(fitness)


def element_in_dataframe(solution, element):
    duplicates = solution.query(
        f"(source == {element.source} and destination == {element.destination}) or (source == {element.destination} and destination == {element.source})"
    )
    return not duplicates.empty


def replace_worst_element(previous, data):
    solution = previous.copy()
    worst_index = solution["distance"].astype(float).idxmin()
    random_element = data.sample().squeeze()
    while element_in_dataframe(solution=solution, element=random_element):
        random_element = data.sample().squeeze()
    solution.loc[worst_index] = random_element
    return solution, worst_index


def get_random_solution(previous, data):
    solution, worst_index = replace_worst_element(previous, data)
    previous_worst_distance = previous["distance"].loc[worst_index]
    while solution.distance.loc[worst_index] <= previous_worst_distance:
        solution, _ = replace_worst_element(previous=solution, data=data)
    return solution


def explore_neighbourhood(element, data, max_iterations=100000):
    neighbourhood = []
    neighbourhood.append(element)
    for _ in range(max_iterations):
        previous_solution = neighbourhood[-1]
        neighbour = get_random_solution(previous=previous_solution, data=data)
        neighbourhood.append(neighbour)
    return neighbour


def genetic_algorithm(n, m, data):
    first_solution = get_first_random_solution(n=n, m=m)
    best_solution = explore_neighbourhood(
        element=first_solution, data=data, max_iterations=100
    )
    return best_solution