{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a7ada9cf",
   "metadata": {},
   "source": [
    "# Draw 3D Lines and Point\n",
    "\n",
    "[![Click and Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/metadriverse/metadrive/blob/main/documentation/source/points_and_lines.ipynb)\n",
    "s"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "782a29d8",
   "metadata": {},
   "source": [
    "![draw.png](figs/draw.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d33fc3",
   "metadata": {},
   "source": [
    "We usually draw points and lines in the 3D world to debug when developing new features and fixing bugs. There are two types of visualization tools can be used: **line** and **point**. In this section, we will show how to use these tools. First of all, let's make an environment and define the tool function which returns a set of points representing a straight line and the color for each point. The points gradually become pink with the distance increasing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e83bd4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from metadrive.envs.base_env import BaseEnv\n",
    "import numpy as np\n",
    "import os\n",
    "render = not os.getenv('TEST_DOC')\n",
    "\n",
    "# Define a tool function. \n",
    "def make_line(x_offset, height, y_dir=1, color=(1,105/255,180/255)):\n",
    "    points = [(x_offset+x,x*y_dir,height*x/10+height) for x in range(10)]\n",
    "    colors = [np.clip(np.array([*color,1])*(i+1)/11, 0., 1.0) for i in range(10)]\n",
    "    if y_dir<0:\n",
    "        points = points[::-1]\n",
    "        colors = colors[::-1]\n",
    "    return points, colors\n",
    "\n",
    "# create environment\n",
    "env = BaseEnv(dict(use_render=render)) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cec8ab02",
   "metadata": {},
   "source": [
    "## Lines\n",
    "The API we are using takes a list of lines as input. A line is represented by a list of points. Thus we create two lines in the following script and draw them in the scene. Also, colors can be specified for each point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08ee500d",
   "metadata": {},
   "outputs": [],
   "source": [
    "line_1, color_1 = make_line(6, 0.5, 1) # define line 1 for test\n",
    "line_2, color_2 = make_line(6, 0.5, -1) # define line 2 for test\n",
    "lines = [line_1, line_2]\n",
    "colors = [color_1, color_2]\n",
    "\n",
    "env.reset() # launch the simulation\n",
    "try:\n",
    "    drawer = env.engine.make_line_drawer(thickness=5) # create a line drawer\n",
    "    drawer.draw_lines(lines, colors) # draw lines\n",
    "    \n",
    "    for i in range(100):\n",
    "        env.step([0,0])\n",
    "finally:    \n",
    "    env.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77c60c42",
   "metadata": {},
   "source": [
    "A drawer can draw as many lines as possible, but sometimes you may want to clear previous lines and recreate new ones. In this case, you can use `drawer.reset()` to clear previous lines. The following code shows an example, where every step the lines to draw are updated. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8bbf84c",
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset() # launch the simulation\n",
    "try:\n",
    "    drawer = env.engine.make_line_drawer(thickness=5) # create a line drawer\n",
    "    for i in range(100):\n",
    "        \n",
    "        # draw different lines every step\n",
    "        line_1, color_1 = make_line(6, 0.5, 0.01*i) # define line 1 for test\n",
    "        line_2, color_2 = make_line(6, 0.5, -0.01*i) # define line 2 for test\n",
    "        lines = [line_1, line_2]\n",
    "        colors = [color_1, color_2]\n",
    "        drawer.reset()\n",
    "        drawer.draw_lines(lines, colors) # draw lines\n",
    "        \n",
    "        env.step([0,0])\n",
    "finally:    \n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e5b93a6",
   "metadata": {},
   "source": [
    "By default, the drawer drawing lines in the world coordinates. It is also allowed to draw lines in objects' local coordinates. In the following example, you will see that the two lines moves with the vehicle. **Note: there is a bug that when the drawer is set in vehicle coordinates the color parameter is invalid. Thus the line color is white in this example.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92fb9ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset() # launch the simulation\n",
    "try:\n",
    "    drawer = env.engine.make_line_drawer(env.agent.origin, thickness=5)\n",
    "    # rotate the drawer by 90 degree, as +x is at the right side of the car.\n",
    "    drawer.setH(90) \n",
    "    for i in range(100):\n",
    "        \n",
    "        # draw different lines every step\n",
    "        line_1, color_1 = make_line(6, 0.5, 0.01*i) # define line 1 for test\n",
    "        line_2, color_2 = make_line(6, 0.5, -0.01*i) # define line 2 for test\n",
    "        lines = [line_1, line_2]\n",
    "        colors = [color_1, color_2]\n",
    "        drawer.reset()\n",
    "        drawer.draw_lines(lines, colors) # draw lines\n",
    "        \n",
    "        env.step([0.1,0.5])\n",
    "finally:    \n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07109f19",
   "metadata": {},
   "source": [
    "## Points\n",
    "MetaDrive also provides point drawer to visualize a set of points. The usage of this function is almost the same as the line drawer. The following example creates a point drawer and update the point positions every frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a51ac0b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset() # launch the simulation\n",
    "try:\n",
    "    drawer = env.engine.make_point_drawer(scale=1) # create a point drawer\n",
    "    for i in range(100):\n",
    "        \n",
    "        # draw different lines every step\n",
    "        line_1, color_1 = make_line(6, 0.5, 0.01*i) # define line 1 for test\n",
    "        line_2, color_2 = make_line(6, 0.5, -0.01*i) # define line 2 for test\n",
    "        points = line_1 + line_2 # create point list\n",
    "        colors = color_1+ color_2\n",
    "        drawer.reset()\n",
    "        drawer.draw_points(points, colors) # draw points\n",
    "        \n",
    "        env.step([0,0])\n",
    "finally:    \n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd30195c",
   "metadata": {},
   "source": [
    "In addition, the points can work with lines. In the following example, a list of points are visualized with lines, while the other one is visualized with points. Besides, we cancel the *reset* to keep previous drawn points and lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71575db2",
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset() # launch the simulation\n",
    "try:\n",
    "    point_drawer = env.engine.make_point_drawer(scale=1) # create a point drawer\n",
    "    line_drawer = env.engine.make_line_drawer(thickness=5) # create a line drawer\n",
    "    for i in range(100):\n",
    "        \n",
    "        if i%5==0:\n",
    "            # draw different lines every step\n",
    "            line_1, color_1 = make_line(6, 0.5, 0.01*i) # define line 1 for test\n",
    "            line_2, color_2 = make_line(6, 0.5, -0.01*i) # define line 2 for test\n",
    "            points = line_1\n",
    "            point_colors = color_1\n",
    "            lines = [line_2]\n",
    "            line_colors = [color_2]\n",
    "            # drawer.reset()\n",
    "            point_drawer.draw_points(points, point_colors) # draw lines\n",
    "            line_drawer.draw_lines(lines, line_colors)\n",
    "        \n",
    "        env.step([0,0])\n",
    "finally:    \n",
    "    env.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}