"""Simple Travelling Salesperson Problem (TSP) between cities."""

from ortools.constraint_solver import routing_enums_pb2
from ortools.constraint_solver import pywrapcp
import timeit

def create_data_model():
    """Stores the data for the problem."""
    data = {}
    data["distance_matrix"] =  [
[8, 8, 15, 7, 11, 17, 8, 3, 4, 11, 9, 6, 2, 7, 15, 6, 16, 5, 9, 14],
[16, 17, 13, 16, 15, 17, 12, 11, 7, 8, 13, 3, 7, 15, 9, 8, 16, 18, 7, 5],
[16, 4, 1, 9, 17, 11, 18, 2, 16, 13, 6, 5, 2, 7, 18, 16, 19, 14, 14, 9],
[9, 9, 7, 12, 16, 12, 3, 5, 16, 18, 10, 7, 6, 12, 4, 5, 14, 11, 5, 5],
[5, 2, 16, 18, 15, 15, 8, 10, 15, 18, 10, 7, 2, 1, 15, 17, 2, 15, 9, 6],
[9, 7, 11, 10, 12, 6, 10, 1, 5, 20, 19, 12, 13, 1, 19, 16, 17, 20, 16, 2],
[15, 7, 3, 15, 20, 5, 10, 4, 20, 13, 5, 11, 20, 16, 13, 4, 7, 11, 13, 17],
[11, 16, 19, 15, 12, 12, 10, 9, 1, 15, 17, 17, 18, 17, 9, 12, 9, 4, 17, 3],
[2, 5, 16, 14, 5, 3, 16, 7, 3, 2, 1, 3, 20, 16, 14, 9, 10, 18, 20, 8],
[9, 4, 20, 2, 11, 19, 5, 3, 11, 2, 13, 20, 17, 5, 12, 17, 17, 9, 5, 4],
[5, 14, 7, 14, 3, 13, 1, 13, 17, 6, 18, 18, 16, 15, 18, 19, 15, 12, 9, 4],
[14, 2, 17, 8, 11, 16, 7, 12, 12, 6, 13, 12, 6, 18, 3, 7, 3, 1, 7, 17],
[8, 10, 8, 13, 15, 20, 2, 11, 9, 1, 4, 14, 1, 19, 15, 13, 18, 16, 10, 2],
[16, 2, 14, 3, 11, 17, 7, 6, 4, 2, 3, 8, 20, 8, 7, 12, 1, 4, 20, 10],
[5, 19, 10, 18, 6, 18, 2, 11, 6, 18, 11, 18, 13, 7, 18, 15, 3, 10, 20, 7],
[16, 19, 12, 19, 15, 12, 14, 8, 7, 11, 4, 13, 19, 12, 14, 13, 3, 7, 8, 20],
[4, 4, 15, 5, 16, 14, 18, 11, 18, 20, 19, 18, 18, 4, 3, 20, 5, 18, 18, 1],
[19, 6, 2, 15, 20, 3, 13, 7, 10, 17, 12, 11, 10, 19, 6, 2, 2, 19, 15, 1],
[16, 16, 3, 7, 6, 12, 6, 3, 2, 1, 11, 8, 13, 16, 14, 5, 20, 20, 10, 3],
[3, 14, 13, 20, 20, 2, 7, 14, 2, 6, 2, 19, 20, 20, 14, 7, 7, 5, 15, 18],
]
    data["num_vehicles"] = 1
    data["depot"] = 0
    return data


def print_solution(manager, routing, solution):
    """Prints solution on console."""
    print(f"Objective: {solution.ObjectiveValue()} miles")
    index = routing.Start(0)
    plan_output = "Route for vehicle 0:\n"
    route_distance = 0
    while not routing.IsEnd(index):
        plan_output += f" {manager.IndexToNode(index) + 1} ->"
        previous_index = index
        index = solution.Value(routing.NextVar(index))
        route_distance += routing.GetArcCostForVehicle(previous_index, index, 0)
    plan_output += f" {manager.IndexToNode(index)}\n"
    plan_output += f"Route distance: {route_distance} miles\n"
    print(plan_output)


def main():
    """Entry point of the program."""
    # Instantiate the data problem.
    data = create_data_model()

    # Create the routing index manager.
    manager = pywrapcp.RoutingIndexManager(
        len(data["distance_matrix"]), data["num_vehicles"], data["depot"]
    )

    # Create Routing Model.
    routing = pywrapcp.RoutingModel(manager)


    def distance_callback(from_index, to_index):
        """Returns the distance between the two nodes."""
        # Convert from routing variable Index to distance matrix NodeIndex.
        from_node = manager.IndexToNode(from_index)
        to_node = manager.IndexToNode(to_index)
        return data["distance_matrix"][from_node][to_node]

    transit_callback_index = routing.RegisterTransitCallback(distance_callback)

    # Define cost of each arc.
    routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)

    # Setting first solution heuristic.
    # Local Optimum
    # search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    # search_parameters.first_solution_strategy = (
    #     routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
    # )
    # Global Optimum
    search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    search_parameters.local_search_metaheuristic = (
        routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH)
    search_parameters.time_limit.seconds = 30
    search_parameters.log_search = True

    # Solve the problem.
    solution = routing.SolveWithParameters(search_parameters)

    # Print solution on console.
    if solution:
        print_solution(manager, routing, solution)


if __name__ == "__main__":
    start = timeit.default_timer()
    main()
    stop = timeit.default_timer()
    print('Time: ', stop - start)