## for data
import numpy as np
import pandas as pd
## for plotting
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
## for machine learning
from sklearn import preprocessing, cluster
import scipy
## for simple routing
import osmnx as ox  #1.2.2
## for advanced routing
from ortools.constraint_solver import pywrapcp  #9.6
from ortools.constraint_solver import routing_enums_pb2

dtf = pd.read_csv(r'/Users/vittorioguglielmoglave/PycharmProjects/Thesis/Updated_temp.csv')
dtf = dtf[["New Request","Lat","Lng"]].reset_index(drop=True)
# create a new column 'id' to identify each point in the dtf
dtf = dtf.reset_index().rename(columns={"index":"id", "Lat":"Lat", "Lng":"Lng"}) # Latitude=Y axis and Longitude=X axis
dtf.head()

i = 0
dtf["base"] = dtf["id"].apply(lambda x: 1 if x==i else 0)
start = dtf[dtf["base"]==1][["Lat","Lng"]].values[0]

########################################## DIVIDE LOCATIONS IN ONLY 1 CLUSTER = NOT CONSIDERING CLUSTERING ######################################
k = 1
model = cluster.KMeans(n_clusters=k, init='k-means++')
X = dtf[dtf["base"]==0][["Lat","Lng"]]

dtf_X = X.copy()

dtf_X["cluster"] = model.fit_predict(X)
dtf["cluster"] = dtf_X["cluster"]
dtf.sample(5)


## plot
fig, ax = plt.subplots()
palette_personalizzata = ["#FF5733", "#33FF57", "#3366FF", "#FF33A1", "#33FFFF", "#FF3366", "#FFFF33", "#9933FF", "#FF9933", "#33FF99"]
sns.scatterplot(x="Lat", y="Lng", data=dtf,
                palette=sns.color_palette(palette_personalizzata,k),
                hue='cluster', legend=False, ax=ax).set_title('GLS without clustering')

ax.set_xlabel('Latitude')
ax.set_ylabel('Longitude')
ax.scatter(start[0], start[1], c='black', marker='^')
plt.show()


# create network graph
G = ox.graph_from_point(start, dist=10000,network_type="drive")
G = ox.add_edge_speeds(G)
G = ox.add_edge_travel_times(G)

############################################################ GLS WITHOUT CLUSTERING ############################################################
# get the node for each location (both depot and customers)
# create a new column 'node' in dtf with the node value for each location
dtf["node"] = dtf[["Lat","Lng"]].apply(lambda x: ox.nearest_nodes(G, x[1], x[0]), axis=1)
# used to eliminate all duplicated nodes
# keep only first nodes that appear in the dataframe
dtf = dtf.drop_duplicates("node", keep='first')
dtf.head()


# this function computes the distance shortest path between each node
def shortest_distance(a,b):
    try:
        d = nx.shortest_path_length(G, source=a, target=b, method='dijkstra', weight='travel_time')
    except:
        d = np.nan
    return d

lst_for_dist_matrix = dtf[dtf["base"]==1]["node"].tolist()
lst_for_dist_matrix += dtf[dtf["cluster"]==0]["node"].tolist()
lst_id_nodes = dtf[dtf["base"]==1]["id"].tolist()
lst_id_nodes += dtf[dtf["cluster"]==0]["id"].tolist()

my_dict = dict(zip(lst_id_nodes, lst_for_dist_matrix))


# it important to associate the shortest route computed by routing model ortools
# with the right 'id' of nodes in the dataframe dtf
size = list(range(0,len(lst_for_dist_matrix)))

dict_route_nodes = dict(zip(size, lst_id_nodes))

# dm stands for 'distance matrix' and it creates the matrix of distances
dm = np.asarray([[shortest_distance(a,b) for b in lst_for_dist_matrix] for a in lst_for_dist_matrix])
# int values are needed for the routing model
dm = (np.rint(dm)).astype(int)

