"""Simple Vehicles Routing Problem."""

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


def create_data_model():
    """Stores the data for the problem."""
    data = {}
    data["distance_matrix"] = [
        # fmt: off
        [0.0, 10055.107404417138, 3721.69821382381, 6420.976150377692, 10055.107404417138, 11494.252586094712, 9509.686718064775, 7194.254050220477, 8705.952660049456, 6379.850683975058, 3886.900964162668, 1822.1360080334186, 3514.397119387665, 3610.3415490478765, 4450.3356475054825],
[10055.107404417138, 0.0, 6352.487948803106, 3865.3718588710813, 0.0, 3114.6631911491004, 10068.043552617626, 10377.997097781765, 14786.349611907537, 13486.132557208786, 6253.518796049527, 8261.902076161898, 8160.72865983126, 6491.504886964134, 5605.700303676604],
[3721.69821382381, 6352.487948803106, 0.0, 2730.044895787093, 6352.487948803106, 8078.121076188516, 8719.693526224311, 7282.548022588702, 10492.857847079005, 8541.725451568233, 360.96767380841925, 2000.8925577888222, 3238.5044012376, 774.9043974937284, 775.4208008567286],
[6420.976150377692, 3865.3718588710813, 2730.044895787093, 0.0, 3865.3718588710813, 6218.598554216611, 9614.791265876564, 8881.003492652062, 12690.637615496404, 10949.313370896472, 2534.078159796496, 4730.651505366733, 5438.512889903392, 3143.0180785142356, 2124.888787887561],
[10055.107404417138, 0.0, 6352.487948803106, 3865.3718588710813, 0.0, 3114.6631911491004, 10068.043552617626, 10377.997097781765, 14786.349611907537, 13486.132557208786, 6253.518796049527, 8261.902076161898, 8160.72865983126, 6491.504886964134, 5605.700303676604],
[11494.252586094712, 3114.6631911491004, 8078.121076188516, 6218.598554216611, 3114.6631911491004, 0.0, 8587.837026595082, 9691.228569020373, 14301.09183819817, 13476.252318321056, 8107.3462921088385, 9682.100007629035, 8789.80103415498, 7927.193029999964, 7307.725343561157],
[9509.686718064775, 10068.043552617626, 8719.693526224311, 9614.791265876564, 10068.043552617626, 8587.837026595082, 0.0, 2723.4335025409605, 6559.769557118661, 6881.23742757047, 9047.410868946314, 8527.023016095842, 6145.117859044644, 7972.602982178097, 8454.126339805342],
[7194.254050220477, 10377.997097781765, 7282.548022588702, 8881.003492652062, 10377.997097781765, 9691.228569020373, 2723.4335025409605, 0.0, 4614.784858405044, 4308.139315054496, 7639.660704478896, 6556.960595239801, 4204.749390769965, 6508.164370346345, 7246.068597913849],
[8705.952660049456, 14786.349611907537, 10492.857847079005, 12690.637615496404, 14786.349611907537, 14301.09183819817, 6559.769557118661, 4614.784858405044, 0.0, 2395.635698408829, 10844.49614996829, 9034.5034090655, 7302.614530267333, 9789.122518627675, 10733.938152600293],
[6379.850683975058, 13486.132557208786, 8541.725451568233, 10949.313370896472, 13486.132557208786, 13476.252318321056, 6881.23742757047, 4308.139315054496, 2395.635698408829, 0.0, 8878.125391048365, 6901.8917190574975, 5517.974514751043, 7897.332824366355, 8891.714263023221],
[3886.900964162668, 6253.518796049527, 360.96767380841925, 2534.078159796496, 6253.518796049527, 8107.3462921088385, 9047.410868946314, 7639.660704478896, 10844.49614996829, 8878.125391048365, 0.0, 2238.569086322884, 3598.221078392358, 1131.6846740391532, 838.9685870948458],
[1822.1360080334186, 8261.902076161898, 2000.8925577888222, 4730.651505366733, 8261.902076161898, 9682.100007629035, 8527.023016095842, 6556.960595239801, 9034.5034090655, 6901.8917190574975, 2238.569086322884, 0.0, 2397.491113642382, 1790.1374118031774, 2675.8725837926613],
[3514.397119387665, 8160.72865983126, 3238.5044012376, 5438.512889903392, 8160.72865983126, 8789.80103415498, 6145.117859044644, 4204.749390769965, 7302.614530267333, 5517.974514751043, 3598.221078392358, 2397.491113642382, 0.0, 2500.849919010973, 3431.7098964898996],
[3610.3415490478765, 6491.504886964134, 774.9043974937284, 3143.0180785142356, 6491.504886964134, 7927.193029999964, 7972.602982178097, 6508.164370346345, 9789.122518627675, 7897.332824366355, 1131.6846740391532, 1790.1374118031774, 2500.849919010973, 0.0, 1019.8137083490479],
[4450.3356475054825, 5605.700303676604, 775.4208008567286, 2124.888787887561, 5605.700303676604, 7307.725343561157, 8454.126339805342, 7246.068597913849, 10733.938152600293, 8891.714263023221, 838.9685870948458, 2675.8725837926613, 3431.7098964898996, 1019.8137083490479, 0.0],
        # fmt: on
    ]
    data["demands"] = [0, 10, 20, 5, 5, 2, 3, 3, 5, 20, 2, 4, 1, 0]
    data["vehicle_capacities"] = [20, 20, 20, 20]
    data["num_vehicles"] = 4
    data["starts"] = [0, 0, 0, 0]
    data["ends"] = [14, 14, 14, 14]

    return data

