{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0b81363d-4768-481d-a0d2-3f3d8468fcba",
   "metadata": {},
   "source": [
    "© Copyright, 2025 G. Schaer.\n",
    "\n",
    "SPDX-License-Identifier: GPL-3.0-only"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c58045bb-4b36-4caf-9cb6-039bec6c63ef",
   "metadata": {},
   "source": [
    "# Project Example 3: The Double Pendulum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bac8b42f-4da7-4c1e-a777-8953244829b2",
   "metadata": {},
   "source": [
    "## Set up the Visualizer, Animator, and Physics Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8874740a-5c75-4115-8ad7-75edd8ef7cd6",
   "metadata": {},
   "source": [
    "Import necessary `condynsate` modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc40d60b-49c6-4c54-8abf-ba7b801cf370",
   "metadata": {},
   "outputs": [],
   "source": [
    "from condynsate import Project\n",
    "from condynsate import __assets__ as assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "088afcf1-b682-4de0-ae59-3788e872a9e1",
   "metadata": {},
   "source": [
    "Create a instance of `condynsate.Project` that uses the visualizer and animator. Also, set `simulator_dt` to 0.003 seconds so that the simulation fidelty is higher. Note, this will result in increased run time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00f52642-97ac-43ca-957d-4102c79cba87",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj = Project(visualizer = True, animator = True, simulator_dt = 0.003)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ee16f7-0781-4f5b-9ceb-822e469c1278",
   "metadata": {},
   "source": [
    "Load the double pendulum into the simulator. Set `fixed = True` so that the base link of the pendulum has 0 degrees of freedom. Also, set the initial angular velocity of one of the joints to some non-zero value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6bff914-6125-404a-921d-5b76c1e9995e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and orient a double pendulum.\n",
    "dp = proj.load_urdf(assets['double_pendulum.urdf'], fixed=True)\n",
    "\n",
    "# Set the initial state\n",
    "dp.joints['mass_1_to_arm_2'].set_initial_state(omega=0.1) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d78ea2e-871f-4aef-8fc5-c71c8c97d3b5",
   "metadata": {},
   "source": [
    "Now we will adjust the visualizer scene. First we turn off the axes visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a3b0c40-f891-4c6d-90a8-6f8948cdb93f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Turn off the axes and grid visualization.\n",
    "proj.visualizer.set_axes(False) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6859aa7e-794e-4cae-9229-57fd84837575",
   "metadata": {},
   "source": [
    "Then we will move the camera's position up just a slight amount from its default position and then tell it to look directly at the center of mass of the pendulum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df1c2b24-bea2-4487-b8bf-8e31d09a709f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the camera's position\n",
    "proj.visualizer.set_cam_position((0, -4, 2.5)) # This returns 0 on success\n",
    "\n",
    "# Focus the camera on the gyro\n",
    "proj.visualizer.set_cam_target(dp.center_of_mass) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50139d0b-a87c-4755-a800-628278b0d235",
   "metadata": {},
   "source": [
    "Now we will remove all friction from the double pendulum. To do so, we need to set the damping of each joint to 0 and the air resistance of each link to 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3024cf5-72a4-44e5-9871-7488ebd87473",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the damping of all joints to 0\n",
    "for joint in dp.joints.values():\n",
    "    joint.set_dynamics(damping=0.0,) # This returns 0 on success\n",
    "\n",
    "# Set the air resistance of all links to 0\n",
    "for link in dp.links.values():\n",
    "    link.set_dynamics(linear_air_resistance=0.0, angular_air_resistance=0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "727182de-55c6-43bb-ac0d-c2e46936e853",
   "metadata": {},
   "source": [
    "Add a line plot to the animator to plot the double pendulum's state in phase space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85f562cd-0e30-409b-b35b-e49f477c38e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The return value of add_lineplot is either a list of or single \n",
    "# hex value that is used to reference each line on the lineplot\n",
    "# In this case, n_lines=2, so the return value is a tuple\n",
    "# of hex codes used to reference each line in the plot\n",
    "(arm_1, arm_2) = proj.animator.add_lineplot(n_lines=2, \n",
    "                                            color=('r', 'b'), \n",
    "                                            label=('Arm 1', 'Arm 2'),\n",
    "                                            line_width=1.0,\n",
    "                                            title='Pendulum State',\n",
    "                                            x_label='Angle [rad]',\n",
    "                                            y_label='Angular Rate [rad / sec]',\n",
    "                                            x_lim=(-25, 25),\n",
    "                                            y_lim=(-25, 25),\n",
    "                                            h_zero_line=True,\n",
    "                                            v_zero_line=True,)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33985970-c3cc-413a-b068-a78eeb262844",
   "metadata": {},
   "source": [
    "Note the animator does not appear yet. This is because the GUI is not started yet. There are two ways to start the GUI:\n",
    "1. `proj.reset`\n",
    "2. `proj.animator.start`\n",
    "\n",
    "Either is valid. In this case, we will just use the `reset` function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a04d5a02-6da4-4542-8da1-74c874fa0470",
   "metadata": {},
   "source": [
    "## Running the Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb722940-b23c-47ec-9013-1871c7fa2172",
   "metadata": {},
   "source": [
    "Before running the simulation, we reset the project. This ensure that everything is started, updated, and in the desired initial state. `proj.reset` should be called before a simulation every time you run one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50902142-1f44-40f5-a671-2f8cad02561c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset the project to its initial state. This is required to\n",
    "# reset the simulation, reset the visualizer, and reset/start the\n",
    "# animator.\n",
    "proj.reset() # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9e2d466-5008-451e-89c5-2216ecbd675d",
   "metadata": {},
   "source": [
    "**Important Note**: After `reset` is called, the animator GUI will open. It will be unresponsive until it is refreshed by `proj`. This refreshes occur automatically at every `proj.reset`, `proj.step`, `proj.await_anykeys`, or `proj.await_keypress` call. Therefore, while the simulation loop is running, or while waiting for user keyboard input, the GUI will be responsive. You can also manually refresh it using the `proj.refresh_animator` function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1626abcf-d925-4c9c-875f-c24341bfa7d0",
   "metadata": {},
   "source": [
    "Next we create and run a simulation loop. In this loop, on every step we\n",
    "1. Plot the pendulum's state in phase space\n",
    "4. Take a simulation step in real time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1cfe1f9-a181-4c27-89e4-a7463aa59223",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run a 20 second simulation loop\n",
    "while proj.simtime <= 20.:\n",
    "    # Get the state of each joint\n",
    "    state_1 = dp.joints['base_to_arm_1'].state\n",
    "    angle_1 = state_1.angle\n",
    "    omega_1 = state_1.omega\n",
    "    state_2 = dp.joints['mass_1_to_arm_2'].state\n",
    "    angle_2 = state_2.angle\n",
    "    omega_2 = state_2.omega\n",
    "    \n",
    "    # Plot the state\n",
    "    proj.animator.lineplot_append_point(arm_1, angle_1, omega_1) # This returns 0 on success\n",
    "    proj.animator.lineplot_append_point(arm_2, angle_2, omega_2) # This returns 0 on success\n",
    "\n",
    "    # Attempt to run in real time. The small step size may result in\n",
    "    # slower than real time simulation\n",
    "    proj.step(real_time=True, stable_step=False) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a650f1e-a228-45d4-a7fe-cf4aa10e43f1",
   "metadata": {},
   "source": [
    "Finally, we terminate the project. This is required to save any videos that are recorded and gracefully exit all the children threads (including the animator window). `proj.terminate()` should be called when done with any member of the `condynsate.Project` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef6e1f9e-3a5e-4cb9-bdcf-dd151129397a",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj.terminate() # This returns 0 on success"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68cee8e-b623-4439-bd80-6f68756bfc52",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.14.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}