{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2e253605",
   "metadata": {},
   "source": [
    "# Closed-loop assignment\n",
    "This notebook servers as a simple assignment (with answers) to help you understand how to build a closed-loop kdFlex simulation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0678d955",
   "metadata": {},
   "source": [
    "This notebook assumes that you've already gone through the [4-bar notebook](../example_4_bar/notebook.ipynb) and all subsequent notebooks.\n",
    "\n",
    "<img src=../resources/nb_images/loop_assignment.png align=center width=400>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "22ae75ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import Karana.Dynamics as kd\n",
    "import Karana.Dynamics.SOADyn_types as kdt\n",
    "from Karana.KUtils.Sim import SimDS\n",
    "import Karana.Models as kmdl\n",
    "import Karana.Integrators as ki\n",
    "import atexit\n",
    "import Karana.Math as km\n",
    "from Karana.Math.Kquantities import ureg\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfd3d9dc",
   "metadata": {},
   "source": [
    "We'll begin by loading up a 3-link pendulum system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "68fd102c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[WebUI] Web server is running on port 41059\n",
      "You may be able to connect in your browser at:\n",
      "\t\u001b[1mhttp://newton:41059\u001b[0m\n",
      "[WebUI] Karana Viewer is running on port 29534\n",
      "Open this URL to connect:\n",
      "\t\u001b[1mhttp://newton:29534\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"http://newton:29534\">\n",
       "            <b>Click to open Karana Viewer in a new tab</b>\n",
       "        </a>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sim_ds_path = Path(\"..\") / \"resources\" / \"3_link_pendulum\" / \"3_link_pendulum.h5\"\n",
    "sim = SimDS.fromFile(sim_ds_path).toSim()\n",
    "sim.mb.resetData()\n",
    "sim.mb.getScene().viewAroundFrame(\n",
    "    sim.mb.virtualRoot(), np.array([0.0, 8.0, 0.0]), np.zeros(3), np.array([0.0, 0.0, 1.0])\n",
    ")\n",
    "sim.mb.getScene().update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54c6482e",
   "metadata": {},
   "source": [
    "# Your assignment\n",
    "The assignment is to do the following:\n",
    "1. clone this system so that there are two 3-link pendulum systems.\n",
    "1. Create a constraint between the bottom pendulum bobs of both systems to form a loop. In other words, the two yellow pendulum bobs should be connected via a constraint.\n",
    "1. Set an initial state for the system that satisfies the constraints.\n",
    "1. Add a Gravity KModel. Use uniform gravity with an acceleration of 9.81 m/s^2 in the negative z-direction.\n",
    "1. Simulate the system for 10 seconds."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e83f05af",
   "metadata": {},
   "source": [
    "## Assignment answer\n",
    "All cells below this one walk through the answer to this assignment. You should try this assignment on your own before looking at these, and only use the cells below if you need some help along the way and/or to confirm your answers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc8739df",
   "metadata": {},
   "source": [
    "## Cloning the 3-bar system\n",
    "There are many ways to clone the 3-bar system. Here, we will use a `DataStruct` to do so."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4931632f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a DataStruct of the current system\n",
    "sg_ds = sim.mb.toDS()\n",
    "\n",
    "# Now, let's rename the bodies to clearly differential them from the old system\n",
    "sg_ds.renameBodies(\"pendulum_2_\")\n",
    "\n",
    "# We will move the attachment to the root 2 meters in the +x direction\n",
    "sg_ds.base_bodies[0].body.inb_to_joint = np.array([2.0, 0.0, 0.0])\n",
    "\n",
    "# Finally, let's add these cloned bodies to the multibody\n",
    "new_sg = sg_ds.toSubGraph(sim.mb, sim.mb.virtualRoot())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfa42324",
   "metadata": {},
   "source": [
    "Let's take a look at our new system. We'll set all coordinates to 0, and then update the scene to see it in the GUI above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "05941c21",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.mb.resetData()\n",
    "sim.mb.getScene().update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12b9282f",
   "metadata": {},
   "source": [
    "## Adding a constraint\n",
    "Now, we will add a constraint between the bottom pendulum bobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "69f10ffa",
   "metadata": {},
   "outputs": [],
   "source": [
    "p1_bd3 = sim.mb.getBody(\"bd3\")\n",
    "p2_bd3 = sim.mb.getBody(\"pendulum_2_bd3\")\n",
    "\n",
    "cn1 = kd.ConstraintNode.lookupOrCreate(\"cn_1\", p1_bd3)\n",
    "cn1.setBodyToNodeTransform(km.HomTran())\n",
    "cn2 = kd.ConstraintNode.lookupOrCreate(\"cn_2\", p2_bd3)\n",
    "cn2.setBodyToNodeTransform(km.HomTran())\n",
    "\n",
    "lc = kd.LoopConstraintCutJoint(\"constraint\", sim.mb, cn1.frameToFrame(cn2), kd.HingeType.REVOLUTE)\n",
    "lc.hinge().subhinge(0).setUnitAxis(np.array([0.0, 1.0, 0.0]))\n",
    "lc.hinge().subhinge(0).setQ(0.0)\n",
    "lc.hinge().subhinge(0).setU(0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f82fbc5",
   "metadata": {},
   "source": [
    "## Initial state\n",
    "Here, we will set the initial state of the system. Let's give both top pendulums an initial position of 0.2 radians in opposite directions along with a velocity of 0.2 radians in opposite directions.\n",
    "\n",
    "Then, we will use the `ConstraintKinematicsSolver` to solve for a pose that satisfies the constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b2f57e4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "bd1 = sim.mb.getBody(\"bd1\")\n",
    "p2_bd1 = sim.mb.getBody(\"pendulum_2_bd1\")\n",
    "\n",
    "bd1.parentHinge().subhinge(0).setQ(0.2)\n",
    "bd1.parentHinge().subhinge(0).setU(0.2)\n",
    "\n",
    "p2_bd1.parentHinge().subhinge(0).setQ(-0.2)\n",
    "p2_bd1.parentHinge().subhinge(0).setU(-0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b06bb0ba",
   "metadata": {},
   "source": [
    "We need to use the `ConstraintKinematicsSolver` now to solve for a new pose. By default, this is allowed to change any coordinate. However, we want to keep the pose for the two subhinges we just set. This is not strictly necessary to finish the assignment; however, it is something useful to know how to do, so we show it here. We'll freeze the coordinates we don't want the `ConstraintKinematicsSolver` to change, and then unfreeze them once we have our pose set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "21d5ae44",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.mb.cks().freezeCoord(bd1.parentHinge().subhinge(0), 0)\n",
    "sim.mb.cks().freezeCoord(p2_bd1.parentHinge().subhinge(0), 0)\n",
    "assert sim.mb.cks().solveQ() < 1e-8\n",
    "assert sim.mb.cks().solveU() < 1e-8\n",
    "sim.mb.cks().clearFrozenCoords()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52ea491f",
   "metadata": {},
   "source": [
    "Let's take a look at our system now that the pose is set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6ef701ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.mb.getScene().update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6777465c",
   "metadata": {},
   "source": [
    "We also need to set this state on the `StatePropagator`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a517df86",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.sp.ensureHealthy()\n",
    "\n",
    "# We need to call a hardReset, since the number of states has changed\n",
    "sim.sp.hardReset()\n",
    "\n",
    "# Finally, set the state itself\n",
    "sim.sp.setState(sim.sp.assembleState())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "177a715a",
   "metadata": {},
   "source": [
    "## Adding gravity\n",
    "Now, we will add gravity using a `Gravity` `KModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "72a7b3ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Karana.Models.Gravity at 0x753eb1dc34b0>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "uniform_gravity = kmdl.UniformGravity(\"uniform_gravity\")\n",
    "uniform_gravity.setGravity(\n",
    "    np.array([0.0, 0.0, -9.81]) * ureg.meter / ureg.second**2,\n",
    "    10.0,\n",
    "    kmdl.OutputUpdateType.PRE_DERIV,\n",
    ")\n",
    "kmdl.Gravity(\"grav_model\", sim.sp, uniform_gravity, sim.mb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eee68baf",
   "metadata": {},
   "source": [
    "## Simulate for 10 seconds\n",
    "Finally, we will simulate the system for 10 seconds. While not strictly necessary, we are also adding a `SyncRealTime` `KModel` and maximum step size. Otherwise, the simulation runs too fast to see in the GUI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1397fb71",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<SpStatusEnum.REACHED_END_TIME: 1>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kmdl.SyncRealTime(\"sync_real_time\", sim.sp, 1.0)\n",
    "sim.sp.setMaxHopSize(0.1)\n",
    "sim.sp.advanceBy(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "550b3b2e",
   "metadata": {},
   "source": [
    "## Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "dbf76e3d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.cleanup()>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def cleanup():\n",
    "    \"\"\"Clean up the simulation.\"\"\"\n",
    "    global sim, sg_ds, new_sg, uniform_gravity, bd1, p2_bd1, p1_bd3, p2_bd3, lc, cn1, cn2\n",
    "    del sim, sg_ds, new_sg, uniform_gravity, bd1, p2_bd1, p1_bd3, p2_bd3, lc, cn1, cn2\n",
    "\n",
    "\n",
    "atexit.register(cleanup)"
   ]
  }
 ],
 "metadata": {
  "docs": {
   "tags": [
    "dynamics",
    "time domain",
    "pendulum",
    "constraints",
    "cutjoint"
   ]
  },
  "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
}
