from preprocessing import parse_file
from pandas import DataFrame
from sys import argv


def get_first_solution(n, data):
    distance_sum = DataFrame(columns=["point", "distance"])
    for element in range(n):
        element_df = data.query(f"source == {element} or destination == {element}")
        distance = element_df["distance"].sum()
        distance_sum = distance_sum.append(
            {"point": element, "distance": distance}, ignore_index=True
        )
    furthest_index = distance_sum["distance"].idxmax()
    furthest_row = distance_sum.iloc[furthest_index]
    furthest_row["distance"] = 0
    return furthest_row


def get_different_element(original, row):
    if row.source == original:
        return row.destination
    return row.source


def get_furthest_element(element, data):
    element_df = data.query(f"source == {element} or destination == {element}")
    furthest_index = element_df["distance"].idxmax()
    furthest_row = data.iloc[furthest_index]
    furthest_point = get_different_element(original=element, row=furthest_row)
    furthest_element = {"point": furthest_point, "distance": furthest_row["distance"]}
    return furthest_element, furthest_index


def greedy_algorithm(n, m, data):
    solutions = DataFrame(columns=["point", "distance"])
    first_solution = get_first_solution(n, data)
    solutions = solutions.append(first_solution, ignore_index=True)
    for _ in range(m):
        last_solution = solutions["point"].tail(n=1)
        centroid, index = get_furthest_element(element=int(last_solution), data=data)
        solutions = solutions.append(dict(centroid), ignore_index=True)
        data = data.drop(index)
    return solutions


# NOTE In each step, switch to the element that gives the least amount
def local_search():
    pass


def show_results(solutions):
    distance_sum = solutions["distance"].sum()
    print(solutions)
    print("Total distance: " + str(distance_sum))


def usage(argv):
    print(f"Usage: python {argv[0]} <file>")
    exit(1)


def main():
    if len(argv) != 2:
        usage(argv)
    n, m, data = parse_file(argv[1])
    solutions = greedy_algorithm(n, m, data)
    show_results(solutions)


if __name__ == "__main__":
    main()