{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f96fd41d",
   "metadata": {},
   "source": [
    "# ATRVjr ROS Integration\n",
    "\n",
    "In this notebook we import the ATRVjr rover model into **kdflex** and interface with ROS. We will publish messages to ROS as the sim advances as well as listen for a ROS message to command the vehicle's driving mode.\n",
    "\n",
    "Requirements:\n",
    "- [ATRVjr driving](../example_atrvjr_drive/notebook.ipynb)\n",
    "\n",
    "In this tutorial we will:\n",
    "- Import and modify the URDF\n",
    "- Setup the **kdFlex** scene\n",
    "- Attach wheel geometries\n",
    "- Setup StatePropagator and models\n",
    "- Setup the Simulation\n",
    "- Integrate with ROS\n",
    "- Run and clean up\n",
    "\n",
    "**Acknowledgements**  \n",
    "The robot URDF model included in this project is Copyright (c) 2018 diasdm and licensed under the BSD 2‑Clause License [(link)](https://github.com/diasdm/atrvjr_description.git). Some modifications were made to adapt the model for use within the kdflex simulation framework, including simplifying wheel geometries due to current rendering and collision handling limitations.\n",
    "\n",
    " For a more in-depth descriptions of **kdflex** concepts see [usage](usage_page)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e146968",
   "metadata": {},
   "source": [
    "First let's confirm that ROS is installed and properly sourced"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cd627362",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import rclpy\n",
    "except ModuleNotFoundError:\n",
    "    raise ModuleNotFoundError(\n",
    "        \"rclpy package is not found!\\n\"\n",
    "        \"Please ensure ROS2 is installed and that you've sourced the appropriate setup script.\\n\"\n",
    "        \"    https://docs.ros.org/en/kilted/Installation.html\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dce602b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import atexit\n",
    "import numpy as np\n",
    "from math import pi\n",
    "from pathlib import Path\n",
    "from Karana.Core import discard\n",
    "from Karana.Frame import FrameContainer\n",
    "from Karana.Dynamics.SOADyn_types import SubGraphDS\n",
    "from Karana.Math import IntegratorType\n",
    "from Karana.Dynamics import (\n",
    "    Multibody,\n",
    "    PhysicalBody,\n",
    "    PhysicalBody,\n",
    "    HingeType,\n",
    "    PhysicalHinge,\n",
    "    PinSubhinge,\n",
    "    StatePropagator,\n",
    "    TimedEvent,\n",
    ")\n",
    "from Karana.Dynamics.SOADyn_types import (\n",
    "    HingeDS,\n",
    "    PinSubhingeDS,\n",
    "    SubhingeDS,\n",
    "    Linear3SubhingeDS,\n",
    "    SphericalSubhingeDS,\n",
    "    BodyWithContextDS,\n",
    "    ScenePartSpecDS,\n",
    ")\n",
    "from Karana.Scene import (\n",
    "    AssimpImporter,\n",
    "    SphereGeometry,\n",
    "    BoxGeometry,\n",
    "    CylinderGeometry,\n",
    "    StaticMeshGeometry,\n",
    "    AbstractStaticGeometry,\n",
    "    Color,\n",
    "    PhysicalMaterialInfo,\n",
    "    PhysicalMaterial,\n",
    "    LAYER_COLLISION,\n",
    "    LAYER_GRAPHICS,\n",
    "    LAYER_ALL,\n",
    ")\n",
    "from Karana.Scene import CoalScene\n",
    "from Karana.Scene import ProxyScene, ProxySceneNode, ProxyScenePart\n",
    "from Karana.Math import UnitQuaternion, SpatialInertia, HomTran\n",
    "from Karana.Models import (\n",
    "    Gravity,\n",
    "    UniformGravity,\n",
    "    UpdateProxyScene,\n",
    "    SyncRealTime,\n",
    "    PenaltyContact,\n",
    "    OutputUpdateType,\n",
    ")\n",
    "from Karana.Collision import FrameCollider, HuntCrossley"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcb35c5b",
   "metadata": {},
   "source": [
    "Build the filepath to the URDF file we will import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fd13f1bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "resources = Path().absolute().parent / \"resources\"\n",
    "urdf_file = resources / \"atrvjr\" / \"atrvjr.urdf\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64ddd559",
   "metadata": {},
   "source": [
    "## Import and Modify the URDF\n",
    "\n",
    "Import the URDF file to create a SOADyn_types.SubGraphDS. The \"DS\" is short for \"Data Structure\". This is kdFlex's native, file-format agnostic way of specifying a multibody. We'll use this later on to instantiate the multibody.\n",
    "\n",
    "It's easier to edit the model while it's in the SubGraphDS format before we convert it to a Multibody. In the cell below we the original hinge between the rover and the root frame to a full 6-DOF hinge, which allows it to move freely. We then convert the hinges between the wheel bodies to pin joints so the wheels can properly rotate.\n",
    "\n",
    "Last we set all the visuals in the SubGraphDS to belong to the graphics layer only so that it can't interact with the wheels which we add later or the ground."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c6ce1a4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_info = SubGraphDS.fromFile(urdf_file)\n",
    "\n",
    "mb_info.model_dump()\n",
    "\n",
    "# convert to full6DOF hinge between root frame and rover.\n",
    "mb_info.base_bodies[0].body.hinge = HingeDS(\n",
    "    hinge_type=HingeType.FULL6DOF,\n",
    "    subhinges=[\n",
    "        Linear3SubhingeDS(prescribed=False),\n",
    "        SphericalSubhingeDS(prescribed=False),\n",
    "    ],\n",
    ")\n",
    "\n",
    "# unlock locked Hinges for wheels from the URDF\n",
    "for i in [0, 1]:\n",
    "    mb_info.base_bodies[0].children[i].body.hinge = HingeDS(\n",
    "        hinge_type=HingeType.REVOLUTE,\n",
    "        subhinges=[PinSubhingeDS(prescribed=False, unit_axis=np.array([0.0, 0.0, 1.0]))],\n",
    "    )\n",
    "for i in [2, 3]:\n",
    "    mb_info.base_bodies[0].children[i].body.hinge = HingeDS(\n",
    "        hinge_type=HingeType.REVOLUTE,\n",
    "        subhinges=[PinSubhingeDS(prescribed=False, unit_axis=np.array([0.0, 0.0, -1.0]))],\n",
    "    )\n",
    "\n",
    "\n",
    "# Remove all collision meshes. We want to add these manually.\n",
    "def removeCollision():\n",
    "    \"\"\"Remove all the collision meshes.\"\"\"\n",
    "\n",
    "    def _removeParts(parts: list[ScenePartSpecDS]):\n",
    "        idx = []\n",
    "        for k, p in enumerate(parts):\n",
    "            if p.layers == LAYER_COLLISION:\n",
    "                idx.append(k)\n",
    "        idx.reverse()\n",
    "        for k in idx:\n",
    "            parts.pop(k)\n",
    "\n",
    "    def _recurse(b: BodyWithContextDS):\n",
    "        _removeParts(b.body.scene_parts)\n",
    "        for c in b.children:\n",
    "            _recurse(c)\n",
    "\n",
    "    for base in mb_info.base_bodies:\n",
    "        _recurse(base)\n",
    "\n",
    "\n",
    "removeCollision()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3698fcc7",
   "metadata": {},
   "source": [
    "Create a {py:class}`Karana.Frame.FrameContainer` to manager the [frames layer](http://kd.internal/doxygen_docs/usage_page.html#frames_layer_sec)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "51a4d9af",
   "metadata": {},
   "outputs": [],
   "source": [
    "fc = FrameContainer(\"root\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fd06632",
   "metadata": {},
   "source": [
    "Instantiate and configure the {py:class}`Karana.Dynamics.Multibody` based on the SubGraphDS description."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8e99d2b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb = mb_info.toMultibody(fc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c4aeea4",
   "metadata": {},
   "source": [
    "Finalize the model to get it ready for use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "115f56a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|-atrvjr_MBVROOT_\n",
      "   |-[FULL6DOF] center_link\n",
      "      |-[PIN] left_front_wheel_link\n",
      "      |-[PIN] left_rear_wheel_link\n",
      "      |-[PIN] right_front_wheel_link\n",
      "      |-[PIN] right_rear_wheel_link\n",
      "      |-[LOCKED] cbimu_frame\n",
      "      |-[LOCKED] mpuimu_frame\n",
      "      |-[LOCKED] arm_base_frame\n",
      "      |-[LOCKED] lrf_frame\n",
      "      |-[LOCKED] novatel\n",
      "\n"
     ]
    }
   ],
   "source": [
    "mb.ensureHealthy()\n",
    "mb.resetData()\n",
    "mb.dumpTree()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97db71f8",
   "metadata": {},
   "source": [
    "## Setup the kdFlex Scene \n",
    "\n",
    "Next we setup **kdflex**'s visualizer by creating a graphics scene and a webserver to view it. We use the helper method setupGraphics() to take care of setting up the scene and webserver for us.\n",
    "\n",
    "See [Visualization](visualization_sec) and [Scene Layer](scene_layer_sec) for more information relating to this section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "024e4376",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[WebUI] Listening at http://newton:39223\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <div style=\"height:300px; resize:vertical; overflow: auto;\">\n",
       "          <iframe src=\"http://newton:39223\" style=\"width:100%; height:100%; border:0; display:block;\"></iframe>\n",
       "        </div>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cleanup_graphics, web_scene = mb.setupGraphics(port=0)\n",
    "web_scene.defaultCamera().pointCameraAt(\n",
    "    offset=np.array([[0.0], [5.0], [10.0]]),\n",
    "    target=np.array([[0.0], [0.0], [0.0]]),\n",
    "    up=np.array([[0.0], [0.0], [1.0]]),\n",
    ")\n",
    "proxy_scene = mb.getScene()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d6cc4149",
   "metadata": {},
   "outputs": [],
   "source": [
    "# attach a collision scene\n",
    "col_scene = CoalScene(\"collision_scene\")\n",
    "proxy_scene.registerClientScene(col_scene, mb.virtualRoot(), layers=LAYER_COLLISION)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b90d05e",
   "metadata": {},
   "source": [
    "Display the bodies in the multibody model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e9ab8bad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "      Body                     Parent              Hinge    Dofs          \n",
      "      ____                     ______              _____    ____          \n",
      "   1. center_link              atrvjr_MBVROOT_    FULL6DOF  6/0           \n",
      "   2. left_front_wheel_link    center_link          PIN     1/6           \n",
      "   3. left_rear_wheel_link     center_link          PIN     1/7           \n",
      "   4. right_front_wheel_link   center_link          PIN     1/8           \n",
      "   5. right_rear_wheel_link    center_link          PIN     1/9           \n",
      "   6. cbimu_frame              center_link         LOCKED   0/10          \n",
      "   7. mpuimu_frame             center_link         LOCKED   0/10          \n",
      "   8. arm_base_frame           center_link         LOCKED   0/10          \n",
      "   9. lrf_frame                center_link         LOCKED   0/10          \n",
      "  10. novatel                  center_link         LOCKED   0/10          \n",
      "                                                            ____          \n",
      "                                                            10            \n",
      "\n"
     ]
    }
   ],
   "source": [
    "mb.displayModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2feb3376",
   "metadata": {},
   "source": [
    "## Attach Wheel Geometries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e4b0d6c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "wheel_geom = CylinderGeometry(0.125, 0.075)\n",
    "\n",
    "# create a ground for the rover to drive on\n",
    "ground_geom = BoxGeometry(30, 30, 0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "79747949",
   "metadata": {},
   "outputs": [],
   "source": [
    "left_front_wheel = mb.getBody(\"left_front_wheel_link\")\n",
    "# left_front_wheel.addScenePartSpec(wheel_obj.parts[0])\n",
    "# left_front_wheel.addScenePartSpec(wheel_obj.parts[1])\n",
    "left_front_wheel_node = ProxyScenePart(\n",
    "    name=\"left_front_wheel_link_node\", scene=proxy_scene, geometry=wheel_geom\n",
    ")\n",
    "left_front_wheel_node.attachTo(left_front_wheel)\n",
    "left_front_wheel_node.setTranslation([0, 0, 0.03])\n",
    "left_front_wheel_node.setUnitQuaternion(UnitQuaternion(pi / 2, [1.0, 0.0, 0.0]))\n",
    "# left_front_wheel_node.graphics().showAxes(1.0)\n",
    "\n",
    "left_rear_wheel = mb.getBody(\"left_rear_wheel_link\")\n",
    "# left_rear_wheel.addScenePartSpec(wheel_obj.parts[0])\n",
    "# left_rear_wheel.addScenePartSpec(wheel_obj.parts[1])\n",
    "left_rear_wheel_node = ProxyScenePart(\n",
    "    \"left_rear_wheel_link_node\", scene=proxy_scene, geometry=wheel_geom\n",
    ")\n",
    "left_rear_wheel_node.attachTo(left_rear_wheel)\n",
    "left_rear_wheel_node.setTranslation([0, 0, 0.03])\n",
    "left_rear_wheel_node.setUnitQuaternion(UnitQuaternion(pi / 2, [1.0, 0.0, 0.0]))\n",
    "# left_rear_wheel_node.graphics().showAxes(1.0)\n",
    "\n",
    "right_front_wheel = mb.getBody(\"right_front_wheel_link\")\n",
    "# right_front_wheel.addScenePartSpec(wheel_obj.parts[0])\n",
    "# right_front_wheel.addScenePartSpec(wheel_obj.parts[1])\n",
    "right_front_wheel_node = ProxyScenePart(\n",
    "    \"right_front_wheel_link_node\", scene=proxy_scene, geometry=wheel_geom\n",
    ")\n",
    "right_front_wheel_node.attachTo(right_front_wheel)\n",
    "right_front_wheel_node.setTranslation([0, 0, 0.03])\n",
    "right_front_wheel_node.setUnitQuaternion(UnitQuaternion(pi / 2, [1.0, 0.0, 0.0]))\n",
    "# right_front_wheel_node.graphics().showAxes(1.0)\n",
    "\n",
    "right_rear_wheel = mb.getBody(\"right_rear_wheel_link\")\n",
    "# right_rear_wheel.addScenePartSpec(wheel_obj.parts[0])\n",
    "# right_rear_wheel.addScenePartSpec(wheel_obj.parts[1])\n",
    "right_rear_wheel_node = ProxyScenePart(\n",
    "    \"right_rear_wheel_link_node\", scene=proxy_scene, geometry=wheel_geom\n",
    ")\n",
    "right_rear_wheel_node.attachTo(right_rear_wheel)\n",
    "right_rear_wheel_node.setTranslation([0, 0, 0.03])\n",
    "right_rear_wheel_node.setUnitQuaternion(UnitQuaternion(pi / 2, [1.0, 0.0, 0.0]))\n",
    "# right_rear_wheel_node.graphics().showAxes(1.0)\n",
    "\n",
    "mat_info = PhysicalMaterialInfo()\n",
    "mat_info.color = Color.WHITESMOKE\n",
    "gray = PhysicalMaterial(mat_info)\n",
    "\n",
    "ground = ProxyScenePart(\n",
    "    \"ground\", scene=proxy_scene, geometry=ground_geom, material=gray, layers=LAYER_ALL\n",
    ")\n",
    "ground.setTranslation([0, 0, -0.1])\n",
    "\n",
    "del ground, mat_info, wheel_geom, ground_geom, gray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20509dd8",
   "metadata": {},
   "source": [
    "## Setup StatePropagator and models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "27ce425f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up state propagator\n",
    "sp = StatePropagator(mb, IntegratorType.RK4)\n",
    "integ = sp.getIntegrator()\n",
    "\n",
    "# Create a UniformGravity model, and set it's gravitational acceleration.\n",
    "ug = Gravity(\"grav_model\", sp, UniformGravity(\"uniform_gravity\"), mb)\n",
    "ug.getGravityInterface().setGravity(np.array([0, 0, -9.81]), 0.0, OutputUpdateType.PRE_HOP)\n",
    "del ug\n",
    "\n",
    "# Makes sure the visualization scene is updated after each state change.\n",
    "UpdateProxyScene(\"update_proxy_scene\", sp, proxy_scene)\n",
    "\n",
    "# sync the simulation time with real-time.\n",
    "SyncRealTime(\"sync_real_time\", sp, 1.0)\n",
    "\n",
    "# add penalty contact model and register its bodies and parameters\n",
    "hc = HuntCrossley(\"hunt_crossley_contact\")\n",
    "hc.params.kp = 100000\n",
    "hc.params.kc = 20000\n",
    "hc.params.mu = 0.3\n",
    "hc.params.n = 1.5\n",
    "hc.params.linear_region_tol = 1e-3\n",
    "PenaltyContact(\"penalty_contact\", sp, mb, [FrameCollider(proxy_scene, col_scene)], hc)\n",
    "del col_scene, hc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2791a2fa",
   "metadata": {},
   "source": [
    "Because forces are cleared out in between simulation steps, we create a callback function to apply forces depending on the intended direction. We attach it as a {py:attr}`Karana.Dynamics.SpFunctions.pre_deriv_fns`, since it is called before the multibody derivative which ensures the forces are used in the state update calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a63bd6fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "direction = \"coast\"\n",
    "\n",
    "\n",
    "def pre_deriv_fn(t, x):\n",
    "    \"\"\"Set torques based on the driving direction.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    t : float\n",
    "        The current time.\n",
    "    x : NDArray[np.float64]\n",
    "       The current state.\n",
    "    \"\"\"\n",
    "    # manually set torques in pre_deriv function\n",
    "    global direction\n",
    "    match direction:\n",
    "        case \"coast\":\n",
    "            right_rear_wheel.parentHinge().subhinge(0).setT([0])\n",
    "            left_rear_wheel.parentHinge().subhinge(0).setT([0])\n",
    "            right_front_wheel.parentHinge().subhinge(0).setT([0])\n",
    "            left_front_wheel.parentHinge().subhinge(0).setT([0])\n",
    "        case \"forward\":\n",
    "            right_rear_wheel.parentHinge().subhinge(0).setT([4])\n",
    "            left_rear_wheel.parentHinge().subhinge(0).setT([4])\n",
    "            right_front_wheel.parentHinge().subhinge(0).setT([4])\n",
    "            left_front_wheel.parentHinge().subhinge(0).setT([4])\n",
    "        case \"backward\":\n",
    "            right_rear_wheel.parentHinge().subhinge(0).setT([-4])\n",
    "            left_rear_wheel.parentHinge().subhinge(0).setT([-4])\n",
    "            right_front_wheel.parentHinge().subhinge(0).setT([-4])\n",
    "            left_front_wheel.parentHinge().subhinge(0).setT([-4])\n",
    "        case \"rotate_right\":\n",
    "            right_rear_wheel.parentHinge().subhinge(0).setT([-5])\n",
    "            left_rear_wheel.parentHinge().subhinge(0).setT([5])\n",
    "            right_front_wheel.parentHinge().subhinge(0).setT([-5])\n",
    "            left_front_wheel.parentHinge().subhinge(0).setT([5])\n",
    "        case \"rotate_left\":\n",
    "            right_rear_wheel.parentHinge().subhinge(0).setT([5])\n",
    "            left_rear_wheel.parentHinge().subhinge(0).setT([-5])\n",
    "            right_front_wheel.parentHinge().subhinge(0).setT([5])\n",
    "            left_front_wheel.parentHinge().subhinge(0).setT([-5])\n",
    "\n",
    "\n",
    "sp.fns.pre_deriv_fns[\"apply_torque\"] = pre_deriv_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf704275",
   "metadata": {},
   "source": [
    "## Setup the ROS interface\n",
    "\n",
    "`EasyPubSub` is a wrapper for a ROS node and executor which:\n",
    "\n",
    "- Simplifies publishing and subscribing to messages\n",
    "- Simplifies integration with the `StatePropagator` using a `KModel`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d57f2986",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Karana.KUtils.ros import EasyPubSub\n",
    "\n",
    "pubsub = EasyPubSub(\"sim\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82206c1f",
   "metadata": {},
   "source": [
    "## Create a `KModel` to communicate with ROS\n",
    "\n",
    "This model will periodically check for "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f1bd512e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Karana.Models import KModel\n",
    "\n",
    "from geometry_msgs.msg import PoseStamped\n",
    "from std_msgs.msg import String\n",
    "from builtin_interfaces.msg import Time\n",
    "\n",
    "\n",
    "class RosComm(KModel):\n",
    "    \"\"\"ROS communication KModel.\"\"\"\n",
    "\n",
    "    def __init__(self, name: str):\n",
    "        \"\"\"Create a RosComm instance.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        name : str\n",
    "            The name of the RosComm KModel.\n",
    "        \"\"\"\n",
    "        super().__init__(name, sp)\n",
    "        sp.registerModel(self)\n",
    "        self.pubsub = pubsub\n",
    "        pubsub.subscribe(\"/command\", String, self._handleCommand)\n",
    "\n",
    "    def _handleCommand(self, cmd: String):\n",
    "        if cmd.data not in [\"coast\", \"forward\", \"backward\", \"rotate_right\", \"rotate_left\"]:\n",
    "            return\n",
    "        global direction\n",
    "        direction = cmd.data\n",
    "\n",
    "    def _buildPoseStamped(self, t):\n",
    "        ns = int(t.astype(\"timedelta64[ns]\").astype(int))\n",
    "        sec = ns // 1_000_000_000\n",
    "        nsec = ns % 1_000_000_000\n",
    "        root = mb.virtualRoot()\n",
    "        chassis = mb.getBody(\"center_link\")\n",
    "        r2c = root.frameToFrame(chassis).relTransform()\n",
    "        r2c_qvec = r2c.getUnitQuaternion().toVector4()\n",
    "\n",
    "        msg = PoseStamped()\n",
    "        msg.header.stamp = Time(sec=sec, nanosec=nsec)\n",
    "        msg.header.frame_id = chassis.name()\n",
    "        msg.pose.position.x = r2c.getTranslation()[0]\n",
    "        msg.pose.position.y = r2c.getTranslation()[1]\n",
    "        msg.pose.position.z = r2c.getTranslation()[2]\n",
    "        msg.pose.orientation.x = r2c_qvec[0]\n",
    "        msg.pose.orientation.y = r2c_qvec[1]\n",
    "        msg.pose.orientation.z = r2c_qvec[2]\n",
    "        msg.pose.orientation.w = r2c_qvec[3]\n",
    "        return msg\n",
    "\n",
    "    def postModelStep(self, t, x):\n",
    "        \"\"\"Process ROS subscriptions and publish the pose message.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        t : float\n",
    "            The current time.\n",
    "        x : NDArray[np.float64]\n",
    "           The current state.\n",
    "        \"\"\"\n",
    "        pubsub.processSubscriptions()\n",
    "        pose = self._buildPoseStamped(t)\n",
    "        self.pubsub.publish(\"/pose\", pose)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "084579aa",
   "metadata": {},
   "source": [
    "## Setup the Simulation\n",
    "Finish by setting the initial states and timed events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1817817a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set integrator state\n",
    "t_init = np.timedelta64(0, \"ns\")\n",
    "x_init = sp.assembleState()\n",
    "sp.setTime(t_init)\n",
    "sp.setState(x_init)\n",
    "\n",
    "# syncs up the graphics\n",
    "proxy_scene.update()\n",
    "\n",
    "# register timed event\n",
    "h = np.timedelta64(int(1e7), \"ns\")\n",
    "t = TimedEvent(\"hop_size\", h, lambda _: None, False)\n",
    "t.period = h\n",
    "sp.registerTimedEvent(t)\n",
    "\n",
    "roscomm = RosComm(\"ros_comm\")\n",
    "roscomm.setPeriod(0.1)\n",
    "\n",
    "del t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51d0efef",
   "metadata": {},
   "source": [
    "## Subscribe to and print the `/pose` topic\n",
    "\n",
    "This is to demonstrate the `PoseStamped` messages being published by the sim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1ac98528",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'e54bf411-fafb-48a8-8dd7-7a8fb8c7e830'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def printPosition(pose):\n",
    "    \"\"\"Print the position from the given pose.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    pose :\n",
    "        The pose whose position will be printed.\n",
    "    \"\"\"\n",
    "    sec = pose.header.stamp.sec\n",
    "    nsec = pose.header.stamp.nanosec\n",
    "    t = sec + 1e-9 * nsec\n",
    "    position = np.array([pose.pose.position.x, pose.pose.position.y, pose.pose.position.z])\n",
    "    print(f\"{t=:.5g} {position}\")\n",
    "\n",
    "\n",
    "pubsub.subscribe(\"/pose\", PoseStamped, printPosition)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a95b814",
   "metadata": {},
   "source": [
    "## Run and Clean Up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b675b6b7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t=0.1 [6.21659043e-06 1.25010413e-05 1.25902848e-01]\n",
      "t=0.2 [ 3.66838531e-05 -1.51171786e-04  1.49291143e-01]\n",
      "t=0.3 [ 2.15317170e-05 -1.70409927e-04  1.28074954e-01]\n",
      "t=0.4 [-9.26186880e-05 -2.23504113e-04  1.46480777e-01]\n",
      "t=0.5 [-8.03233143e-05 -2.02847167e-04  1.30706999e-01]\n",
      "t=0.6 [-9.61293833e-05 -1.74686984e-04  1.43464872e-01]\n",
      "t=0.7 [ 4.44683424e-05 -1.67816961e-04  1.33616218e-01]\n",
      "t=0.8 [ 0.00019843 -0.00019242  0.14036566]\n",
      "t=0.9 [ 0.00042372 -0.00015653  0.1363341 ]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<SpStatusEnum.REACHED_END_TIME: 1>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Nodes\n",
      " ---\n"
     ]
    }
   ],
   "source": [
    "sp.advanceBy(1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db1a630d",
   "metadata": {},
   "source": [
    "## Publish a new driving command\n",
    "\n",
    "The `RosComm` model will receive this message and update the driving mode upon the next simulation step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0029c07a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "publisher: beginning loop\r\n",
      "publishing #1: std_msgs.msg.String(data='forward')\r\n",
      "\r\n"
     ]
    }
   ],
   "source": [
    "!ros2 topic pub --once /command std_msgs/msg/String \"data: forward\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e6d8687b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t=1 [ 0.00067868 -0.0001608   0.13759262]\n",
      "t=1.1 [ 1.00348693e-03 -1.24044077e-04  1.38446879e-01]\n",
      "t=1.2 [ 3.16128932e-03 -1.01202302e-04  1.35636614e-01]\n",
      "t=1.3 [ 0.00903308 -0.00030654  0.13982312]\n",
      "t=1.4 [1.93564319e-02 4.94279856e-05 1.34550902e-01]\n",
      "t=1.5 [0.03382264 0.00055203 0.14015249]\n",
      "t=1.6 [0.05241513 0.00080046 0.13459761]\n",
      "t=1.7 [0.07501351 0.00070825 0.13975692]\n",
      "t=1.8 [0.10148733 0.00051264 0.13509404]\n",
      "t=1.9 [0.13229783 0.000572   0.13908449]\n",
      "t=2 [0.16693539 0.00064599 0.135986  ]\n",
      "t=2.1 [0.20597232 0.00104475 0.13807206]\n",
      "t=2.2 [0.24913292 0.00162681 0.1368598 ]\n",
      "t=2.3 [0.29629743 0.00172614 0.13741043]\n",
      "t=2.4 [0.34736463 0.00228526 0.13741304]\n",
      "t=2.5 [0.40265159 0.00224405 0.13680302]\n",
      "t=2.6 [0.46176137 0.0021196  0.13793199]\n",
      "t=2.7 [0.52509566 0.0016836  0.1364957 ]\n",
      "t=2.8 [0.59257062 0.0012306  0.13799747]\n",
      "t=2.9 [0.664086   0.00084548 0.13651509]\n",
      "t=3 [7.39702558e-01 3.07849564e-04 1.37969850e-01]\n",
      "t=3.1 [ 8.19332345e-01 -2.34610044e-05  1.36587207e-01]\n",
      "t=3.2 [ 9.03141931e-01 -4.34932976e-04  1.37760275e-01]\n",
      "t=3.3 [ 9.90866895e-01 -8.30675657e-04  1.36915116e-01]\n",
      "t=3.4 [ 1.08268293 -0.00116761  0.137431  ]\n",
      "t=3.5 [ 1.17872395 -0.00142385  0.1371321 ]\n",
      "t=3.6 [ 1.2785882  -0.00161248  0.13729714]\n",
      "t=3.7 [ 1.38246175 -0.00174713  0.13729834]\n",
      "t=3.8 [ 1.49054171 -0.00172436  0.13708652]\n",
      "t=3.9 [ 1.60271744 -0.00207118  0.13745296]\n",
      "t=4 [ 1.71897401 -0.00199916  0.13704639]\n",
      "t=4.1 [ 1.83928576e+00 -1.80944815e-03  1.37413413e-01]\n",
      "t=4.2 [ 1.96367247e+00 -1.56610367e-03  1.37045459e-01]\n",
      "t=4.3 [ 2.09209076e+00 -1.27890947e-03  1.37461748e-01]\n",
      "t=4.4 [ 2.22446623e+00 -1.19796977e-03  1.37038496e-01]\n",
      "t=4.5 [ 2.36122349e+00 -5.97530773e-04  1.37379644e-01]\n",
      "t=4.6 [ 2.50197298e+00 -1.15804405e-07  1.37157541e-01]\n",
      "t=4.7 [2.64684918e+00 6.09325254e-04 1.37288237e-01]\n",
      "t=4.8 [2.79585573e+00 1.20567381e-03 1.37193755e-01]\n",
      "t=4.9 [2.94881131e+00 1.54930840e-03 1.37247962e-01]\n",
      "t=5 [3.10592827e+00 2.12318854e-03 1.37283282e-01]\n",
      "t=5.1 [3.26717218e+00 2.59836175e-03 1.37151740e-01]\n",
      "t=5.2 [3.43250499e+00 3.11107483e-03 1.37314879e-01]\n",
      "t=5.3 [3.60190120e+00 3.57585951e-03 1.37182003e-01]\n",
      "t=5.4 [3.77533707 0.00404745 0.13728283]\n",
      "t=5.5 [3.95282637 0.00426584 0.13717628]\n",
      "t=5.6 [4.13439729 0.00443055 0.13729539]\n",
      "t=5.7 [4.32013097 0.00495711 0.13718959]\n",
      "t=5.8 [4.50995117 0.00520905 0.13724915]\n",
      "t=5.9 [4.70371602 0.0050414  0.13723777]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<SpStatusEnum.REACHED_END_TIME: 1>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.advanceBy(5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7d1bfbcc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.cleanup()>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Cleanup\n",
    "def cleanup():\n",
    "    \"\"\"Cleanup the simulation.\"\"\"\n",
    "    global \\\n",
    "        left_front_wheel, \\\n",
    "        left_front_wheel_node, \\\n",
    "        left_rear_wheel, \\\n",
    "        left_rear_wheel_node, \\\n",
    "        right_rear_wheel, \\\n",
    "        right_rear_wheel_node, \\\n",
    "        right_front_wheel, \\\n",
    "        right_front_wheel_node, \\\n",
    "        proxy_scene, \\\n",
    "        mb_info, \\\n",
    "        integ, \\\n",
    "        web_scene\n",
    "    del (\n",
    "        left_front_wheel,\n",
    "        left_front_wheel_node,\n",
    "        left_rear_wheel,\n",
    "        left_rear_wheel_node,\n",
    "        right_rear_wheel,\n",
    "        right_rear_wheel_node,\n",
    "        right_front_wheel,\n",
    "        right_front_wheel_node,\n",
    "        proxy_scene,\n",
    "        web_scene,\n",
    "        mb_info,\n",
    "        integ,\n",
    "    )\n",
    "\n",
    "    discard(sp)\n",
    "    cleanup_graphics()\n",
    "\n",
    "    discard(mb)\n",
    "    discard(fc)\n",
    "\n",
    "\n",
    "atexit.register(cleanup)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdfe83ae",
   "metadata": {},
   "source": [
    "## Summary\n",
    "You now understand how to drive a ATRVjr model by applying torque forces in the PreDeriv call."
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
