Problem 3 - MAPP-Network-Flow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import time\n",
    "import matplotlib\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "import math\n",
    "import numpy as np\n",
    "from shapely.geometry.polygon import Polygon\n",
    "from shapely.geometry import LineString\n",
    "from shapely.geometry import Point\n",
    "from descartes import PolygonPatch\n",
    "import matplotlib.animation as animation\n",
    "import io\n",
    "import base64\n",
    "from IPython.display import HTML\n",
    "from test import *\n",
    "import yaml\n",
    "from src.environments import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tutorial for Network Flow Based Muti-agent Path Planning Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we begin, we encourage you to read the tutorial of the network flow based multi-agent path planning algorithm, which is located in ```/lab2-trajectory-planning/docs```. If you want to learn more about the algorithm in detail, with more approximations and theoretical analysis, please take a look at the following paper.\n",
    "\n",
    "Yu, Jingjin, and Steven M. LaValle. \"Multi-agent path planning and network flow.\" Algorithmic foundations of robotics X. Springer, Berlin, Heidelberg, 2013. 157-173."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to Networkx Package\n",
    "\n",
    "We'll be using the [networkx](https://networkx.github.io) Python package to represent time-expanded graph. We recommend you skim the documentation if you're not familiar with this package.\n",
    "\n",
    "All you should know about the networkx package for this implementation are 1) adding edge (with some properties), 2) adding node (with some properties) and 3) accessing the node or edge to get value for some property.\n",
    "\n",
    "First, we can initialize the directed graph as follow (note that time-expanded graph is directed graph) :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = nx.DiGraph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you can add edge with property using 'add_edge' function. Let add edge from node \"A\" to \"B\", with properties as weight=1 and capacity=1. Note that, non-existing nodes are automatically generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G.add_edge('A', 'B', weight=1, capacity=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, you can add node with property using 'add_node' function. Let add node \"C\", with properties as pos=(0,0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G.add_node('C', pos=(0,0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, you can get the value of any property of edge or node as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To get property of an edge:\n",
    "Weight_from_A_to_B = G['A']['B']['weight']\n",
    "print(Weight_from_A_to_B)\n",
    "\n",
    "# To get property of an node:\n",
    "Pos_of_C = G.node['C']['pos']\n",
    "print(Pos_of_C)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you are ready! You just need to add edges with proper weight and capacity and solve the problem with the provided solver. Let's start to write the functions for constructing time expanded graph."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function for Adding Gadget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will provide the function to expand one edge of CUG into a gadget as illustrated in the below image. Please have a look how we assign weight (refered to as \"cost\" in the tutorial) and capacity to each directed edge in the gadget, in accordance with explanation in the tutorial. We have to use \"weight\" instead of \"cost\" as the name of edge property, because it is a requirement of a built-in function we will use in the later section. For example, CUG.add_edge(A, B, weight=0, capacity=1).\n",
    "\n",
    "Also, it is important to follow the naming convention as used in the below image (also for the later section in this material). It will help the test function and visualization tools work correctly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![title](img/expand_gadget.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def expand_gadget(Time_Expanded_G, CUG, Goal_Nodes, t):\n",
    "    \"\"\" Expand each edge of a CUG into a gadget\n",
    "    \n",
    "    This function should iterate over all edges of the CUG, \n",
    "    and expand one link of the CUG into a gadget of 5 links which will be stored \n",
    "    in the time-expanded graph.\n",
    "    Each node in the expand gadget must follow the name convention in the above diagram.\n",
    "    Moreover, there are weight and capacity attached to each link of the gadget.\n",
    "    \n",
    "    Input:\n",
    "        Time_Expanded_G - Time-expanded graph\n",
    "        CUG - Collision Free Unit Distance Graph which is the originial encoding of the problem\n",
    "        Goal_Nodes - Set of goal nodes (this parameter is not used in this function)\n",
    "        t - Time step that the CUG edge will be expaned from t' to t+1\n",
    "        \n",
    "    Return value: (none) but the input parameter Time_Expanded_G will be updated\n",
    "    \"\"\"\n",
    "\n",
    "    CUG_Edges = list(CUG.edges.data())\n",
    "\n",
    "    for e in CUG_Edges:\n",
    "        u = str(e[0])\n",
    "        v = str(e[1])\n",
    "        u_v = str(e[0]) + \"_\" + str(e[1])\n",
    "        \n",
    "        u_t_p = str(u) + \"_\" + str(t) + \"_p\"\n",
    "        v_t_p = str(v) + \"_\" + str(t) + \"_p\"\n",
    "        u_v_t_p = str(u_v) + \"_\" + str(t) + \"_p\"\n",
    "        \n",
    "        u_tplus1 = str(u) + \"_\" + str(t+1)\n",
    "        v_tplus1 = str(v) + \"_\" + str(t+1)\n",
    "        u_v_tplus1 = str(u_v) + \"_\" + str(t+1)        \n",
    "        \n",
    "        Time_Expanded_G.add_edge(u_t_p, u_v_t_p, weight=0, capacity=1)\n",
    "        Time_Expanded_G.add_edge(v_t_p, u_v_t_p, weight=0, capacity=1)\n",
    "        Time_Expanded_G.add_edge(u_v_t_p, u_v_tplus1, weight=1, capacity=1)\n",
    "        Time_Expanded_G.add_edge(u_v_tplus1, u_tplus1, weight=0, capacity=1)\n",
    "        Time_Expanded_G.add_edge(u_v_tplus1, v_tplus1, weight=0, capacity=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 1: Function for Adding Parallel Edges\n",
    "\n",
    "Please write a function that add parallel edges as shonw in the following figure. From the below example, the added edges are in green and blue color. The green color represents the edges between the t' and t+1 time step while the blue color represent the links between the t+1 and (t+1)' time step. You should assign weight and capacity to each edge. Recall that weigth(cost) for the statianary action from goal node should be 0, because it alreay achieved its goal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![title](img/expand_stationary.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "a0dc86c7cd970af296b78d6c822e193e",
     "grade": false,
     "grade_id": "cell-63373a3827018989",
     "locked": false,
     "schema_version": 1,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def expand_parallel(Time_Expanded_G, CUG, Goal_Nodes, t):\n",
    "    \"\"\" Add the links that represent each agent staying at the same node \n",
    "    between adjacent time step to the time expanded graph\n",
    "    \n",
    "    This function should iterate over all nodes of the CUG and add a link \n",
    "    with appropriate weight and capacity.\n",
    "    Each node in the expand gadget must follow the name convention in the above diagram.\n",
    "    \n",
    "    Input:\n",
    "        Time_Expanded_G - Time-expanded graph\n",
    "        CUG - Collision Free Unit Distance Graph which is the originial encoding of the problem\n",
    "        Goal_Nodes - Set of goal nodes\n",
    "        t - Time step that the CUG edge will be expaned from t' to t+1 and t+1 to (t+1)' \n",
    "        \n",
    "    Return value: (none) but the input parameter Time_Expanded_G will be updated\n",
    "    \"\"\"\n",
    "    ### BEGIN SOLUTION\n",
    "    CUG_Nodes = list(CUG.nodes())\n",
    "\n",
    "    for n in CUG_Nodes:\n",
    "        n_t_p = str(n) + \"_\" + str(t) + \"_p\"\n",
    "        n_tplus1 = str(n) + \"_\" + str(t+1)\n",
    "        \n",
    "        if n in Goal_Nodes:\n",
    "            Time_Expanded_G.add_edge(n_t_p, n_tplus1, weight=0, capacity=1)\n",
    "        else:\n",
    "            Time_Expanded_G.add_edge(n_t_p, n_tplus1, weight=1, capacity=1)\n",
    "            \n",
    "    for n in CUG_Nodes:\n",
    "        n_tplus1 = str(n) + \"_\" + str(t+1)\n",
    "        n_tplus1_p = n_tplus1 + \"_p\"    \n",
    "        \n",
    "        Time_Expanded_G.add_edge(n_tplus1, n_tplus1_p, weight=0, capacity=1)  \n",
    "    ### END SOLUTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "c75297f7aa3101978e71417978ecd427",
     "grade": true,
     "grade_id": "cell-2716cd2827a1b231",
     "locked": true,
     "points": 0,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "test_expand_parallel(expand_parallel)\n",
    "test_ok()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function for Setting Initial and Goal Nodes\n",
    "\n",
    "We will provide the function to specify which nodes are the initial and goal nodes in the time expanded graph by setting its demand attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_terminals(Time_Expanded_G, Initial_Nodes, Goal_Nodes,T):\n",
    "    \"\"\" Set the node attribute \"demand\" of the initial node and goal node\n",
    "    \n",
    "    This function should iterate over all the initial and goal nodes.\n",
    "    Then set the node attribute \"demand\" to -1 for the initial node\n",
    "    and set that of the goal node to 1\n",
    "    \n",
    "    Input:\n",
    "        Time_Expanded_G - Time-expanded graph\n",
    "        Initial_Nodes - Set of initial nodes\n",
    "        Goal_Nodes - Set of goal nodes\n",
    "        T - Maximum number of time steps \n",
    "        \n",
    "    Return value: (none) but the input parameter Time_Expanded_G will be updated\n",
    "    \"\"\"\n",
    "    \n",
    "    for n in Initial_Nodes:\n",
    "        n_0_p = str(n) + \"_0_p\"\n",
    "        Time_Expanded_G.add_node(n_0_p, demand=-1)\n",
    "        \n",
    "    for n in Goal_Nodes:\n",
    "        n_Tminus1 = str(n) + \"_\" + str(T)\n",
    "        Time_Expanded_G.add_node(n_Tminus1, demand=1)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 2: Function for Expanding CUG (Assembling everything)\n",
    "\n",
    "Now let's put everything all together. Then you will get a function to expand the whole CUG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "4dd60c79100d343f02f156ca802e7282",
     "grade": false,
     "grade_id": "cell-e0a5bc016796dfbe",
     "locked": false,
     "schema_version": 1,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def expand_CUG(CUG, T, Initial_Nodes, Goal_Nodes):\n",
    "    \"\"\" Expand the CUG into a time-expanded graph\n",
    "    \n",
    "    This function should iterate over all the time step t from 0 to T-1.\n",
    "    At each time step, expand the gadget of CUG links and the stationary nodes of CUG\n",
    "    \n",
    "    Input:\n",
    "        CUG - Collision Free Unit Distance Graph which is the originial encoding of the problem\n",
    "        T - Maximum number of time steps \n",
    "        Initial_Nodes - Set of initial nodes    \n",
    "        Goal_Nodes - Set of goal nodes\n",
    "        \n",
    "        \n",
    "    Return value: Time_Expanded_G - Time-expanded graph which is created and updated inside this function\n",
    "    \"\"\"\n",
    "\n",
    "    ### BEGIN SOLUTION\n",
    "    Time_Expanded_G = nx.DiGraph()\n",
    "\n",
    "    for t in range(T):\n",
    "        expand_gadget(Time_Expanded_G, CUG, Goal_Nodes, t)\n",
    "        expand_parallel(Time_Expanded_G, CUG, Goal_Nodes, t)\n",
    "\n",
    "    set_terminals(Time_Expanded_G, Initial_Nodes, Goal_Nodes,T)\n",
    "\n",
    "    return Time_Expanded_G\n",
    "    ### END SOLUTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "bd55dde5c9ff5e2d5c503890ec20f0b5",
     "grade": true,
     "grade_id": "cell-65d9e0bfe2d89b01",
     "locked": true,
     "points": 0,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "test_expand_CUG(expand_CUG)\n",
    "test_ok()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use the expanded CUG for solving multi-agent path planning problem. All the remaining helper and simulation functions are provided for you to see how the agents will move without colliding with each other. \n",
    "\n",
    "We also provide the make_grid_CUG function for you to create a CUG with specified numbers of rows and columns of the grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_grid_CUG(n,m):\n",
    "\n",
    "    CUG = nx.DiGraph()\n",
    "\n",
    "    k = 0\n",
    "    for i in range(m):\n",
    "        for j in range(n):\n",
    "            k=k+1\n",
    "            CUG.add_node(k, pos=(j,i), cat=3)\n",
    "\n",
    "            if j<n-1 and i<m-1:\n",
    "                CUG.add_edge(k,k+1)\n",
    "                CUG.add_edge(k,k+n)\n",
    "            elif j==n-1 and i<m-1:\n",
    "                CUG.add_edge(k,k+n)\n",
    "            elif j<n-1 and i==m-1:\n",
    "                CUG.add_edge(k,k+1)\n",
    "\n",
    "    return CUG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Feel free to change the value of n and m to create different layouts of the CUG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make grid CUG\n",
    "n=3\n",
    "m=3\n",
    "CUG = make_grid_CUG(n,m)\n",
    "\n",
    "#Visualize CUG grid\n",
    "fig = plt.figure(1, figsize=(30,30), dpi=90)\n",
    "ax = fig.add_subplot(111)\n",
    "pos=nx.get_node_attributes(CUG,'pos')\n",
    "nx.draw(CUG,pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The below code will randomly select the initial and goal positions for your agents. After specifying the maximum number of time step T, we can expand the CUG and visualize it. Feel free to change num_agent to any number of agent you want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# randomly select the initial and goal nodes\n",
    "num_agent = 3\n",
    "Initial_Nodes = []\n",
    "Goal_Nodes = []\n",
    "while(len(Initial_Nodes)<num_agent):\n",
    "    rand_num = np.random.randint(n*m) +1\n",
    "    if rand_num not in Initial_Nodes:\n",
    "        Initial_Nodes.append(rand_num)\n",
    "\n",
    "while(len(Goal_Nodes)<num_agent):\n",
    "    rand_num = np.random.randint(n*m) +1\n",
    "    if rand_num not in Goal_Nodes and rand_num not in Initial_Nodes:\n",
    "        Goal_Nodes.append(rand_num)\n",
    "              \n",
    "print(\"\\nInitial Nodes: \")\n",
    "print(Initial_Nodes)\n",
    "print(\"\\nGoal Nodes: \")\n",
    "print(Goal_Nodes)\n",
    "\n",
    "#Specify the maximum number of time steps\n",
    "T = m+n\n",
    "\n",
    "# convert the CUG to the time-expanded graph\n",
    "Time_Expanded_G = expand_CUG(CUG, T, Initial_Nodes, Goal_Nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The draw_expaned_CUG is provided for visualizing the expanded CUG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_expanded_CUG(Time_Expanded_G, CUG, m, n):\n",
    "    Time_Expanded_G_Draw = Time_Expanded_G\n",
    "    Expanded_Edges = list(Time_Expanded_G_Draw.edges.data())\n",
    "    Expanded_Nodes = list(Time_Expanded_G_Draw.nodes.data())\n",
    "    CUG_Nodes = list(CUG.nodes())\n",
    "\n",
    "    for nd in Expanded_Nodes:\n",
    "        node_id = nd[0]           \n",
    "        node_name_components = node_id.split(\"_\")\n",
    "        if len(node_name_components) == 2:\n",
    "            t = int(node_name_components[1])\n",
    "            pos = CUG.node[int(node_name_components[0])]['pos']\n",
    "            pos = ( (n-1)*2 + 4*(t-1)*(n-1) + pos[0], pos[1] + 0.5 + (t-1) )\n",
    "            Time_Expanded_G_Draw.add_node(node_id, pos=pos, cat=1)        \n",
    "\n",
    "    for nd in Expanded_Nodes:\n",
    "        node_id = nd[0]           \n",
    "        node_name_components = node_id.split(\"_\")\n",
    "        if len(node_name_components)==3:\n",
    "            t = int(node_name_components[1])\n",
    "            pos = CUG.node[int(node_name_components[0])]['pos']\n",
    "            pos = ( 4*t*(n-1) + pos[0], pos[1] + t*1)            \n",
    "            Time_Expanded_G_Draw.add_node(node_id, pos=pos, cat=2)\n",
    "\n",
    "    for nd in Expanded_Nodes:\n",
    "        node_id = nd[0] \n",
    "        node_name_components = node_id.split(\"_\")\n",
    "        if len(node_name_components)==4:\n",
    "            t = int(node_name_components[2])\n",
    "                       \n",
    "            pos11 = Time_Expanded_G_Draw.node[(node_name_components[0] + \"_\" + str(t) + \"_p\")]['pos']\n",
    "            pos12 = Time_Expanded_G_Draw.node[(node_name_components[1] + \"_\" + str(t) + \"_p\")]['pos']\n",
    "            pos21 = Time_Expanded_G_Draw.node[(node_name_components[0] + \"_\" + str(t+1))]['pos']\n",
    "            pos22 = Time_Expanded_G_Draw.node[(node_name_components[1] + \"_\" + str(t+1))]['pos']\n",
    "            \n",
    "            pos1 = ( 1/2*(pos11[0] + pos12[0]), 1/2*(pos11[1] + pos12[1]) )\n",
    "            pos2 = ( 1/2*(pos21[0] + pos22[0]), 1/2*(pos21[1] + pos22[1]) )\n",
    "\n",
    "            if len(node_name_components)==4:\n",
    "                pos = ( pos1[0] + 1/3*(pos2[0] - pos1[0]), pos1[1] + 1/3*(pos2[1] - pos1[1]) )\n",
    "\n",
    "            elif (len(node_name_components)==3 and node_name_components[2] != 'p'):     \n",
    "                pos = ( pos1[0] + 2/3*(pos2[0] - pos1[0]), pos1[1] + 2/3*(pos2[1] - pos1[1]) )\n",
    "\n",
    "            Time_Expanded_G_Draw.add_node(node_id, pos=pos, cat=3)\n",
    "            \n",
    "        if len(node_name_components)==3 and node_name_components[2] != 'p':\n",
    "            t = int(node_name_components[2])\n",
    "                       \n",
    "            pos11 = Time_Expanded_G_Draw.node[(node_name_components[0] + \"_\" + str(t-1) + \"_p\")]['pos']\n",
    "            pos12 = Time_Expanded_G_Draw.node[(node_name_components[1] + \"_\" + str(t-1) + \"_p\")]['pos']\n",
    "            pos21 = Time_Expanded_G_Draw.node[(node_name_components[0] + \"_\" + str(t))]['pos']\n",
    "            pos22 = Time_Expanded_G_Draw.node[(node_name_components[1] + \"_\" + str(t))]['pos']\n",
    "            \n",
    "            pos1 = ( 1/2*(pos11[0] + pos12[0]), 1/2*(pos11[1] + pos12[1]) )\n",
    "            pos2 = ( 1/2*(pos21[0] + pos22[0]), 1/2*(pos21[1] + pos22[1]) )\n",
    "\n",
    "            if len(node_name_components)==4:\n",
    "                pos = ( pos1[0] + 1/3*(pos2[0] - pos1[0]), pos1[1] + 1/3*(pos2[1] - pos1[1]) )\n",
    "\n",
    "            elif (len(node_name_components)==3 and node_name_components[2] != 'p'):     \n",
    "                pos = ( pos1[0] + 2/3*(pos2[0] - pos1[0]), pos1[1] + 2/3*(pos2[1] - pos1[1]) )\n",
    "\n",
    "            Time_Expanded_G_Draw.add_node(node_id, pos=pos, cat=3)        \n",
    "\n",
    "    return Time_Expanded_G_Draw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can visualize the time expanded CUG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Visualize the time expanded CUG\n",
    "Time_Expanded_G_Draw = draw_expanded_CUG(Time_Expanded_G, CUG, m, n)\n",
    "fig2 = plt.figure(2, figsize=(30,30), dpi=90)\n",
    "ax2 = fig2.add_subplot(111)\n",
    "pos=nx.get_node_attributes(Time_Expanded_G_Draw,'pos')\n",
    "color_map = {1:'b', 2:'#FF0099', 3:'#660066'}\n",
    "nx.draw(Time_Expanded_G_Draw,pos,node_color=[color_map[Time_Expanded_G_Draw.node[node]['cat']] \n",
    "                                             for node in Time_Expanded_G_Draw])\n",
    "plt.show(block=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then use the built-in network_simplex function to find a solution of this network flow problem. One thing to be noted is that the goal position is not fixed for each agent at the beginning. This means that any agent can go to any goal node but they must go to a different one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solving \"permutation invariant\" multi agent path planning problem using \n",
    "# built-in function of networkx package\n",
    "flowCost, flowDict = nx.network_simplex(Time_Expanded_G)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The provided extract_path function will give you the path of each agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_path(flowDict,Initial_Nodes,CUG,return_type=\"node\"):\n",
    "\n",
    "    Path = []\n",
    "    \n",
    "    for n in Initial_Nodes:\n",
    "        Path_temp = [n]\n",
    "        node_name = str(n) + \"_0_p\";\n",
    "        Connected_Nodes = flowDict[node_name]     \n",
    "\n",
    "        while(True):\n",
    "            flag = 0\n",
    "\n",
    "            for Node in Connected_Nodes:\n",
    "                if Connected_Nodes[Node]==1:\n",
    "                    node_name_components = Node.split(\"_\")\n",
    "                    if len(node_name_components) == 2:\n",
    "                        Path_temp.append(int(node_name_components[0]))\n",
    "                    Connected_Nodes = flowDict[Node]\n",
    "                    flag = 1\n",
    "                    break\n",
    "\n",
    "            if flag == 0:\n",
    "                break\n",
    "\n",
    "        Path.append(Path_temp)\n",
    "        \n",
    "    \n",
    "    for n in range(len(Path[0])-1):\n",
    "        flag = 0\n",
    "        for p in Path:\n",
    "            if p[n]!=p[n+1]:\n",
    "                flag = 1\n",
    "                break\n",
    "        if flag==0:\n",
    "            break\n",
    "\n",
    "    Path_truncated = []\n",
    "    for p in Path:\n",
    "        Path_truncated.append(p[0:n+1])\n",
    "        \n",
    "    Path = Path_truncated\n",
    "            \n",
    "            \n",
    "    if return_type==\"waypoint\":\n",
    "        \n",
    "        nodes = CUG.nodes()\n",
    "\n",
    "        Path_dict = dict()\n",
    "        for i in range(len(Path)):\n",
    "            for j in range(len(Path[i])):\n",
    "                if j==0:\n",
    "                    Path_dict['agent'+str(i)+'_x'] = [nodes.get(Path[i][j])['pos'][0]]\n",
    "                    Path_dict['agent'+str(i)+'_y'] = [nodes.get(Path[i][j])['pos'][1]]\n",
    "                else:\n",
    "                    Path_dict['agent'+str(i)+'_x'].append(nodes.get(Path[i][j])['pos'][0])\n",
    "                    Path_dict['agent'+str(i)+'_y'].append(nodes.get(Path[i][j])['pos'][1])\n",
    "                    \n",
    "        return Path_dict\n",
    "\n",
    "\n",
    "    return Path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then extract the path into a series of nodes at each time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# retrieve path\n",
    "Path = extract_path(flowDict, Initial_Nodes, CUG, return_type=\"node\")\n",
    "\n",
    "print (\"\\nextract path:\")\n",
    "print(Path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's see how your agents will move by using a simulation function below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate(fig,ax, CUG, Path):\n",
    "    N=20\n",
    "    nodes = CUG.nodes()\n",
    "\n",
    "    x_set = []\n",
    "    y_set = []\n",
    "    for i in range(len(Path)):\n",
    "        x_temp = []\n",
    "        y_temp = []\n",
    "        for j in range(len(Path[i])-1):\n",
    "            x1 = nodes.get(Path[i][j])['pos'][0]\n",
    "            y1 = nodes.get(Path[i][j])['pos'][1]\n",
    "            x2 = nodes.get(Path[i][j+1])['pos'][0]\n",
    "            y2 = nodes.get(Path[i][j+1])['pos'][1]\n",
    "\n",
    "            x_inc = (x2-x1)/N\n",
    "            y_inc = (y2-y1)/N\n",
    "\n",
    "            for n in range(N):\n",
    "                x_temp.append(x1 + x_inc * (n+1))\n",
    "                y_temp.append(y1 + y_inc * (n+1))\n",
    "\n",
    "        x_set.append(x_temp)\n",
    "        y_set.append(y_temp)\n",
    "\n",
    "\n",
    "    pos=nx.get_node_attributes(CUG,'pos')\n",
    "    xmin = 10**10\n",
    "    xmax = -10**10\n",
    "    ymin = 10**10\n",
    "    ymax = -10**10\n",
    "    for p in pos:\n",
    "        if pos[p][0] < xmin:\n",
    "            xmin = pos[p][0]\n",
    "        if pos[p][0] > xmax:\n",
    "            xmax = pos[p][0]\n",
    "        if pos[p][1] < ymin:\n",
    "            ymin = pos[p][1]\n",
    "        if pos[p][1] > ymax:\n",
    "            ymax = pos[p][1]\n",
    "\n",
    "    xrange = [xmin-1, xmax+1]\n",
    "    yrange = [ymin-1, ymax+1]\n",
    "\n",
    "    ax.set_xlim(*xrange)\n",
    "    ax.set_ylim(*yrange)\n",
    "\n",
    "    scat = []        \n",
    "    scat = ax.scatter(x_set[0][0],y_set[0][0],color='blue',s=1000)\n",
    "    \n",
    "    def animate(i):        \n",
    "        offsets = []\n",
    "        for z in range(len(Path)):\n",
    "            offsets.append([x_set[z][i], y_set[z][i]])\n",
    "                          \n",
    "        scat.set_offsets(tuple(offsets))\n",
    "        return scat\n",
    "\n",
    "    ani = animation.FuncAnimation(fig, animate, interval=100, frames = len(x_set[0]))\n",
    "    \n",
    "    return ani"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize\n",
    "fig = plt.figure(1, figsize=(10,10), dpi=90)\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "BLUE = '#6699cc'\n",
    "for p in Path:\n",
    "    line = []\n",
    "    for n in p:\n",
    "        line.append(CUG.node[int(n)]['pos'])\n",
    "\n",
    "    linestring = LineString(line)\n",
    "    dilated = linestring.buffer(0.3)\n",
    "    patch1 = PolygonPatch(dilated, fc=BLUE, ec=BLUE, alpha=0.5, zorder=2)\n",
    "    ax.add_patch(patch1)\n",
    "\n",
    "pos=nx.get_node_attributes(CUG,'pos')\n",
    "color_map = {1:'b', 2:'#FF0099', 3:'#660066'}\n",
    "nx.draw(CUG,pos,node_color=[color_map[CUG.node[node]['cat']] for node in CUG])\n",
    "\n",
    "ani = simulate(fig,ax,CUG,Path)\n",
    "HTML(ani.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Optional) Application to Real Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's work with the real environment. For this, you should create your own CUG from the given environment file. Then use all of the pre-developed functions to solve the multi-agent path planning problem for the environmnet. Finally, you can extract each agent's waypoint using \"extract_path\" function by setting \"return_type\" as \"waypoint\". Let's do this step by step for the following cells.\n",
    "\n",
    "To load and utilize the environment yaml file, you will use ```Environment``` class which was provided to you. We encourage you to skim the ```environments.py``` in the ```/lab2-trajetory-planning/src``` folder. The following cell loads the environment file and visualize it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize environment class object\n",
    "env = Environment()\n",
    "\n",
    "# load example maze environment\n",
    "env.load_from_yaml_file(\"src/maze.yaml\")\n",
    "\n",
    "# draw environment\n",
    "ax = env.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First of all, you have to define the function which makes CUG from environment file. The following function is incomplete function for making CUG from environment yaml file. The missing parts are removing edges and nodes which are not available due to obstacles. (It will be good to see what is the problem and what you should do, by simply using this incomplete function for the further cells.) Please fill the missing parts to make the valid CUG. To do this, you may have to access ```obstacles``` property of the ```environment``` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "cell-e8288834b499137c",
     "locked": false,
     "points": 0,
     "schema_version": 1,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def make_grid_CUG_with_Env(env,edge_length,margin=(0,0,0,0)):\n",
    "    \"\"\" Make a grid CUG map for the given environment file.\n",
    "    The edges of the CUG should not traverse any obstacle in the environment. \n",
    "    Also, you may want to remove nodes that are not necessary for the path planning.\n",
    "    \n",
    "    Input:\n",
    "        env - environment yaml file\n",
    "        edge_length - the length of the edge\n",
    "        margin (optional) - bound margins (xmin_margin, ymin_margin, xmax_margin, ymax_margin)\n",
    "        \n",
    "    Return value: CUG, m, n\n",
    "    \"\"\"\n",
    "    \n",
    "\n",
    "    CUG = nx.DiGraph()\n",
    "    \n",
    "    # get bounds of the environment\n",
    "    bounds = env.bounds\n",
    "    \n",
    "    x_min = bounds[0]+margin[0]\n",
    "    x_max = bounds[2]+margin[2]\n",
    "    y_min = bounds[1]+margin[1]\n",
    "    y_max = bounds[3]+margin[3]\n",
    "    \n",
    "    k = 0\n",
    "    \n",
    "    # find the number of grids\n",
    "    m = math.floor((x_max - x_min) / edge_length)\n",
    "    n = math.floor((y_max - y_min) / edge_length)\n",
    "    \n",
    "    for i in range(m):\n",
    "        for j in range(n):\n",
    "            k=k+1\n",
    "            CUG.add_node(k, pos=(x_min+i*edge_length,y_min+j*edge_length), cat=3)\n",
    "\n",
    "            if j<n-1 and i<m-1:\n",
    "                CUG.add_edge(k,k+1)\n",
    "                CUG.add_edge(k,k+n)\n",
    "            elif j==n-1 and i<m-1:\n",
    "                CUG.add_edge(k,k+n)\n",
    "            elif j<n-1 and i==m-1:\n",
    "                CUG.add_edge(k,k+1)\n",
    "    \n",
    "    Pruned_CUG = CUG.copy()  # deep copy of the graph\n",
    "    nodes = CUG.nodes()\n",
    "    edges = CUG.edges()\n",
    "                    \n",
    "    # prune edges which traverse any obstacle\n",
    "    \n",
    "    ### BEGIN SOLUTION\n",
    "    for e in edges:\n",
    "        path = LineString([nodes.get(e[0])['pos'], nodes.get(e[1])['pos']])\n",
    "        for polygon in env.obstacles:\n",
    "            if path.intersects(polygon):\n",
    "                Pruned_CUG.remove_edge(e[0],e[1])\n",
    "                break\n",
    "                \n",
    "    ### END SOLUTION\n",
    "                \n",
    "        \n",
    "        \n",
    "    # prune nodes with degree 0\n",
    "    \n",
    "    ### BEGIN SOLUTION\n",
    "    CUG = Pruned_CUG.copy()\n",
    "    nodes = CUG.nodes()\n",
    "    edges = CUG.edges()\n",
    "    for n in nodes:\n",
    "        if CUG.degree(n)==0:\n",
    "            Pruned_CUG.remove_node(n)\n",
    "            \n",
    "    ### END SOLUTION\n",
    "    \n",
    "\n",
    "    return Pruned_CUG, m, n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you can make CUG which fits to the environment. It is also useful to define goal and start region, to define the intial and goal nodes. We gave example regions, but feel free to change them to have different initial and goal locations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# draw environment\n",
    "ax = env.draw()\n",
    "\n",
    "# make grid CUG which fit to the environment\n",
    "CUG, m, n = make_grid_CUG_with_Env(env,4,margin=(4,4,0,0))\n",
    "\n",
    "# add goal region\n",
    "goal_region = Polygon([(-25,-15), (-15,-15), (-15,-5), (-25,-5)])\n",
    "\n",
    "# add start region\n",
    "start_region = Polygon([(15,5), (25,5), (25,15), (15,15)])\n",
    "\n",
    "# Visualize CUG grid\n",
    "patch1 = PolygonPatch(goal_region, facecolor=[0,0,0.5], edgecolor=[0,0,0], alpha=0.7, zorder=2)\n",
    "patch2 = PolygonPatch(start_region, facecolor=[0,0,0.5], edgecolor=[0,0,0], alpha=0.7, zorder=2)\n",
    "ax.add_patch(patch1)\n",
    "ax.add_patch(patch2)\n",
    "\n",
    "pos=nx.get_node_attributes(CUG,'pos')\n",
    "nx.draw(CUG,pos,node_size=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, you should select initial and goal nodes. If you defined initial and goal regions, you can use them to select initial and goal locations. Otherwise, feel free to select your own initial and goal locations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose initial and goal nodes\n",
    "# randomly select the initial and goal nodes\n",
    "num_agent = 4\n",
    "Initial_Nodes = []\n",
    "Goal_Nodes = []\n",
    "\n",
    "nodes = CUG.nodes()\n",
    "for n in nodes:\n",
    "    point = Point(nodes.get(n)['pos'])\n",
    "    if point.within(start_region):\n",
    "        if len(Initial_Nodes)<num_agent:\n",
    "            Initial_Nodes.append(n)\n",
    "    elif point.within(goal_region):\n",
    "        if len(Goal_Nodes)<num_agent:\n",
    "            Goal_Nodes.append(n)\n",
    "\n",
    "    if len(Initial_Nodes)==num_agent and len(Goal_Nodes)==num_agent:\n",
    "        break\n",
    "              \n",
    "print(\"\\nInitial Nodes: \")\n",
    "print(Initial_Nodes)\n",
    "print(\"\\nGoal Nodes: \")\n",
    "print(Goal_Nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, you can solve the problem and visualize it! (It might take some time for both solving the problem and visualizing the solution. If it takes to much time, try different resolution by changing edge length when you make the CUG.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Specify the maximum number of time steps\n",
    "T = m+n\n",
    "\n",
    "# convert the CUG to the time-expanded graph\n",
    "Time_Expanded_G = expand_CUG(CUG, T, Initial_Nodes, Goal_Nodes)\n",
    "\n",
    "# solving \"permutation invariant\" multi agent path planning problem using \n",
    "# built-in function of networkx package\n",
    "flowCost, flowDict = nx.network_simplex(Time_Expanded_G)\n",
    "\n",
    "# retrieve path\n",
    "Path = extract_path(flowDict, Initial_Nodes, CUG, return_type=\"node\")\n",
    "\n",
    "print (\"\\nextract path:\")\n",
    "print(Path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize\n",
    "fig = plt.figure(1, figsize=(10,10), dpi=90)\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "BLUE = '#6699cc'\n",
    "for p in Path:\n",
    "    line = []\n",
    "    for n in p:\n",
    "        line.append(CUG.node[int(n)]['pos'])\n",
    "\n",
    "    linestring = LineString(line)\n",
    "    dilated = linestring.buffer(0.3)\n",
    "    patch1 = PolygonPatch(dilated, fc=BLUE, ec=BLUE, alpha=0.5, zorder=2)\n",
    "    ax.add_patch(patch1)\n",
    "\n",
    "pos=nx.get_node_attributes(CUG,'pos')\n",
    "color_map = {1:'b', 2:'#FF0099', 3:'#660066'}\n",
    "\n",
    "nx.draw(CUG,pos,node_color=[color_map[CUG.node[node]['cat']] for node in CUG])\n",
    "\n",
    "ani = simulate(fig,ax,CUG,Path)\n",
    "HTML(ani.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or, you can extract path and create yaml file for waypoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# retrieve path\n",
    "Path = extract_path(flowDict, Initial_Nodes, CUG, return_type=\"waypoint\")\n",
    "\n",
    "with open('waypoints.yml','w') as outfile:   \n",
    "    yaml.dump(Path, outfile, default_flow_style=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the end of this work! We encourage you to change the functions and use different settings for various scenarios!"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Create Assignment",
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}