#!/usr/bin/env python3
"""Capacited Vehicles Routing Problem (CVRP)."""
from __future__ import print_function
from ortools.constraint_solver import routing_enums_pb2
from ortools.constraint_solver import pywrapcp

# DM-1: ORIGINAL DISTANCE MATRIX FROM THE EXAMPLE (WORKS)
distance_matrix = [[0, 0, 0, 0, 121, 47, 72, 25, 4],
 [0, 0, 0, 0, 121, 47, 72, 25, 4],
 [0, 0, 0, 0, 121, 47, 72, 25, 4],
 [0, 0, 0, 0, 121, 47, 72, 25, 4],
 [124, 124, 124, 124, 0, 89, 177, 137, 131],
 [50, 50, 50, 50, 92, 0, 102, 63, 56],
 [81, 81, 81, 81, 188, 114, 0, 90, 75],
 [40, 40, 40, 40, 137, 63, 92, 0, 37],
 [2, 2, 2, 2, 113, 38, 75, 28, 0]]

# TM-1: ORIGINAL TIME MATRIX FROM THE EXAMPLE (WORKS)
time_matrix = [[0, 0, 0, 0, 14, 7, 11, 8, 2],
 [0, 0, 0, 0, 14, 7, 11, 8, 2],
 [0, 0, 0, 0, 14, 7, 11, 8, 2],
 [0, 0, 0, 0, 14, 7, 11, 8, 2],
 [14, 14, 14, 14, 0, 12, 20, 19, 16],
 [7, 7, 7, 7, 13, 0, 13, 11, 9],
 [11, 11, 11, 11, 22, 15, 0, 17, 13],
 [12, 12, 12, 12, 21, 13, 17, 0, 11],
 [1, 1, 1, 1, 13, 5, 12, 9, 0]]

# DM-2: MY DISTANCE MATRIX WITH REAL DISTANCES FROM REAL ADDRESSES (SOLUTIONS: "NONE")
distance_matrix = [
    [0, 9125, 27576, 4892, 14603, 10224, 5021, 9647, 11653, 11103], 
    [8762, 0, 24718, 6507, 12812, 6467, 4229, 6289, 18614, 18064], 
    [27067, 24852, 0, 25371, 19170, 21079, 25748, 18951, 36978, 36428], 
    [5122, 6822, 25773, 0, 16483, 7921, 2113, 7344, 14974, 14424], 
    [13576, 12060, 25681, 15598, 0, 9180, 15728, 9634, 23488, 22938], 
    [9664, 6628, 18780, 7409, 9849, 0, 7787, 454, 19517, 18966], 
    [6450, 6002, 24908, 4524, 13446, 7101, 0, 6478, 17514, 16964], 
    [9618, 6844, 19158, 7363, 10068, 3925, 7740, 0, 19470, 18920], 
    [10488, 18682, 37103, 14449, 24130, 19781, 14579, 19204, 0, 1296], 
    [10250, 18445, 36865, 14211, 23892, 19543, 14341, 18967, 1476, 0]
]

# DM-3: DISABLED MATRIX, COPY FROM ORIGINAL, BUT WITH 10 DESTINATIONS (WORKS)
#distance_matrix = [
#    [0, 0, 0, 0, 121, 47, 72, 25, 4, 13],
#    [0, 0, 0, 0, 121, 47, 72, 25, 4, 9],
#    [0, 0, 0, 0, 121, 47, 72, 25, 4, 7],
#    [0, 0, 0, 0, 121, 47, 72, 25, 4, 11],
#    [124, 124, 124, 124, 0, 89, 177, 137, 131, 84],
#    [50, 50, 50, 50, 92, 0, 102, 63, 56, 47],
#    [81, 81, 81, 81, 188, 114, 0, 90, 75, 79],
#    [40, 40, 40, 40, 137, 63, 92, 0, 37, 42],
#    [2, 2, 2, 2, 113, 38, 75, 28, 0, 1],
#    [40, 40, 40, 40, 137, 63, 92, 0, 37, 28],
#]

# TM-2: MY TIME MATRIX WITH REAL TIMES FROM THE REAL ADDRESSES (WORKS WHEN USING DM-3)
time_matrix = [
    [0, 12, 20, 7, 13, 11, 7, 10, 16, 16], 
    [11, 0, 23, 10, 17, 11, 10, 11, 22, 21], 
    [20, 23, 0, 21, 19, 18, 22, 18, 30, 30], 
    [7, 11, 21, 0, 16, 10, 5, 9, 17, 17], 
    [11, 14, 20, 13, 0, 9, 14, 11, 22, 21], 
    [11, 12, 18, 10, 11, 0, 12, 2, 22, 22], 
    [11, 13, 24, 9, 18, 13, 0, 12, 22, 22], 
    [11, 12, 19, 10, 12, 7, 11, 0, 21, 21], 
    [13, 22, 30, 17, 24, 21, 17, 20, 0, 7], 
    [16, 24, 33, 19, 26, 23, 20, 22, 7, 0]
]

