{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e081a7ed",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Flexible Double-Wishbone\n",
    "This example builds on the [rigid double-wishbone](../example_double_wishbone/notebook.ipynb). This example and makes one of its links flexible, adds a spring-damper, and runs a full simulation of the double-wishbone system. Similar to the [Flexible slider crank](../example_flexible_slider_crank/notebook.ipynb), this is an example of a multibody that has flexible bodies and constraints.\n",
    "\n",
    "In this example, the multibody has already been prepared and saved to an H5 file, so it is easy to create. The [Flexible slider crank](../example_flexible_slider_crank/notebook.ipynb) shows how a simulation with flexible bodies can be built of procedurally. \n",
    "\n",
    "Requirements:\n",
    "- [double-wishbone](../example_double_wishbone/notebook.ipynb)\n",
    "\n",
    "In this tutorial we will:\n",
    "- Import the Multibody\n",
    "- Initialize the state and constraints\n",
    "- Setup the **kdFlex** Scene\n",
    "- Setup the state propagator and models\n",
    "- Run the simulation\n",
    "- Clean up the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6e7233df",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import atexit\n",
    "import numpy as np\n",
    "from numpy.typing import NDArray\n",
    "from typing import cast\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.Models as kmdl\n",
    "import Karana.Dynamics.SOADyn_types as kdt\n",
    "from Karana.KUtils.Sim import Sim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74711448",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Import the multibody\n",
    "Here, we create the multibody. It is already saved as an H5 file, so we can just load it in via a SOADyn_types.SubGraphDS DataStruct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "01a2b03b",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "mbody_ds = kdt.SubGraphDS.fromFile(\"../resources/double_wishbone/flexible_double_wishbone.h5\")\n",
    "sp_ds = kdt.StatePropagatorDS(\n",
    "    integrator_type=ki.IntegratorType.CVODE_STIFF,\n",
    "    solver_type=kd.MMSolverType.TREE_AUGMENTED_DYNAMICS,\n",
    "    models=[],\n",
    ")\n",
    "sim = Sim(mbody_ds, sp_ds)\n",
    "del mbody_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0db7ac30",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Initialize the State and Constraints\n",
    "We are going to initialize the state of this system. This system contains both constraints and flexible bodies. For the constraints, we will need to use a {py:class}`Karana.Dynamics.ConstraintKinematicsSolver`, just as we did in the [Rigid double-wishbone](../example_double_wishbone/notebook.ipynb) example.\n",
    "\n",
    "We could finish the state initialization here. However, when we first start moving, there may be some transient behavior, as the flexible body itself is not in the correct deformed state. Here, we create a simple equilibrium solver to deform the flexible body such that it's initial modal coordinate acceleration is zero. In this case, the solution is trivial, as we are starting from rest, but this showcases how it can be done in general."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9932a95",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Satisfying the constraints\n",
    "We being by setting an initial guess for the {py:class}`Karana.Dynamics.ConstraintKinematicsSolver`. This is done by setting all joints, body coordinates, and constraint coordinates to 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "43300e8c",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "sim.mb.ensureHealthy()\n",
    "sim.mb.resetData()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36c913e4",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Now, we use the Multibody's ConstraintKinematicsSolver to solve for coordinates and velocities that satisfy the constraints. We want `lower_control_arm`'s subhinge to remain unchanged, so we freeze this coordinate. This effectively removes it from the coordinates that the ConstraintKinematicsSolver is allowed to change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "caf1f42d",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "bd = cast(kd.PhysicalBody, sim.mb.getBody(\"lower_control_arm\"))\n",
    "sim.mb.cks().freezeCoord(bd.parentHinge().subhinge(0), 0)\n",
    "\n",
    "# Solve for joint angles that satisfy the constraints\n",
    "assert sim.mb.cks().solveQ() < 1e-10\n",
    "assert sim.mb.cks().solveU() < 1e-10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aed42f33",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Ensure everything is ready before continuing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "06c7229c",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set the gravity and external forces to zero\n",
    "sim.mb.setUniformGravAccel(np.zeros(3))\n",
    "cast(kd.PhysicalBody, sim.mb.getBody(\"wheel\")).getNode(\"WheelContactNode\").setExternalSpForce(\n",
    "    km.SpatialForce(np.zeros(3), np.zeros(3))\n",
    ")\n",
    "\n",
    "# Ensure everything has been set and is ready.\n",
    "kc.allReady()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "807e2158",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Equilibrium\n",
    "As mentioned earlier, we solve for the flexible body's equilibrium. Here, the solution is trivial since we are starting from rest (so the solution should be a vector of all zeros). Regardless, we include it here to showcase how solving for equilibrium can be done.\n",
    "\n",
    "To solve for equilibrium we use a {py:class}`Karana.Dynamics.NonlinearSolver`. These solvers will be covered in detail in an upcoming example. For now, it is sufficient to note that they take in a function that takes in a guess for the current solution, `x`, and outputs values, `v` for that guess. The values must be modified in-place, which is why the code below uses `v[:] = ...` rather than `v = ...`. Here, our equilibrium function is simple: it takes in a guess for the flexible body coordinates, `q`, of `upper_control_arm`, runs the forward dynamics, and outputs the modal accelerations of `upper_control_arm`. The nonlinear solver tries to drive these outputs to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cad27111",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "uca_bd = cast(kd.PhysicalModalBody, sim.mb.getBody(\"upper_control_arm\"))\n",
    "\n",
    "\n",
    "def equilibrium(x: NDArray, v: NDArray) -> None:\n",
    "    \"\"\"Solve for equilibrium.\n",
    "\n",
    "    The incoming x values will be the new Q states of the beam.\n",
    "    The output v values should be the accelerations of the beam as a\n",
    "    result of the incoming Q values.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : NDArray\n",
    "        New guess for beam Q values.\n",
    "    v : NDArray\n",
    "        Udot values of the beam after running the dynamics.\n",
    "    \"\"\"\n",
    "    uca_bd.setQ(x)\n",
    "    kd.Algorithms.evalForwardDynamics(sim.mb)\n",
    "    v[:] = uca_bd.getUdot()\n",
    "\n",
    "\n",
    "solver = kd.NonlinearSolver(\"solver\", uca_bd.nQ(), uca_bd.nU(), equilibrium)\n",
    "solver.createSolver(kd.SolverType.HYBRID_NON_LINEAR_SOLVER)\n",
    "q = np.array(uca_bd.getQ())\n",
    "solver.solve(q)\n",
    "\n",
    "uca_bd.setQ(q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97369c48",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Similar to the constraint solution, we can again check that the solution here is sufficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a50911f5",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "udot = np.empty(q.size)\n",
    "equilibrium(q, udot)\n",
    "assert np.linalg.norm(udot) < 1e-10\n",
    "\n",
    "del uca_bd\n",
    "del solver"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f73332f",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Setup the kdFlex Scene\n",
    "Start 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": 8,
   "id": "028cd091",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <div style=\"height:300px; resize:vertical; overflow: auto;\">\n",
       "          <iframe src=\"http://newton:44635\" 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": [
    "_, web_scene = sim.setupGraphics(port=0, axes=0.5)\n",
    "web_scene.defaultCamera().pointCameraAt(\n",
    "    np.array([1.5, -0.2, 0.0]), np.array([0.0, -0.2, 0.0]), np.array([0.0, 0.0, -1.0])\n",
    ")\n",
    "sim.mb.createStickParts()\n",
    "del web_scene"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "878d0420",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Set the State Propagator state and KModels\n",
    "As in other examples, we use a {py:class}`Karana.Dynamics.StatePropagator` to drive the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "180c841a",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Initialize state.\n",
    "t_init = np.timedelta64(0, \"ns\")\n",
    "sim.sp.setTime(t_init)\n",
    "sim.sp.setState(sim.sp.assembleState())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20b82157",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Shock absorber\n",
    "Register a model for the shock absorber."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e455d72b",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "ssd = kmdl.SubhingeSpringDamper(\n",
    "    \"shock_absorber_spring\",\n",
    "    sim.sp,\n",
    "    cast(kd.PhysicalSubhinge, sim.mb.getBody(\"shock_absorber_upper\").parentHinge().subhinge(0)),\n",
    ")\n",
    "ssd.params.k = 300\n",
    "ssd.params.d = 100\n",
    "ssd.params.setpoint = np.array([0.0])\n",
    "del ssd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc793bcd",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Constraint error projection\n",
    "Even though the individual forward dynamics calls themselves should not introduce much constraint error, as we integrate, constraint error will build up. Hence, periodically, we need to use a method to eliminate this error. Here, we use the {py:class}`Karana.Models.ProjectConstraintError` model, which uses a ConstraintKinematicsSolver to project out the error once it gets too large."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "84f0120c",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "pce = kmdl.ProjectConstraintError(\"project_error\", sim.sp, sim.mb, sim.mb.cks())\n",
    "pce.params.tol = 1e-8\n",
    "h = np.timedelta64(int(1e7), \"ns\")\n",
    "pce.setPeriod(h)\n",
    "del pce"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06ae9b1f",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Prescribing upper_control_arm acceleration\n",
    "Currently, our simulation won't move, as nothing is changing the acceleration of the prescribed hinge for `upper_control_arm`. Here, we add a function that will modify the acceleration so that we increase the velocity for the first two seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e325d7d3",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Move the system\n",
    "t_r = 4.0\n",
    "\n",
    "\n",
    "def move(t_in, _):\n",
    "    \"\"\"Apply prescribed Udot to move the system.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    t_in : float\n",
    "        The current time.\n",
    "    _ : NDArray[np.float64]\n",
    "       The current state (unused).\n",
    "    \"\"\"\n",
    "    t = float(t_in) / 1e9\n",
    "    if t > t_r:\n",
    "        accel = -10 * np.pi / 180 * np.sin(t)\n",
    "    else:\n",
    "        accel = (\n",
    "            10\n",
    "            * np.pi\n",
    "            / 180\n",
    "            * (\n",
    "                (6 / t_r**2 - 12 * t / t_r**3) * np.sin(t)\n",
    "                + 2 * (6 * t / t_r**2 - 6 * t**2 / t_r**3) * np.cos(t)\n",
    "                - (3 * (t / t_r) ** 2 - 2 * (t / t_r) ** 3) * np.sin(t)\n",
    "            )\n",
    "        )\n",
    "    bd.parentHinge().subhinge(0).setUdot(accel)\n",
    "\n",
    "\n",
    "sim.sp.fns.pre_deriv_fns[\"move\"] = move"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00fb7a7b",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Run the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4dd60c74",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<SpStatusEnum.REACHED_END_TIME: 1>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim.sp.advanceTo(4 * np.pi)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec916def",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Clean Up the Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c9326a56",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.cleanup()>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def cleanup():\n",
    "    \"\"\"Cleanup the simulation.\"\"\"\n",
    "    import gc\n",
    "\n",
    "    global bd, sim\n",
    "    del bd, sim\n",
    "\n",
    "    gc.collect()\n",
    "\n",
    "\n",
    "atexit.register(cleanup)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca6e36bd",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Summary\n",
    "You now understand how to set the state and constraints, solve them, and register the shock, constraint error, and acceleration models to make the simulation work.\n",
    "\n",
    "## Further Readings\n",
    "- [flexible slider-crank](../example_flexible_slider_crank/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
}