def print_solution(data, manager, routing, solution):
    """Prints solution on console."""
    print(f"Objective: {solution.ObjectiveValue()}")
    # Display routes
    total_distance = 0
    total_load = 0
    for vehicle_id in range(data["num_vehicles"]):
        index = routing.Start(vehicle_id)
        plan_output = f"Route for vehicle {vehicle_id}:\n"
        route_distance = 0
        route_load = 0
        while not routing.IsEnd(index):
            node_index = manager.IndexToNode(index)
            route_load += data["demands"][node_index]
            plan_output += f" {node_index} Load({route_load}) -> "
            previous_index = index
            index = solution.Value(routing.NextVar(index))
            route_distance += data["distance_matrix"][manager.IndexToNode(previous_index)][manager.IndexToNode(index)]
        plan_output += f" {manager.IndexToNode(index)} Load({route_load})\n"
        plan_output += f"Distance of the route: {route_distance}m\n"
        plan_output += f"Load of the route: {route_load}\n"
        print(plan_output)
        total_distance += route_distance
        total_load += route_load
    print(f"Total distance of all routes: {total_distance}m")
    print(f"Total load of all routes: {total_load}")


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["starts"], data["ends"]
    )

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

    # Create and register a transit callback.
    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)
        distance = data["distance_matrix"][from_node][to_node]

         # Debug printing
        print(f"From Index: {from_index}, To Index: {to_index}")
        print(f"From Node: {from_node}, To Node: {to_node}")
        print(f"Distance: {distance}")

        return distance

    transit_callback_index = routing.RegisterTransitCallback(distance_callback)

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

    # Add Capacity constraint.
    def demand_callback(from_index):
        """Returns the demand of the node."""
        # Convert from routing variable Index to demands NodeIndex.
        from_node = manager.IndexToNode(from_index)
        return data["demands"][from_node]

    demand_callback_index = routing.RegisterUnaryTransitCallback(demand_callback)
    routing.AddDimensionWithVehicleCapacity(
        demand_callback_index,
        0,  # null capacity slack
        data["vehicle_capacities"],  # vehicle maximum capacities
        True,  # start cumul to zero
        "Capacity",
    )

    # Convert maximum travel distance from kilometers to meters (assuming 3000 km).
    max_travel_distance_km = 2
    max_travel_distance_m = max_travel_distance_km * 1000  # Convert km to meters

    # Add Distance constraint.
    dimension_name = "Distance"
    routing.AddDimension(
        transit_callback_index,
        0,  # no slack
        max_travel_distance_m,  # vehicle maximum travel distance
        True,  # start cumul to zero
        dimension_name,
    )
    distance_dimension = routing.GetDimensionOrDie(dimension_name)
    distance_dimension.SetGlobalSpanCostCoefficient(100)



    # Setting first solution heuristic.
    search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    search_parameters.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
    )

    search_parameters.time_limit.seconds = 60
    search_parameters.log_search = True

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

    print("Solver status: ", routing.status())

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


if __name__ == "__main__":
    main()