# Упрощенная версия модели без ограничений типа (2)
from ortools.sat.python import cp_model
import pandas as pd
import numpy as np

# Считаем входные данные
df_dist = df_cab = pd.read_csv("https://raw.githubusercontent.com/Lozkins/mos/master/data/06_tsp_dist.csv", sep=";", encoding="cp1251")
df_location = df_cab = pd.read_csv("https://raw.githubusercontent.com/Lozkins/mos/master/data/06_tsp_coord.csv", sep=";", encoding="cp1251")

obj_type = "cost"  # Тип целевой функции. Допустимые значения ["cost", "travel_time"]

m = cp_model.CpModel()

# Инициализация переменных
df_dist["var_id"] = np.ogrid[:df_dist.shape[0]]
df_dist["var"] = df_dist["var_id"].apply(lambda x: m.NewBoolVar(f"x_{x}"))

# Создаем ограничения типа (1)
dct_constr_1 = df_dist.groupby("to")["var"].sum().to_dict()
for from_node, var_sum in dct_constr_1.items():
  m.Add(var_sum == 2)

# Создаем ограничения типа (3)
dct_constr_3 = df_dist.set_index(["from", "to"]).to_dict()["var"]
for (from_node,  to_node), var_ in dct_constr_3.items():
  m.Add(var_ == dct_constr_3[to_node, from_node])

m.Minimize(sum(df_dist[obj_type] * df_dist["var"]))

# Инициализация solver
solver = cp_model.CpSolver()
solver.parameters.log_search_progress = True

# Решение задачи
status = solver.Solve(m)

# Проверяем статус
if status == cp_model.OPTIMAL:
  print("-" * 40)
  gap_val = abs((solver.ObjectiveValue() - solver.BestObjectiveBound()) / max(1, solver.ObjectiveValue()))
  # Извлекаем результат
  df_dist["sol"] = df_dist["var"].apply(lambda x: solver.Value(x))
  df_dist_opt = df_dist[df_dist["sol"] > 0]

  print(f"-- Found solution {solver.ObjectiveValue()}")
  print(f"-- GAP: {round(gap_val, 2)}%")
  print(f"-- Solved in {solver.WallTime()} sec.")
  print(f"-- Costs {df_dist_opt['cost'].sum() / 2} rub")
  print(f"-- Duration {df_dist_opt['travel_time'].sum() / 2} h")

# Нанесение результата на карту
import plotly.graph_objects as go

def plot_route(df_graph: pd.DataFrame):
  # Начало и окончание маршрута в МСК
  df_location["color"] = np.where(df_location["id"] == 0, "red", "blue")
  df_location["size"] = np.where(df_location["id"] == 0, 14, 10)
  df_location['text'] = df_location["id"].astype(str) + ", " + df_location["name"]

  # Добавляем координаты и наименование точек отправления и прибытия
  df_loc_from = df_location.copy()
  df_loc_from.columns = [col + "_from" for col in df_loc_from.columns]
  df_loc_to = df_location.copy()
  df_loc_to.columns = [col + "_to" for col in df_loc_to.columns]

  df_travel_map = df_graph.merge(df_loc_from, how="left", left_on="from",
                              right_on="id_from")
  df_travel_map = df_travel_map.merge(df_loc_to, how="left", left_on="to",
                              right_on="id_to")

  # Наносим вершины
  fig = go.Figure(data=go.Scattergeo(
          lon = df_location['lon'],
          lat = df_location['lat'],
          text = df_location['text'],
          mode = 'markers',
          name = "Локации",
          marker = dict(
              size = df_location["size"],
              opacity = 0.8,
              color = df_location["color"],
              ),
          ))

  # Наносим ребра
  for row in df_travel_map.itertuples():
    fig.add_trace(
            go.Scattergeo(
                lon = [row.lon_from, row.lon_to],
                lat = [row.lat_from, row.lat_to],
                mode = 'lines',
                line = dict(width = 1,color = 'blue'),
                opacity = 0.8,
                name=row.name_from + " - " + row.name_to
            ),
        )

  fig.update_layout(
          title = 'Travel route in Russia <br>',
          geo = dict(
              showland = True,
              landcolor = "rgb(250, 250, 250)",
              subunitcolor = "rgb(217, 217, 217)",
              countrycolor = "rgb(217, 217, 217)",
              countrywidth = 0.5,
              subunitwidth = 0.5
          ),
      )

  fig.show()

# Итеративное добавление ограничений типа (2)
from ortools.sat.python import cp_model
import networkx as nx
import pandas as pd
import numpy as np

# Считаем входные данные
df_dist = df_cab = pd.read_csv("https://raw.githubusercontent.com/Lozkins/mos/master/data/06_tsp_dist.csv", sep=";", encoding="cp1251")
df_location = df_cab = pd.read_csv("https://raw.githubusercontent.com/Lozkins/mos/master/data/06_tsp_coord.csv", sep=";", encoding="cp1251")

obj_type = "travel_time"  # Тип целевой функции. Допустимые значения ["cost", "travel_time"]

