{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a7ce0132-dab5-4307-ac0e-babed341f452",
   "metadata": {},
   "source": [
    "© Copyright, 2025 G. Schaer.\n",
    "\n",
    "SPDX-License-Identifier: GPL-3.0-only"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f75d8d6c-b2e4-4235-8fd9-70c5ff4bfdf5",
   "metadata": {},
   "source": [
    "# Project Example 1: The Cart"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9de80984-6f68-4a85-af3b-8f9a6cb7066b",
   "metadata": {},
   "source": [
    "Import necessary `condynsate` modules. Also import `numpy` for array management. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a93a1982-107a-4a85-a825-99d1cbb7d443",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from condynsate import Project\n",
    "from condynsate import __assets__ as assets\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b9129d1-883b-41f2-8883-91e3aae37ef8",
   "metadata": {},
   "source": [
    "Create a instance of `condynsate.Project` that does not use the keyboard but does use the animator and the visualizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11c00a7e-876d-4285-8a8c-c7ce62de283a",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj = Project(keyboard = False,\n",
    "               visualizer = True,\n",
    "               animator = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63405e52-239a-47e2-9b07-c7180600eed6",
   "metadata": {},
   "source": [
    "Load four medium planes into the simulator to represent the ground and walls. Set `fixed = True` so that they are fixed in space and adjust their initial states so they are oriented properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "036c9a01-2b8c-43ca-9891-a560677d6805",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a plane with a carpet texture for the ground\n",
    "ground = proj.load_urdf(assets['plane_medium.urdf'], fixed=True)\n",
    "ground.links['plane'].set_texture(assets['carpet.png'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6402c2f6-67ab-44bc-ae19-ecedf972935b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and orient a plane with a windowed wall texture for the left wall\n",
    "left_wall = proj.load_urdf(assets['half_plane_medium.urdf'], fixed=True)\n",
    "left_wall.links['plane'].set_texture(assets['window_wall.png'])\n",
    "left_wall.set_initial_state(roll=1.5708, yaw=1.5708, position=(-5,0,2.5)) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef51031e-e703-4532-a66d-28fc60a64ab0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and orient a plane with a doored wall texture for the right wall\n",
    "right_wall = proj.load_urdf(assets['half_plane_medium.urdf'], fixed=True)\n",
    "right_wall.links['plane'].set_texture(assets['door_wall.png'])\n",
    "right_wall.set_initial_state(roll=1.5708, yaw=-1.5708, position=(5,0,2.5)) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aba23e37-7b5f-45e4-a11a-e1683ed8fe01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and orient a plane with a classroom wall texture for the back wall\n",
    "back_wall = proj.load_urdf(assets['half_plane_medium.urdf'], fixed=True)\n",
    "back_wall.links['plane'].set_texture(assets['classroom_wall.png'])\n",
    "back_wall.set_initial_state(roll=1.5708, position=(0,5,2.5)) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59e2aa59-3a1a-42d0-afc2-5801ea5e6997",
   "metadata": {},
   "source": [
    "Load the cart into the simulator. Additionally, set the initial position to `[0.0, 0.0, 0.126]` so that its wheels are just above the ground plane. Finally, set the initial angle of the pendulum to a non zero value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1c75c92-214f-4dd1-933c-1d67c50eaf5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and orient a cart carrying an inverted pendulum. Set the initial\n",
    "# angle of the pendulum to a non-zero angle.\n",
    "cart = proj.load_urdf(assets['cart.urdf'])\n",
    "cart.set_initial_state(position=(0,0,0.126)) # This returns 0 on success\n",
    "cart.joints['chassis_to_arm'].set_initial_state(angle=0.25) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d5c6aee-3ce5-40a2-aeb0-9a79f80417c5",
   "metadata": {},
   "source": [
    "Note that the cart's position and pendulum angle updates are not yet reflected in the visualizer. To update this manually, we can call the `proj.refresh_visualizer` function. Note, however, this is also done automatically every time `proj.load_urdf`, `proj.reset`, or `proj.step` is called. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eca3eb9-c9d0-4412-9cd6-30ee7cc0d6bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj.refresh_visualizer() # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88e05faf-c2c0-4cea-bc9d-952fe4d6d3ee",
   "metadata": {},
   "source": [
    "Now we create some extra local vars for the logistics of this specific project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f55ce7ad-f81e-4280-8dff-8c205b86af54",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store the name of each wheel joint for easy iteration\n",
    "wheel_joint_names = ('chassis_to_wheel_1', 'chassis_to_wheel_2',\n",
    "                     'chassis_to_wheel_3', 'chassis_to_wheel_4',)\n",
    "\n",
    "# Set control constants to keep the pendulum upright\n",
    "k = np.array([[ 0.664, -0.024, 4.0, -0.0064]])\n",
    "m_e = np.zeros(4)\n",
    "n_e = np.zeros(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "588b288a-5117-4967-a339-1d32b8aa3e73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Turn off the axes and grid visualization. Turn on the spotlight\n",
    "proj.visualizer.set_axes(False) # This returns 0 on success\n",
    "proj.visualizer.set_grid(False) # This returns 0 on success\n",
    "proj.visualizer.set_spotlight(on=True) # This returns 0 on success\n",
    "\n",
    "# Focus the camera on the cart\n",
    "proj.visualizer.set_cam_target(cart.center_of_mass) # This returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56ce5e80-b314-4f29-8769-f7ba10e26b83",
   "metadata": {},
   "source": [
    "We will add two lines plots to the animator. Each line plot will have exactly one line one it. These will not appear yet, only after starting the animator GUI will they appear. 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": "code",
   "execution_count": null,
   "id": "50781c20-8163-4aeb-b669-44962c88a628",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a line plot to the animator to track the pendulum angle\n",
    "# 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=1, so the return value is just a hex code.\n",
    "plot1 = proj.animator.add_lineplot(n_lines=1, \n",
    "                                   color='b', \n",
    "                                   line_width=2.5,\n",
    "                                   y_lim=(-30., 30.),\n",
    "                                   title='Pendulum',\n",
    "                                   x_label='Time [seconds]',\n",
    "                                   y_label='Angle [degrees]',\n",
    "                                   h_zero_line=True,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f0cf1e1-ecb0-4b8f-a591-a7497aa96c00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add another line plot to the animator to track the cart x position\n",
    "plot2 = proj.animator.add_lineplot(n_lines=1, \n",
    "                                   color='r', \n",
    "                                   line_width=2.5,\n",
    "                                   y_lim=(-5., 5.),\n",
    "                                   title='Cart',\n",
    "                                   x_label='Time [seconds]',\n",
    "                                   y_label='Position [meters]',\n",
    "                                   h_zero_line=True,)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c3f52c-65f0-4134-8798-1c2a81222f54",
   "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": "396571f6-ccf4-4010-b927-eb8baef058e2",
   "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": "b2c63423-a3fc-490d-9437-0ca9d9aecff4",
   "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": "7cd2ae43-682e-45f0-8500-f86864aced0d",
   "metadata": {},
   "source": [
    "Next we create and run a simulation loop. In this loop, on every step we\n",
    "1. Collect state information about the cart and the pendulum angle\n",
    "2. Note if the pendulum has fallen down or not\n",
    "3. Calculate a torque to apply to the wheels that will keep the pendulum upright\n",
    "4. Apply this calculate torque\n",
    "5. Add the current pendulum angle and cart x position to the associated line plots\n",
    "6. Take a simulation step in real time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "915b6a32-8465-443f-898e-c94ebfe9d07b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run a 10 second simulation loop\n",
    "start = time.time()\n",
    "while proj.simtime <= 10.:\n",
    "    # Read the states of the pendulum and each wheel\n",
    "    pen_state = cart.joints['chassis_to_arm'].state\n",
    "    wheel_states=tuple(cart.joints[n].state for n in wheel_joint_names)\n",
    "\n",
    "    # If the pendulum angle exceeds 90 degrees, a failure condition is\n",
    "    # met. Terminate the simulation loop.\n",
    "    if abs(pen_state.angle) > 1.5708:\n",
    "        print('The pendulum fell.')\n",
    "        break\n",
    "\n",
    "    # Do controls calculations to determine what torque, when applied\n",
    "    # to each wheel, will keep the pendulum upright.\n",
    "    m = np.array([pen_state.omega,\n",
    "                  np.mean([s.omega for s in wheel_states]),\n",
    "                  pen_state.angle,\n",
    "                  np.mean([s.angle for s in wheel_states])])\n",
    "    torque = float(np.clip((-k@(m - m_e) + n_e)[0], -0.75, 0.75))\n",
    "\n",
    "    # Apply the torque we calculated to each wheel\n",
    "    for joint_name in wheel_joint_names:\n",
    "        # This will offset a drawn torque arrow out of the center of\n",
    "        # the wheels so we can see them. It is required to be\n",
    "        # different between the front wheels (1 and 2) and the rear\n",
    "        # wheels (3 and 4) because they are oriented 180 degrees apart\n",
    "        offset = ('3' in joint_name or '4' in joint_name)*0.1-0.05\n",
    "        cart.joints[joint_name].apply_torque(torque,\n",
    "                                             draw_arrow=True,\n",
    "                                             arrow_scale=0.33,\n",
    "                                             arrow_offset=offset)\n",
    "\n",
    "    # Plot the pendulum angle against the current simulation time\n",
    "    angle_deg = pen_state.angle*180./np.pi\n",
    "    proj.animator.lineplot_append_point(plot1, proj.simtime, angle_deg) # This returns 0 on success\n",
    "\n",
    "    # Plot the cart's x position against the current simulation time\n",
    "    cart_xpos = cart.state.position[0]\n",
    "    proj.animator.lineplot_append_point(plot2, proj.simtime, cart_xpos) # This returns 0 on success\n",
    "\n",
    "    # Take a simulation step that attempts real time simulation\n",
    "    proj.step(real_time=True, stable_step=False) # This returns 0 on success\n",
    "\n",
    "# Note how long the simulation took\n",
    "print(f\"Simuation took {time.time()-start:.2f} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18de874a-a52a-4704-8e73-c16e34c9e233",
   "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": "bf7c73ba-f228-4665-810a-fe200aa604f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj.terminate() # This returns 0 on success"
   ]
  }
 ],
 "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
}