def create_data_model():
    """Stores the data for the problem."""
    data = {}
    data['time_matrix'] = time_matrix

    # TW-1: ORIGINAL TIME WINDOWS FROM WORKING EXAMPLE
    data['time_windows'] = [
        (0, 1000),  # depot
        (0, 1000),  # unload depot 1
        (0, 1000),  # unload depot 2
        (0, 1000),  # unload depot 3
        (0, 1000),  # 1
        (0, 1000),  # 2
        (0, 1000),  # 3
        (0, 1000),  # 4
        (0, 1000),  # 5
    ]

    # TW-2: MODIFIED TIME WINDOWS WITH 10 ADDRESSES
    data['time_windows'] = [
        (0, 1000),  # depot
        (0, 1000),  # unload depot 1
        (0, 1000),  # unload depot 2
        (0, 1000),  # unload depot 3
        (0, 1000),  # 1
        (0, 1000),  # 2
        (0, 1000),  # 3
        (0, 1000),  # 4
        (0, 1000),  # 5
        (0, 1000),  # 6
    ]

    data['distance_matrix'] = distance_matrix

    # D-1: ORIGINAL DEMANDS FROM WORKING EXAMPLE
    data['demands'] = [0,
                       -15,
                       -15,
                       -15,
                       10, #1
                       10, #2
                       10, #3
                       10, #4
                       10] #5

    # D-2: MODIFIED DEMAND WITH 10 ADDRESSES
    data['demands'] = [
        10,
        10,
        10,
        10,
        10, #1
        10, #2
        10, #3
        10, #4
        10, #5
        10, #6
    ]

    data['vehicle_capacities'] = [15, 15,15, 15]
    data['num_vehicles'] = 4

    # MODIFIED VEHICLES SETTINGS FOR TESTING
    data['vehicle_capacities'] = [500, 500]
    data['num_vehicles'] = 2

    data['depot'] = 0
    return data


def print_solution(data, manager, routing, solution):
    """Prints solution on console."""
    total_distance = 0
    time_dimension = routing.GetDimensionOrDie('Time')
    capacity_dimension = routing.GetDimensionOrDie('Capacity')
    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)
            load_var = capacity_dimension.CumulVar(index)
            route_load = solution.Value(load_var)
            plan_output += f' {node_index} Load({route_load}) -> '
            previous_index = index
            index = solution.Value(routing.NextVar(index))
            route_distance += routing.GetArcCostForVehicle(previous_index, index, vehicle_id)
        node_index = manager.IndexToNode(index)
        load_var = capacity_dimension.CumulVar(index)
        route_load = solution.Value(load_var)
        plan_output += f' {node_index} Load({route_load})\n'
        plan_output += f'Distance of the route: {route_distance}m\n'
        print(plan_output)
        total_distance += route_distance
    print(f'Total distance of all routes: {total_distance}m')


def main():
    """Solve the CVRP problem."""
    # 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)


    # 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)
        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)

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

    # 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)
    capacity = 'Capacity'
    routing.AddDimensionWithVehicleCapacity(
        demand_callback_index,
        15,  # null capacity slack
        data['vehicle_capacities'],  # vehicle maximum capacities
        True,  # start cumul to zero
        capacity)
    capacity_dimension = routing.GetDimensionOrDie(capacity)
    # Set slack to zero for each location except depot.
    for location_idx in range(len(data['demands'])):
        index = manager.NodeToIndex(location_idx)
        if location_idx < 4:
            capacity_dimension.SlackVar(index).SetRange(0, 15)
        else:
            capacity_dimension.SlackVar(index).SetRange(0, 0)

    ############# Time ##############

    # Create and register a transit callback.
    def time_callback(from_index, to_index):
        """Returns the travel time between the two nodes."""
        # Convert from routing variable Index to time matrix NodeIndex.
        from_node = manager.IndexToNode(from_index)
        to_node = manager.IndexToNode(to_index)
        return data['time_matrix'][from_node][to_node]

    # Add Time Windows constraint.
    time = 'Time'
    routing.AddDimension(
        transit_callback_index,
        30,  # allow waiting time
        10000,  # maximum time per vehicle
        False,  # Don't force start cumul to zero.
        time)
    time_dimension = routing.GetDimensionOrDie(time)
    # Add time window constraints for each location except depot.
    for location_idx, time_window in enumerate(data['time_windows']):
        if location_idx == 0:
            continue
        index = manager.NodeToIndex(location_idx)
        time_dimension.CumulVar(index).SetRange(time_window[0], time_window[1])


    # Add time window constraints for each vehicle start node.
    for vehicle_id in range(data['num_vehicles']):
        index = routing.Start(vehicle_id)
        time_dimension.CumulVar(index).SetRange(data['time_windows'][0][0],
                                                data['time_windows'][0][1])

    # Instantiate route start and end times to produce feasible times.
    for i in range(data['num_vehicles']):
        routing.AddVariableMinimizedByFinalizer(
            time_dimension.CumulVar(routing.Start(i)))
        routing.AddVariableMinimizedByFinalizer(
            time_dimension.CumulVar(routing.End(i)))

    ############# Time ##############

    # Setting first solution heuristic.
    search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    search_parameters.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC)
    search_parameters.local_search_metaheuristic = (
        routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH)
    search_parameters.time_limit.FromSeconds(1)

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

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


if __name__ == '__main__':
    main()