m = cp_model.CpModel()

# Инициализация переменных
df_dist["var_id"] = np.ogrid[:df_dist.shape[0]]
df_dist["var"] = df_dist["var_id"].apply(lambda x: m.NewBoolVar(f"x_{x}"))

# Создаем ограничения типа (1)
dct_constr_1 = df_dist.groupby("to")["var"].sum().to_dict()
for from_node, var_sum in dct_constr_1.items():
  m.Add(var_sum == 2)

# Создаем ограничения типа (3)
dct_constr_3 = df_dist.set_index(["from", "to"]).to_dict()["var"]
for (from_node,  to_node), var_ in dct_constr_3.items():
  m.Add(var_ == dct_constr_3[to_node, from_node])

m.Minimize(sum(df_dist[obj_type] * df_dist["var"]))

# Инициализация solver
solver = cp_model.CpSolver()
solver.parameters.log_search_progress = True

# Решение задачи
status = solver.Solve(m)

# Извлекаем результат
df_dist["sol"] = df_dist["var"].apply(lambda x: solver.Value(x))
df_dist_opt = df_dist[df_dist["sol"] > 0]

# ----------- Итерации по добавлению ограничений типа (2). НАЧАЛО ------------------
# Кол-во связных подграфов
route_tmp = nx.from_pandas_edgelist(df_dist_opt, "from", "to")
lst_sg = list(nx.connected_components(route_tmp))  # Список узлов в связных подграфах

# Для статистики
cnt_constrs = 0  # Кол-во добавленных ограничений типа (2)
cnt_iter = 0  # Кол-во итераций решения задачи ILP
solver_time = 0  # Суммарное время решения задач солвером

# Повторять, пока имеем более одного несвязного подграфа
while len(lst_sg) > 1:
  cnt_iter += 1
  for sg in lst_sg:
    cnt_constrs += 1

    # Отбираем необходимые переменные для суммирования
    filt_i = df_dist["from"].isin(sg)
    filt_j = df_dist["to"].isin(sg)
    filt_i_less_j = df_dist["from"] < df_dist["to"]
    filt_all = filt_i_less_j & filt_i & filt_j
    var_sum = df_dist[filt_all]["var"].sum()

    # Добавляем ограничение типа (2) в модель
    m.Add(var_sum <= len(sg) - 1)

  status = solver.Solve(m)
  solver_time += solver.WallTime()

  # Извлекаем результат
  df_dist["sol"] = df_dist["var"].apply(lambda x: solver.Value(x))
  df_dist_opt = df_dist[df_dist["sol"] > 0]

  # Кол-во связных подграфов
  route_tmp = nx.from_pandas_edgelist(df_dist_opt, "from", "to")
  lst_sg = list(nx.connected_components(route_tmp))  # Список узлов в связных подграфах
# ----------- Итерации по добавлению ограничений типа (2). КОНЕЦ ------------------

# Проверяем статус
if status == cp_model.OPTIMAL:
  print("-" * 40)
  gap_val = abs((solver.ObjectiveValue() - solver.BestObjectiveBound()) / max(1, solver.ObjectiveValue()))
  print(f"-- Found solution {solver.ObjectiveValue()}")
  print(f"-- GAP: {round(gap_val, 2)}%")
  print(f"-- Solved in {solver_time} sec.")
  print(f"-- Constraints (2) added {cnt_constrs}")
  print(f"-- Iterations count {cnt_iter}")
  print(f"-- Costs {df_dist_opt['cost'].sum() / 2} rub")
  print(f"-- Duration {df_dist_opt['travel_time'].sum() / 2} h")

# Полная модель - инициализируем все ограничения типа (2)
from ortools.sat.python import cp_model
import pandas as pd
import numpy as np
from itertools import chain, combinations

# Считаем входные данные
df_dist = df_cab = pd.read_csv("https://raw.githubusercontent.com/Lozkins/mos/master/data/06_tsp_dist.csv", sep=";", encoding="cp1251")
df_location = df_cab = pd.read_csv("https://raw.githubusercontent.com/Lozkins/mos/master/data/06_tsp_coord.csv", sep=";", encoding="cp1251")

obj_type = "travel_time"  # Тип целевой функции. Допустимые значения ["cost", "travel_time"]

m = cp_model.CpModel()

# Инициализация переменных
df_dist["var_id"] = np.ogrid[:df_dist.shape[0]]
df_dist["var"] = df_dist["var_id"].apply(lambda x: m.NewBoolVar(f"x_{x}"))

# Создаем ограничения типа (1)
dct_constr_1 = df_dist.groupby("to")["var"].sum().to_dict()
for from_node, var_sum in dct_constr_1.items():
  m.Add(var_sum == 2)

# Создаем ограничения типа (2)
def powerset(iterable):
  s = list(iterable)
  return chain.from_iterable(combinations(s, r) for r in range(len(s)+1) if 1 < r < len(s) - 1)