# Parameters
driver = 1
# we need the equivalent node in the graph
start_node = ox.nearest_nodes(G, start[1], start[0])
print("start node:", start_node, "| total locations to visit:", len(lst_for_dist_matrix)-1, "| drivers:", driver, "\n")
driver_capacity = [500]
demands = [0] + [1]*(len(lst_for_dist_matrix)-1)
max_distance = 100000

# Create the routing index manager
manager = pywrapcp.RoutingIndexManager(len(lst_for_dist_matrix), driver, lst_for_dist_matrix.index(start_node))
# 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 dm[from_node][to_node]

transit_callback_index = routing.RegisterTransitCallback(distance_callback)

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

## The constraint about capacity
def get_demand(from_index):
    return demands[from_index]

demand = routing.RegisterUnaryTransitCallback(get_demand)

routing.AddDimensionWithVehicleCapacity(demand, slack_max=0,
                                     vehicle_capacities=driver_capacity,
                                     fix_start_cumul_to_zero=True,
                                     name='Capacity')

## The constraint about distance
name = 'Distance'
routing.AddDimension(transit_callback_index, slack_max=0, capacity=max_distance,
                   fix_start_cumul_to_zero=True, name=name)
distance_dimension = routing.GetDimensionOrDie(name)
distance_dimension.SetGlobalSpanCostCoefficient(100)

## Initial solution that minimizes costs
parameters = pywrapcp.DefaultRoutingSearchParameters()
parameters.first_solution_strategy = (routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC)
# Metaheuristic optimization of initial solution
parameters.local_search_metaheuristic = (routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH)
# The solver is configured to run for a maximum of 1 second.
# This means that the solver will attempt to find the best solution within the specified time limit
# and may terminate if the time limit is reached, even if an optimal solution hasn't been found yet.
parameters.time_limit.FromSeconds(1)
solution = routing.SolveWithParameters(parameters)


dic_routes = {}
total_distance = 0
total_load = 0

index = routing.Start(0)
route_idx = []
route_distance = 0
route_load = 0
while not routing.IsEnd(index):
    node_index = manager.IndexToNode(index)
    route_idx.append(manager.IndexToNode(index))
    previous_index = index
    index = solution.Value(routing.NextVar(index))
    route_distance += distance_callback(previous_index, index)
    route_load += demands[node_index]  ## for data
# in route_idx there is a sequence of descending numbers that build the route
# in my_route, the 'id' of the nodes of the route are extracted and
# the same route of 'route-idx' now it's shown in my_route but using the 'id' of nodes
route_idx.append(manager.IndexToNode(index))
my_route = [dict_route_nodes[x] for x in route_idx]
print(my_route)
dic_routes[0] = my_route
# route_distance has the distance in meters and so it is divided by 1000
# And then the result has 2 decimals
print(f'distance: {round(route_distance / 1000, 2)} km')
print(f'load: {round(route_load, 2)}', "\n")
total_distance += route_distance
total_load += route_load

print(f'Total distance: {round(total_distance / 1000, 2)} km')
print(f'Total load: {total_load}')

################################################# ROUTING REPRESENTATION #############################################
fig, ax = plt.subplots()
palette_personalizzata = ["#FF5733", "#33FF57", "#3366FF", "#FF33A1", "#33FFFF", "#FF3366", "#FFFF33", "#9933FF", "#FF9933", "#33FF99"]

# Scatter plot
sns.scatterplot(x="Lat", y="Lng", data=dtf,
                palette=sns.color_palette(palette_personalizzata, k),
                hue='cluster', legend=False, ax=ax).set_title('GLS without clustering')

# Add the start point (depot)
ax.scatter(start[0], start[1], c='black', marker='^')

ax.set_xlabel('Latitude')
ax.set_ylabel('Longitude')

for k,v in dic_routes.items():
    route_coordinates = dtf.loc[v, ["Lat", "Lng"]]
    ax.plot(route_coordinates["Lat"], route_coordinates["Lng"], linestyle='--', color='black')

plt.show()
