{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Sequence\n",
    "\n",
    "\n",
    "def find_subtours(x: Sequence[Sequence[int]]) -> Sequence[Sequence[int]]:\n",
    "    \"\"\"\n",
    "    Extracts subtours from the given 2D-array.\n",
    "\n",
    "    Args:\n",
    "        x: A two-dimensional array corresponding to the x variable in the TSP formulation, where\n",
    "            x[i][j] is 1 if arc (i, j) is used.\n",
    "\n",
    "    Returns:\n",
    "        A list of subtours, where each subtour is a list.\n",
    "    \"\"\"\n",
    "    N = len(x)\n",
    "    marked = [False] * N\n",
    "    subtours: list[list[int]] = []\n",
    "\n",
    "    while not all(marked):\n",
    "        # Index of the first non-marked city:\n",
    "        istart = min(range(N), key=lambda j: marked[j])\n",
    "\n",
    "        # Create the subtour:\n",
    "        subtour = [istart]\n",
    "\n",
    "        while True:\n",
    "            i = istart\n",
    "            for i, b in enumerate(x[subtour[-1]]):\n",
    "                if b:\n",
    "                    break\n",
    "\n",
    "            if i == istart:\n",
    "                break\n",
    "\n",
    "            subtour.append(i)\n",
    "\n",
    "        for i in subtour:\n",
    "            marked[i] = True\n",
    "\n",
    "        subtours.append(subtour)\n",
    "\n",
    "    return subtours\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "699.0\n"
     ]
    }
   ],
   "source": [
    "from docplex.mp.model import Model\n",
    "from docplex.mp.callbacks.cb_mixin import ConstraintCallbackMixin\n",
    "from cplex.callbacks import LazyConstraintCallback\n",
    "from docplex.mp.solution import SolveSolution\n",
    "import numpy as np\n",
    "from numpy.typing import NDArray\n",
    "\n",
    "import tsp.data\n",
    "\n",
    "\n",
    "class SubtourConstraintCallback(ConstraintCallbackMixin, LazyConstraintCallback):\n",
    "    x: NDArray\n",
    "\n",
    "    def __init__(self, env):\n",
    "        LazyConstraintCallback.__init__(self, env)\n",
    "        ConstraintCallbackMixin.__init__(self)\n",
    "\n",
    "        self.nb_callback = 0\n",
    "        self.nb_subtours = 0\n",
    "\n",
    "        self.single_subtour_per_call = False\n",
    "\n",
    "    def __call__(self):\n",
    "        n = len(self.x)\n",
    "\n",
    "        x_list = self.x.flatten().tolist()\n",
    "\n",
    "        sol: SolveSolution = self.make_solution_from_vars(x_list)\n",
    "\n",
    "        xvals = np.array(sol.get_values(x_list)).reshape(self.x.shape).astype(int)\n",
    "        subtours = find_subtours(xvals)\n",
    "\n",
    "        if len(subtours) == 1:\n",
    "            return\n",
    "\n",
    "        self.nb_callback += 1\n",
    "\n",
    "        for subtour in subtours:\n",
    "            xs = [\n",
    "                self.x[i, j].index\n",
    "                for i in subtour\n",
    "                for j in range(n)\n",
    "                if j not in subtour\n",
    "            ]\n",
    "            self.add([xs, [1] * len(xs)], \"G\", 1)\n",
    "\n",
    "            self.nb_subtours += 1\n",
    "\n",
    "            if self.single_subtour_per_call:\n",
    "                break\n",
    "\n",
    "\n",
    "with Model(name=\"TSP\") as m:\n",
    "    dist = tsp.data.grid42\n",
    "    n = len(dist)\n",
    "\n",
    "    x = np.array([m.binary_var_list(n, name=f\"x_{i}\") for i in range(n)])\n",
    "\n",
    "    for i in range(n):\n",
    "        m.add_constraint(x[i, :].sum() == 1)\n",
    "        m.add_constraint(x[:, i].sum() == 1)\n",
    "        m.add_constraint(x[i, i] == 0)\n",
    "\n",
    "    m.set_objective(\"min\", (x * dist).sum())\n",
    "\n",
    "    cb: SubtourConstraintCallback = m.register_callback(SubtourConstraintCallback)\n",
    "    cb.x = x\n",
    "\n",
    "    s = m.solve()\n",
    "\n",
    "    print(cb.nb_subtours)\n",
    "    print(s.get_objective_value())\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "INSA",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}