{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f96fd41d",
   "metadata": {},
   "source": [
    "# Importing a robot arm multibody model from a URDF file\n",
    "The purpose of this notebook is to\n",
    "- demonstrate the creation of a {py:class}`Karana.Dynamics.Multibody` instance for a robot arm from a URDF model file \n",
    "- show the structure of the loaded model\n",
    "- show a 3D visualization of the robot arm\n",
    "- serve as a `hello world` script for verifying a proper `kdFlex` installation\n",
    "\n",
    "Resources\n",
    "- [Other notebook examples](https://portal.karanadyn.com/docs/examples_page.html) \n",
    "- [kdFlex Documentation](https://portal.karanadyn.com/docs/index.html)\n",
    "\n",
    "**Acknowledgements**\n",
    "The Barrett WAM urdf is sourced from the public repository [barrett-ros2-pkg](https://git.barrett.com/software/barrett-ros2-pkg.git). Thank you to the authors and contributors."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5978859c",
   "metadata": {},
   "source": [
    "Import **kdFlex** Python modules objects for use later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dce602b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import atexit\n",
    "\n",
    "import Karana.Core as kc\n",
    "import Karana.Frame as kf\n",
    "import Karana.Dynamics.SOADyn_types as kdt\n",
    "import Karana.KUtils.BasicPrefab as kub\n",
    "from Karana.KUtils.Sim import Sim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcb35c5b",
   "metadata": {},
   "source": [
    "Build the filepath to the URDF file we will import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fd13f1bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "resources = Path().absolute().parent / \"resources\"\n",
    "urdf_file = resources / \"wam7dof\" / \"wam7dof_hand.urdf\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8b05075",
   "metadata": {},
   "source": [
    "Import the URDF file to create a BasicPrefabDS. The \"DS\" is short for \"Data Structure\". This is kdFlex's native, file-format agnostic way of specifying a multibody and associated KModels. We'll use this next to create our simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c6ce1a4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "basic_prefab_ds = kub.BasicPrefabDS.fromFile(urdf_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fd06632",
   "metadata": {},
   "source": [
    "Instantiate and configure the {py:class}`Karana.Dynamics.Multibody` based on the BasicPrefabDS description."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8e99d2b9",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'Handle' object has no attribute '__dict__'",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mAttributeError\u001b[39m                            Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[4]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m sim, prefab = \u001b[43mSim\u001b[49m\u001b[43m.\u001b[49m\u001b[43mfromBasicPrefabDS\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbasic_prefab_ds\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/src/sboxes/sbox/x86_64-ubuntu24.04/install/lib/PYTHON/Karana/KUtils/Sim.py:354\u001b[39m, in \u001b[36mSim.fromBasicPrefabDS\u001b[39m\u001b[34m(cls, basic_prefab_ds)\u001b[39m\n\u001b[32m    351\u001b[39m fc = FrameContainer(\u001b[33m\"\u001b[39m\u001b[33mroot\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m    352\u001b[39m bc.at_exit_fns[\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33m_discard_fc_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfc.id()\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m] = _discard(fc)\n\u001b[32m--> \u001b[39m\u001b[32m354\u001b[39m bp, sp = \u001b[43mBasicPrefab\u001b[49m\u001b[43m[\u001b[49m\u001b[43mExtraInfoType\u001b[49m\u001b[43m]\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcreateStandalone\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbasic_prefab_ds\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfc\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    355\u001b[39m mb = sp.getSubTree().multibody()\n\u001b[32m    356\u001b[39m bc.at_exit_fns[\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33m_discard_mb_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mmb.id()\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m] = _discard(mb)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/src/sboxes/sbox/x86_64-ubuntu24.04/install/lib/PYTHON/Karana/KUtils/BasicPrefab.py:1056\u001b[39m, in \u001b[36mBasicPrefab.createStandalone\u001b[39m\u001b[34m(cls, basic_prefab_ds, fc)\u001b[39m\n\u001b[32m   1054\u001b[39m \u001b[38;5;66;03m# Create the BasicPrefab\u001b[39;00m\n\u001b[32m   1055\u001b[39m basic_prefab_ds.context.model_manager = sp\n\u001b[32m-> \u001b[39m\u001b[32m1056\u001b[39m bp = cast(Self, \u001b[43mbasic_prefab_ds\u001b[49m\u001b[43m.\u001b[49m\u001b[43mtoPrefab\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m)\n\u001b[32m   1058\u001b[39m sp = cast(StatePropagator, basic_prefab_ds.context.model_manager)\n\u001b[32m   1060\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m bp, sp\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/src/sboxes/sbox/x86_64-ubuntu24.04/install/lib/PYTHON/Karana/KUtils/Prefab.py:221\u001b[39m, in \u001b[36mPrefabDS.toPrefab\u001b[39m\u001b[34m(self, parent_or_top, context)\u001b[39m\n\u001b[32m    219\u001b[39m     obj = klass(\u001b[38;5;28mself\u001b[39m.name, \u001b[38;5;28mself\u001b[39m.config, context, \u001b[38;5;28mself\u001b[39m.params, parent)  \u001b[38;5;66;03m# pyright: ignore - Pyright false positive.\u001b[39;00m\n\u001b[32m    220\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m221\u001b[39m     obj = \u001b[43mklass\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparent\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# pyright: ignore - Pyright false positive.\u001b[39;00m\n\u001b[32m    223\u001b[39m \u001b[38;5;66;03m# Add children\u001b[39;00m\n\u001b[32m    224\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m c \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.children:\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/src/sboxes/sbox/x86_64-ubuntu24.04/install/lib/PYTHON/Karana/KUtils/Prefab.py:302\u001b[39m, in \u001b[36mPrefab.__init__\u001b[39m\u001b[34m(self, name, config, context, params, parent_prefab)\u001b[39m\n\u001b[32m    300\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.parent \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mor\u001b[39;00m (\u001b[38;5;28mself\u001b[39m.parent \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m.parent._initialized):\n\u001b[32m    301\u001b[39m     \u001b[38;5;28mself\u001b[39m._addChildPrefabs()\n\u001b[32m--> \u001b[39m\u001b[32m302\u001b[39m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_addMultibodyObjects\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    303\u001b[39m     \u001b[38;5;28mself\u001b[39m._addKModels()\n\u001b[32m    304\u001b[39m     \u001b[38;5;28mself\u001b[39m._connectKModels()\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/src/sboxes/sbox/x86_64-ubuntu24.04/install/lib/PYTHON/Karana/KUtils/Prefab.py:482\u001b[39m, in \u001b[36mPrefab._addMultibodyObjects\u001b[39m\u001b[34m(self)\u001b[39m\n\u001b[32m    480\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[34m_addMultibodyObjects\u001b[39m(\u001b[38;5;28mself\u001b[39m) -> \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m    481\u001b[39m \u001b[38;5;250m    \u001b[39m\u001b[33;03m\"\"\"Call addMultibodyObjects on this and all children recursively.\"\"\"\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m482\u001b[39m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43maddMultibodyObjects\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    484\u001b[39m     \u001b[38;5;28;01mfor\u001b[39;00m c \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m._children:\n\u001b[32m    485\u001b[39m         c._addMultibodyObjects()\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/src/sboxes/sbox/x86_64-ubuntu24.04/install/lib/PYTHON/Karana/KUtils/BasicPrefab.py:885\u001b[39m, in \u001b[36mBasicPrefab.addMultibodyObjects\u001b[39m\u001b[34m(self)\u001b[39m\n\u001b[32m    882\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.params.subtree_state_info \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m    883\u001b[39m         \u001b[38;5;28mself\u001b[39m.params.subtree_state_info.toSubTree(st)\n\u001b[32m--> \u001b[39m\u001b[32m885\u001b[39m \u001b[43mdiscard\u001b[49m\u001b[43m(\u001b[49m\u001b[43mst\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[31mAttributeError\u001b[39m: 'Handle' object has no attribute '__dict__'"
     ]
    }
   ],
   "source": [
    "sim, prefab = Sim.fromBasicPrefabDS(basic_prefab_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c4aeea4",
   "metadata": {},
   "source": [
    "Finalize the model to get it ready for use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "115f56a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.mb.ensureHealthy()\n",
    "sim.mb.resetData()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97db71f8",
   "metadata": {},
   "source": [
    "Set up the 3D graphics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "024e4376",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, graphics = sim.setupGraphics(port=0)\n",
    "graphics.defaultCamera().pointCameraAt(\n",
    "    offset=[0.2, 0.5, 1.7],\n",
    "    target=[0, 0, 1],\n",
    "    up=[1, 0, 0],\n",
    ")\n",
    "del graphics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b90d05e",
   "metadata": {},
   "source": [
    "Display the bodies in the multibody model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9ab8bad",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.mb.displayModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a82391",
   "metadata": {},
   "source": [
    "Display the topological structure of the multibody system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e88214b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.mb.dumpTree()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df73d006",
   "metadata": {},
   "source": [
    "Animate the graphics by articulating the individual joints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f575fb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.mb.articulateBodies()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ec19236",
   "metadata": {},
   "source": [
    "Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "136c1c30",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cleanup():\n",
    "    \"\"\"Cleanup the simulation.\"\"\"\n",
    "    global basic_prefab_ds, prefab, sim\n",
    "    del basic_prefab_ds, prefab, sim\n",
    "\n",
    "\n",
    "atexit.register(cleanup)"
   ]
  }
 ],
 "metadata": {
  "docs": {
   "recording": "https://portal.karanadyn.com/tutorial_videos/example_urdf.mp4",
   "tags": [
    "dynamics",
    "urdf",
    "import",
    "robot",
    "WAM"
   ]
  },
  "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"
  },
  "name": "notebook.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}