{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93688312",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# 2-link Pendulum Collision Example\n",
    "\n",
    "This notebook is a more advanced version of the 2-link pendulum example by adding a collision dynamics. This tutorial notebook walks you through the steps to create a 2-link pendulum manually and then adds collision dynamics in addition to extra visualization geometries.\n",
    "\n",
    "Requirements:\n",
    "- [2-link Pendulum](../example_2_link_pendulum/notebook.ipynb)\n",
    "\n",
    "In this example we will:\n",
    " - Create the multibody\n",
    " - Setup the **kdFlex** Scene\n",
    " - Setup geometries\n",
    " - Set up the state propagator and models\n",
    " - Set initial state\n",
    " - Register a timed event\n",
    " - Run the simulation\n",
    " - Clean up the simulation\n",
    "\n",
    "![](../resources/nb_images/pendulum_collision.PNG)\n",
    "\n",
    "\n",
    "For a more in-depth descriptions of **kdflex** concepts see [usage](usage_page)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce9260c1",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from typing import cast\n",
    "import atexit\n",
    "\n",
    "import Karana.Core as kc\n",
    "import Karana.Math as km\n",
    "import Karana.Integrators as ki\n",
    "import Karana.Frame as kf\n",
    "import Karana.Dynamics as kd\n",
    "import Karana.Dynamics.SOADyn_types as kdt\n",
    "import Karana.Scene as ks\n",
    "import Karana.KUtils as ku\n",
    "import Karana.Collision as kcoll\n",
    "import Karana.Models as kmdl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf3ac55b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Create the Multibody\n",
    "\n",
    "See [Multibody](treembody_sec) or [Frames](frames_layer_sec) for more information relating to this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c24f60d7",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fc = kf.FrameContainer(\"root\")\n",
    "\n",
    "\n",
    "def createMbody():\n",
    "    \"\"\"Create the Multibody.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Multibody\n",
    "        The newly created multibody.\n",
    "    \"\"\"\n",
    "    mb = kd.Multibody(\"mb\", fc)\n",
    "\n",
    "    # inertia used for both bodies\n",
    "    spI = km.SpatialInertia(2.0, np.zeros(3), np.diag([3, 2, 1]))\n",
    "\n",
    "    # inboard to joint transform used for both bodies\n",
    "    ib2j = km.HomTran()\n",
    "\n",
    "    # body to joint transform used for both bodies\n",
    "    b2j = km.HomTran(np.array([0, 0, 1]))\n",
    "\n",
    "    # setup body 1\n",
    "    bd1 = kd.PhysicalBody(\"bd1\", mb)\n",
    "    bd1.setSpatialInertia(spI)\n",
    "    bd1_hge = kd.PhysicalHinge(cast(kd.PhysicalBody, mb.virtualRoot()), bd1, kd.HingeType.REVOLUTE)\n",
    "    bd1.setBodyToJointTransform(b2j)\n",
    "    bd1_subhge = cast(kd.PinSubhinge, bd1_hge.subhinge(0))\n",
    "    bd1_subhge.setUnitAxis(np.array([0, 1, 0]))\n",
    "\n",
    "    # setup body 2\n",
    "    bd2 = kd.PhysicalBody(\"bd2\", mb)\n",
    "    bd2.setSpatialInertia(spI)\n",
    "    bd2_hge = kd.PhysicalHinge(bd1, bd2, kd.HingeType.REVOLUTE)\n",
    "    bd2_hge.onode().setBodyToNodeTransform(ib2j)\n",
    "    bd2.setBodyToJointTransform(b2j)\n",
    "    bd2_subhge = cast(kd.PinSubhinge, bd2_hge.subhinge(0))\n",
    "    bd2_subhge.setUnitAxis(np.array([0, 1, 0]))\n",
    "\n",
    "    # finalize and verify the multibody\n",
    "    mb.ensureHealthy()\n",
    "    mb.resetData()\n",
    "    assert kc.allReady()\n",
    "\n",
    "    return mb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "faf0393e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "mb = createMbody()\n",
    "\n",
    "# get the first body\n",
    "bd1 = mb.getBody(\"bd1\")\n",
    "\n",
    "# get the second body\n",
    "bd2 = mb.getBody(\"bd2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e8b2310",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Setup the kdFlex Scene\n",
    "\n",
    "Next we setup **kdflex**'s graphics by calling the setupGraphics helper method on the multibody. This method takes care of setting up the graphics environment. \n",
    "\n",
    "See [Visualization](visualization_sec) and [Scene Layer](scene_layer_sec) for more information relating to this section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b8bc12a4",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <div style=\"height:300px; resize:vertical; overflow: auto;\">\n",
       "          <iframe src=\"http://localhost:41601\" 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, axes=0.5)\n",
    "\n",
    "# position the viewpoint camera of the visualization\n",
    "web_scene.defaultCamera().pointCameraAt([0.4, 4, -0.9], [0, 0, -1], [0, 0, 1])\n",
    "proxy_scene = mb.getScene()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26fa588b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "To our Scene Layer, we register a {py:class}`Karana.Scene.CoalScene` instance to perform collision and distance queries using the geometries populated in {py:class}`Karana.Scene.ProxyScene`. CoalScene is a wrapper of the open source [Coal](https://github.com/coal-library/coal) project.\n",
    "\n",
    "See [Collisions](collision_dynamics_sec) for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e1d780b8",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "col_scene = ks.CoalScene(\"collision_scene\")\n",
    "proxy_scene.registerClientScene(col_scene, mb.virtualRoot())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47dcf2d1",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Setup Geometries\n",
    "\n",
    "We create 3d geometries with associated material properties and attach them to bodies. This is commonly useful for collision detection and 3d visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "debf4014",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# create a ball geometry to be used for both bodies\n",
    "ball_geom = ks.SphereGeometry(0.1)\n",
    "\n",
    "# create an obstacle geometry to be placed in the path of body 2\n",
    "obstacle_geom = ks.BoxGeometry(0.2, 0.2, 0.2)\n",
    "\n",
    "# create red and blue visualize materials to color the bodies\n",
    "mat_info = ks.PhysicalMaterialInfo()\n",
    "mat_info.color = ks.Color.BLUE\n",
    "blue = ks.PhysicalMaterial(mat_info)\n",
    "mat_info.color = ks.Color.RED\n",
    "red = ks.PhysicalMaterial(mat_info)\n",
    "\n",
    "# create a ks.ProxyScenePart and attach it to the body\n",
    "bd1_part = ks.ProxyScenePart(\"bd1_part\", scene=proxy_scene, geometry=ball_geom, material=blue)\n",
    "bd1_part.attachTo(bd1)\n",
    "\n",
    "bd2_part = ks.ProxyScenePart(\"bd2_part\", scene=proxy_scene, geometry=ball_geom, material=blue)\n",
    "bd2_part.attachTo(bd2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "651c7984",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Now, we create a separate obstacle for the 2nd pendulum link to collide with. Since it is separate from the pendulum, it is attached to the root frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "53fbd1cf",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# create some unattached objects (implicitly attached to the root frame)\n",
    "root_scene_node = ks.ProxySceneNode(\"root_scene_node\", scene=proxy_scene)\n",
    "obstacle_part = ks.ProxyScenePart(\"obstacle_part\", scene=proxy_scene, geometry=obstacle_geom)\n",
    "obstacle_part.setTranslation([0.1, 0, -2])\n",
    "\n",
    "del obstacle_geom, mat_info, ball_geom"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2b8f0d6",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "For visual effects, we visualize the axes for the root, body 1, and body 2. Additionally, we add a trail following the motion of body 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5238ac7a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# visualize axes for important frames\n",
    "root_scene_node.graphics().showAxes(0.5)\n",
    "bd1_part.graphics().showAxes(0.5)\n",
    "bd2_part.graphics().showAxes(0.5)\n",
    "\n",
    "# add a trail to track the motion of body 2\n",
    "bd2_part.graphics().trail()\n",
    "\n",
    "proxy_scene.update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20a8ba3d",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Setup State Propagator and Models\n",
    "\n",
    "Now, we setup the {py:class}`Karana.Dynamics.StatePropagator` and register our models as well. \n",
    "\n",
    "See [Models](system_level_models_sec) for more concepts and information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a6cf2e2",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Karana.Models.SyncRealTime at 0x777e26932130>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# set up state propagator\n",
    "sp = kd.StatePropagator(mb, ki.IntegratorType.CVODE)\n",
    "integ = sp.getIntegrator()\n",
    "\n",
    "# Create a UniformGravity model, and set it's gravitational acceleration.\n",
    "ug = kmdl.Gravity(\"grav_model\", sp, kmdl.UniformGravity(\"uniform_gravity\"), mb)\n",
    "ug.getGravityInterface().setGravity(np.array([0, 0, -9.81]), 0.0, kmdl.OutputUpdateType.PRE_HOP)\n",
    "del ug\n",
    "\n",
    "# Makes sure the visualization scene is updated after each state change.\n",
    "kmdl.UpdateProxyScene(\"update_proxy_scene\", sp, proxy_scene)\n",
    "\n",
    "# sync the simulation time with real-time.\n",
    "kmdl.SyncRealTime(\"sync_real_time\", sp, 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6890649",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Because our CoalScene managed distance queries and collision detection, we now attach a {py:class}`Karana.Models.PenaltyContact` using the CoalScene and ProxyScene to apply counter forces whenever registered bodies collide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8e38c908",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Set penalty contact model parameters\n",
    "hc = kcoll.HuntCrossley(\"hunt_crossley_contact\")\n",
    "hc.params.kp = 100000\n",
    "hc.params.kc = 10000\n",
    "hc.params.mu = 0.3\n",
    "hc.params.n = 1.5\n",
    "hc.params.linear_region_tol = 1e-3\n",
    "\n",
    "kmdl.PenaltyContact(\"penalty_contact\", sp, mb, [kcoll.FrameCollider(proxy_scene, col_scene)], hc)\n",
    "\n",
    "del hc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2eb5211",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Now we register a callback function that is run at the end of each timestep. It directly queries for if body 2 is in collision with the obstacle, and turns body 2 red if so."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4d7c172a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def post_step_fn(t, x):\n",
    "    \"\"\"Change material color based on collision status.\n",
    "\n",
    "    If body 2 is colliding, make the material red. Otherwise, make it blue.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    t : float\n",
    "        The current time.\n",
    "    x : NDArray[np.float64]\n",
    "       The current state.\n",
    "    \"\"\"\n",
    "    if bd2_part.collision().collide(obstacle_part.collision()):\n",
    "        bd2_part.setMaterial(red)\n",
    "    else:\n",
    "        bd2_part.setMaterial(blue)\n",
    "\n",
    "\n",
    "sp.fns.post_hop_fns[\"update_and_info\"] = post_step_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5de4d87",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Set Initial State\n",
    "\n",
    "Before we begin our simulation, we set the initial state for our multibody and statepropagator. \n",
    "\n",
    "When accessing or modifying [generalized coordinates](coords_sec) for a subhinge, it is recommended to directly set the subhinge's values rather than for the entire multibody in order to avoid ambiguity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3a569ae5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Modify the initial multibody state.\n",
    "# Here we will set the first pendulum position to 0.5 radians and its velocity to 0.0\n",
    "bd1 = mb.getBody(\"bd1\")\n",
    "bd1.parentHinge().subhinge(0).setQ([0.5])\n",
    "bd1.parentHinge().subhinge(0).setU([0.0])\n",
    "\n",
    "# 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 graphics\n",
    "proxy_scene.update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c71f917",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Register a Timed Event\n",
    "Because the penalty contact model applies forces proportional to the overlap between two objects in collision, it is necessary to shorten the period. Otherwise, if a collision is only detected when 2 objects have significant overlap, unrealistic forces may be applied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "88188c58",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "h = np.timedelta64(int(1e7), \"ns\")\n",
    "t = kd.TimedEvent(\"hop_size\", h, lambda _: None, False)\n",
    "t.period = h\n",
    "\n",
    "# register after time has been initialized\n",
    "sp.registerTimedEvent(t)\n",
    "del t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5406f59",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Run the Simulation\n",
    "\n",
    "{py:meth}`Karana.Math.Integrator.advanceTo` can be increased to lengthen the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0dba2fbc",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t = 0.0s; x = [0.5 0.  0.  0. ]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<SpStatusEnum.REACHED_END_TIME: 1>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# print the initial state\n",
    "print(f\"t = {float(integ.getTime()) / 1e9}s; x = {integ.getX()}\")\n",
    "\n",
    "# run the simulation\n",
    "sp.advanceTo(10.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7bb87d5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Clean Up the Simulation\n",
    "\n",
    "Below, we cleanup our simulation. We first delete local variables, cleanup our visualizer, discard remaining Karana objects, and optionally verify using allDestroyed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2f130538",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.cleanup()>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def cleanup():\n",
    "    \"\"\"Cleanup the simulation.\"\"\"\n",
    "    global integ, proxy_scene, web_scene, col_scene\n",
    "    global root_scene_node, bd1, bd2, obstacle_part, red, blue, bd1_part, bd2_part\n",
    "    del (\n",
    "        integ,\n",
    "        proxy_scene,\n",
    "        web_scene,\n",
    "        col_scene,\n",
    "        root_scene_node,\n",
    "        bd1,\n",
    "        bd2,\n",
    "        obstacle_part,\n",
    "        red,\n",
    "        blue,\n",
    "        bd1_part,\n",
    "        bd2_part,\n",
    "    )\n",
    "\n",
    "    kc.discard(sp)\n",
    "    cleanup_graphics()\n",
    "\n",
    "    kc.discard(mb)\n",
    "    kc.discard(fc)\n",
    "\n",
    "    bc = kc.BaseContainer.singleton()\n",
    "    bc.at_exit_fns.executeAndPopReverse()\n",
    "    del bc\n",
    "\n",
    "\n",
    "atexit.register(cleanup)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec50c60",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Summary\n",
    "You now know how to setup a 2-link pendulum with a obstacle to simulate collisions. By attaching a CoalScene and creating a PenaltyContact model, the simulation is able to model collisions.\n",
    "\n",
    "## Further Readings\n",
    "[Simulate collisions with a n-link pendulum](../example_procedural_collision/notebook.ipynb)  \n",
    "[Drive an ATRVjr rover](../example_atrvjr_drive/notebook.ipynb)"
   ]
  }
 ],
 "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"
  },
  "name": "notebook.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