cnt_constrs = 0
for lst_nodes in powerset(df_location["id"]):
  cnt_constrs += 1

  # Отбираем необходимые переменные для суммирования
  filt_i = df_dist["from"].isin(lst_nodes)
  filt_j = df_dist["to"].isin(lst_nodes)
  filt_i_less_j = df_dist["from"] < df_dist["to"]
  filt_all = filt_i_less_j & filt_i & filt_j
  var_sum = df_dist[filt_all]["var"]

  if len(var_sum) != 0:
    cnt_constrs += 1
    var_sum = var_sum.sum()
    # Добавляем ограничение типа (2) в модель
    m.Add(var_sum <= len(lst_nodes) - 1)

# Создаем ограничения типа (3)
dct_constr_3 = df_dist.set_index(["from", "to"]).to_dict()["var"]
for (from_node,  to_node), var_ in dct_constr_3.items():
  m.Add(var_ == dct_constr_3[to_node, from_node])

m.Minimize(sum(df_dist[obj_type] * df_dist["var"]))

# Инициализация solver
solver = cp_model.CpSolver()
solver.parameters.log_search_progress = True

# Решение задачи
status = solver.Solve(m)

# Проверяем статус
if status == cp_model.OPTIMAL:
  print("-" * 40)
  gap_val = abs((solver.ObjectiveValue() - solver.BestObjectiveBound()) / max(1, solver.ObjectiveValue()))
  print(f"-- Found solution {solver.ObjectiveValue()}")
  print(f"-- GAP: {round(gap_val, 2)}%")
  print(f"-- Solved in {solver.WallTime()} sec.")
  print(f"-- Constraints (2) added {cnt_constrs}")
  print(f"-- Costs {df_dist_opt['cost'].sum() / 2} rub")
  print(f"-- Duration {df_dist_opt['travel_time'].sum() / 2} h")

  # Извлекаем результат
  df_dist["sol"] = df_dist["var"].apply(lambda x: solver.Value(x))
  df_dist_opt = df_dist[df_dist["sol"] > 0]

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

# Считаем входные данные
df_dist = df_cab = pd.read_csv("https://raw.githubusercontent.com/Lozkins/mos/master/data/06_tsp_dist.csv", sep=";", encoding="cp1251")
df_location = df_cab = pd.read_csv("https://raw.githubusercontent.com/Lozkins/mos/master/data/06_tsp_coord.csv", sep=";", encoding="cp1251")

obj_type = "cost"  # Тип целевой функции. Допустимые значения ["cost", "travel_time"]
large_M = 10**5
round_mult = 1

def create_data_model():
    """
      Подготовка данных к запуску
    """
    data = {}
    dct_dist = df_dist.set_index(["from", "to"]).to_dict()[obj_type]
    # Матрица расстояний
    data["distance_matrix"] = [
        [int(dct_dist.get((i, j), 0 if i == j else large_M) * round_mult)  for j in range(df_location.shape[0])]
        for i in range(df_location.shape[0])]
    # Кол-во транспортных средств - для задачи TSP - 1
    data["num_vehicles"] = 1
    # Вершина начала и завершения маршрута
    data["depot"] = 0
    return data

def print_solution(manager, routing, solution):
    """
      Вывод решения
    """
    print(f"Значение целевой функции: {solution.ObjectiveValue()} {obj_type}")
    index = routing.Start(0)
    plan_output = "Маршрут:\n"
    route_distance = 0
    while not routing.IsEnd(index):
        plan_output += f" {manager.IndexToNode(index)} ->"
        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} {obj_type}\n"

def get_routes(solution, routing, manager):
  """
    Извлечение маршрута из модели
  """
  routes = []
  for route_nbr in range(routing.vehicles()):
    index = routing.Start(route_nbr)
    route = [manager.IndexToNode(index)]
    while not routing.IsEnd(index):
      index = solution.Value(routing.NextVar(index))
      route.append(manager.IndexToNode(index))
    routes.append(route)
  return routes

def main():
    """
      Запуск алгоритма
    """
    # Инициализация данных
    data = create_data_model()

    # Создание менеджера маршрутизации
    manager = pywrapcp.RoutingIndexManager(
        df_location.shape[0], data["num_vehicles"], data["depot"])

    # Инициализация модели маршрутизации
    routing = pywrapcp.RoutingModel(manager)

    def distance_callback(from_index, to_index):
        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)

    # Определяем вес ребра (стоимость или время в пути)
    routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)

    # Инициализация начального решений
    search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    search_parameters.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
    )

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

    # Print solution on console.
    if solution:
        print_solution(manager, routing, solution)
        routes = get_routes(solution, routing, manager)
        # Display the routes.
        for i, route in enumerate(routes):
          print('Route', i, route)
          return route

route = main()

# Преобразуем в таблицу
lst_pairs = list(zip(route[:-1], route[1:]))
df_route_ortools = pd.DataFrame(lst_pairs, columns=["from", "to"])
df_route_ortools = df_dist.merge(df_route_ortools, how="right", on=["from", "to"])