{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyPlot Animation with MultibodyPlant Tutorial\n",
    "For instructions on how to run these tutorial notebooks, please see the [index](./index.ipynb).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Selecting matplotlib Backends\n",
    "\n",
    "Jupyter notebooks provide the `%matplotlib` that we will use to select\n",
    "different backends:\n",
    "https://ipython.readthedocs.io/en/stable/interactive/magics.html#magic-matplotlib\n",
    "\n",
    "Generally, the backends you can use are either non-interactive (e.g. `inline`)\n",
    "or interactive (`notebook`, `tk`). This notebook shows some options you can\n",
    "comment and uncomment to try them out.\n",
    "\n",
    "**Note: We recommend using only the `inline` backend on Deepnote and Google Colab.**\n",
    "\n",
    "First, show what options might be available. Note that some of these backends\n",
    "may not have all dependencies installed, so they may not work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib --list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now select **one** of the backends here. These options are more likely\n",
    "to be supported on your system.\n",
    "\n",
    "Note that you can only select a different backend once during a kernel session.\n",
    "If you want to change backends, you will need to restart your session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is non-interactive: it shows static plots inline\n",
    "%matplotlib inline\n",
    "\n",
    "# This is interactive: it shows dynamic plots in the notebook\n",
    "# %matplotlib notebook\n",
    "\n",
    "# This is interactive: it shows dynamic plots in separate GUI windows\n",
    "# %matplotlib tk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "from matplotlib import animation\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.multibody.parsing import Parser\n",
    "from pydrake.multibody.plant import AddMultibodyPlantSceneGraph\n",
    "from pydrake.systems.analysis import Simulator\n",
    "from pydrake.systems.framework import DiagramBuilder\n",
    "from pydrake.systems.planar_scenegraph_visualizer import (\n",
    "    ConnectPlanarSceneGraphVisualizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Pendulum Example\n",
    "\n",
    "This function is consolidated from `run_planar_scenegraph_visualizer.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_pendulum_example(duration=1., playback=True, show=True):\n",
    "    \"\"\"\n",
    "    Runs a simulation of a pendulum.\n",
    "\n",
    "    Arguments:\n",
    "        duration: Simulation duration (sec).\n",
    "        playback: Enable pyplot animations to be produced.\n",
    "    \"\"\"\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, 0.)\n",
    "    parser = Parser(plant)\n",
    "    parser.AddModels(\n",
    "        url=\"package://drake/examples/pendulum/Pendulum.urdf\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    T_VW = np.array([[1., 0., 0., 0.],\n",
    "                     [0., 0., 1., 0.],\n",
    "                     [0., 0., 0., 1.]])\n",
    "    visualizer = ConnectPlanarSceneGraphVisualizer(\n",
    "        builder, scene_graph, T_VW=T_VW, xlim=[-1.2, 1.2],\n",
    "        ylim=[-1.2, 1.2], show=show)\n",
    "    if playback:\n",
    "        visualizer.start_recording()\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.Initialize()\n",
    "    simulator.set_target_realtime_rate(1.)\n",
    "\n",
    "    # Fix the input port to zero.\n",
    "    plant_context = diagram.GetMutableSubsystemContext(\n",
    "        plant, simulator.get_mutable_context())\n",
    "    plant.get_actuation_input_port().FixValue(\n",
    "        plant_context, np.zeros(plant.num_actuators()))\n",
    "    plant_context.SetContinuousState([0.5, 0.1])\n",
    "    simulator.AdvanceTo(duration)\n",
    "\n",
    "    if playback:\n",
    "        visualizer.stop_recording()\n",
    "        ani = visualizer.get_recording_as_animation()\n",
    "        return ani\n",
    "    else:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run without Playback\n",
    "\n",
    "If you have a non-interactive backend, you will not see any animation.\n",
    "Additionally, you will see a UserWarning that it is using a non-GUI backend.\n",
    "\n",
    "If you have an interactive backend, you should see the simulation animation\n",
    "as it happens.\n",
    "\n",
    "If you select a GUI option, this will open a new figure each time you run the\n",
    "following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_pendulum_example(playback=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run with Playback\n",
    "\n",
    "If you have a non-interactive backend, you will not see any animation\n",
    "in the first output.\n",
    "\n",
    "If you have an interactive backend, you will see animation in the first output\n",
    "(as the simulation happens). Additionally, the direct animation plot itself will\n",
    "loop its playback."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ani = run_pendulum_example(playback=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given that you recorded playback, now you can produce an animation (regardless\n",
    "of your backend) either as:\n",
    "\n",
    "* A JavaScript HTML widget - allows for slightly finer-grained control\n",
    "* An HTML5 video - requires ffmpeg, which is not installed as part of Drake's\n",
    "dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if animation.writers.is_available(\"ffmpeg\"):\n",
    "    display(HTML(ani.to_html5_video()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you do not want to render the image (only the animation), then pass `show=False` in to the constructor of `PlanarSceneGraphVisualizer()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ani = run_pendulum_example(playback=True, show=False)\n",
    "HTML(ani.to_jshtml())"
   ]
  }
 ],
 "metadata": {
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}