{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "000ae622-6c87-45d0-a90d-0b133fdbd67e",
   "metadata": {},
   "source": [
    "# Keyboard Interactivity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c58045bb-4b36-4caf-9cb6-039bec6c63ef",
   "metadata": {},
   "source": [
    "## Setting up the Simulation Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8874740a-5c75-4115-8ad7-75edd8ef7cd6",
   "metadata": {},
   "source": [
    "Import necessary `condynsate` modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cc40d60b-49c6-4c54-8abf-ba7b801cf370",
   "metadata": {},
   "outputs": [],
   "source": [
    "from condynsate import Simulator as conSim\n",
    "from condynsate import __assets__ as assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "088afcf1-b682-4de0-ae59-3788e872a9e1",
   "metadata": {},
   "source": [
    "Create a instance of `condynsate.simulator` that does not use the animator but does use the keyboard and the visualizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "00f52642-97ac-43ca-957d-4102c79cba87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You can open the visualizer by visiting the following URL:\n",
      "http://127.0.0.1:7034/static/\n"
     ]
    }
   ],
   "source": [
    "simulator = conSim(keyboard = True,\n",
    "                   visualization = True,\n",
    "                   animation = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64b530be-8f79-4b96-8fb6-2e3cd4d61eed",
   "metadata": {},
   "source": [
    "Load a large plane into the simulator to represent the ground. Set `fixed = True` so that it is fixed in space, and set `update_vis = False`. Because the ground will not move, there is no need to send updates to the visualizer. This saves resources. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bae9d9be-feca-4c75-a546-ab29f0fa637f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the ground\n",
    "ground = simulator.load_urdf(urdf_path = assets['plane_big'],\n",
    "                             fixed = True,\n",
    "                             update_vis = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ee16f7-0781-4f5b-9ceb-822e469c1278",
   "metadata": {},
   "source": [
    "Load the gyro into the simulator. Set `fixed = True` so that the base link of the gyro is fixed in space. Set `update_vis = True` so that its position updates are reflected in the visualizer. Additionally, set the position to `[0.0, 0.0, 0.05]` so that the base of the pendulum rests on the ground plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c6bff914-6125-404a-921d-5b76c1e9995e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the gyro\n",
    "gyro = simulator.load_urdf(urdf_path = assets['gyro'],\n",
    "                           fixed = True,\n",
    "                           update_vis = True,\n",
    "                           position = [0.,0.,0.05])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46d32ae4-bb94-4541-8aa1-4492fe940103",
   "metadata": {},
   "source": [
    "Set the `\"base_to_outer\"`, `\"outer_to_r2\"`, and `\"r2_to_r1\"` joint frictions to some small value. The `\"base_to_outer\"` joint is the rotational joint between the base of the pendulum and the outer gimbal. The `\"outer_to_r2\"` joint is the rotational joint between the outer gimbal and the second gimbal of the gyro. The `\"r2_to_r1\"` joint is the rotational joint between the second gimbal and the inner gimbal of the gyro."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ec5f1e20-b51e-47b6-a280-05d10be738bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set joint damping\n",
    "simulator.set_joint_damping(urdf_obj = gyro,\n",
    "                           joint_name = 'base_to_outer',\n",
    "                           damping = 0.025)\n",
    "simulator.set_joint_damping(urdf_obj = gyro,\n",
    "                           joint_name = 'outer_to_r2',\n",
    "                           damping = 0.025)\n",
    "simulator.set_joint_damping(urdf_obj = gyro,\n",
    "                           joint_name = 'r2_to_r1',\n",
    "                           damping = 0.025)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be049c2b-ca46-4c1c-8660-18ed540457ea",
   "metadata": {},
   "source": [
    "Set the initial speed of the wheel of the gyro to 10 rad/sec. We denote that this is an initial condition with the `initial_cond = True` flag and we make the change instantaneously by using the `physics = False` flag. Further, we would like to color the wheel based on its speed. We do this by passing three additional arguments:\n",
    "1. `color = True`: tells the simulator to color the child link of the joint (in this case the child link of `\"'r1_to_core'\"` is the wheel) based on a minimum speed, maximum speed, and current speed. At minimum speed, the child link will be blue, at maximum speed, it will be red, and half way between the minimum and maximum it will be white.\n",
    "2. `min_vel = wheel_min_vel`: tells the simulator what the minimum allowed velocity is. If the `velocity` argument is less than `min_vel`, `velocity` will be clipped.\n",
    "3. `max_vel = wheel_max_vel`: tells the simulator what the minimum allowed velocity is. If the `velocity` argument is greater than `max_vel`, `velocity` will be clipped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "465142a0-5671-4b14-a2a7-88aa819f3538",
   "metadata": {},
   "outputs": [],
   "source": [
    "wheel_min_vel = -10.0\n",
    "wheel_max_vel = 10.0\n",
    "simulator.set_joint_velocity(urdf_obj = gyro,\n",
    "                            joint_name = 'r1_to_core',\n",
    "                            velocity = 5.0,\n",
    "                            initial_cond = True,\n",
    "                            physics = False,\n",
    "                            color = True,\n",
    "                            min_vel = wheel_min_vel,\n",
    "                            max_vel = wheel_max_vel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08b2f973-4e0b-40d0-a437-bedab2ae8c89",
   "metadata": {},
   "source": [
    "## Enabling Keyboard Interactivity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9f120e3-1287-44cf-9282-921cb2528b14",
   "metadata": {},
   "source": [
    "In this simulation, we want the user to be able to apply torques to each gimbal joint by using keypresses. To enable this, we will create a function that listens for keypresses, converts them to gimbal joint torque commands, and returns these values for the simulator to apply."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8b62f0a9-0274-4485-9fc5-70c1d91aaac8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_torque_keypresses(simulator):\n",
    "    # Determine what torques to apply based on key presses\n",
    "    tau0 = 0.0 # Torque applied to outer gimbal\n",
    "    tau2 = 0.0 # Torque applied to second gimbal\n",
    "    tau1 = 0.0 # Torque applied to inner gimbal\n",
    "    if simulator.is_pressed('q'): tau0 -= 0.1\n",
    "    if simulator.is_pressed('e'): tau0 += 0.1\n",
    "    if simulator.is_pressed('w'): tau1 -= 0.1\n",
    "    if simulator.is_pressed('s'): tau1 += 0.1\n",
    "    if simulator.is_pressed('a'): tau2 -= 0.1\n",
    "    if simulator.is_pressed('d'): tau2 += 0.1\n",
    "    return (tau0, tau1, tau2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a4a1768-53ae-4cc9-91c9-d368f5de8eac",
   "metadata": {},
   "source": [
    "Next we create a function that applies these torques. This is done with the `condynsate.Simulator.set_joint_torque` function. Further, we would like to visualize these torque inputs. We do this by setting the `show_arrow = True` flag and adjusting the `arrow_offset` and `arrow_scale` until the drawn arrows look like we want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a1ea710f-e874-4fe0-9433-78acbb2bdffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_torques(simulator, gyro, torques):\n",
    "    simulator.set_joint_torque(urdf_obj = gyro,\n",
    "                               joint_name = \"base_to_outer\",\n",
    "                               torque = torques[0],\n",
    "                               show_arrow = True,\n",
    "                               arrow_offset = -1.3,\n",
    "                               arrow_scale = 8.)\n",
    "    simulator.set_joint_torque(urdf_obj = gyro,\n",
    "                               joint_name = \"outer_to_r2\",\n",
    "                               torque = torques[1],\n",
    "                               show_arrow = True,\n",
    "                               arrow_offset = 1.1,\n",
    "                               arrow_scale = 8.)\n",
    "    simulator.set_joint_torque(urdf_obj = gyro,\n",
    "                               joint_name = \"r2_to_r1\",\n",
    "                               torque = torques[2],\n",
    "                               show_arrow = True,\n",
    "                               arrow_offset = -0.85,\n",
    "                               arrow_scale = 8.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b5bbad-0e14-4d3b-9050-8a106b6804e6",
   "metadata": {},
   "source": [
    "Finally, we would also like to allow the user the adjust the wheel speed based on keyboard inputs as well. We do this in a similar way. First we listen for keypresses and convert them to a wheel speed iteration value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f0121a5d-9740-458c-b86d-65ba64215e2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_wheel_vel_keypresses(simulator):\n",
    "    iter_val = 0.0\n",
    "    if simulator.is_pressed('f'): iter_val -= 0.1\n",
    "    if simulator.is_pressed('r'): iter_val += 0.1\n",
    "    return iter_val"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c79f9ff8-be87-4ceb-99a5-e72f16d66731",
   "metadata": {},
   "source": [
    "Then we apply the wheel speed iteration by reading the current wheel speed and then iterating its value, remembering to upate the color of the wheel as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b1be6864-c7e7-4296-b3d3-eaff1781fe7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_wheel_vel(simulator, gyro, iter_val, wheel_min_vel, wheel_max_vel):\n",
    "    wheel_state = simulator.get_joint_state(urdf_obj = gyro,\n",
    "                                            joint_name = 'r1_to_core')\n",
    "    wheel_vel = wheel_state['velocity']\n",
    "    simulator.set_joint_velocity(urdf_obj = gyro,\n",
    "                                 joint_name = 'r1_to_core',\n",
    "                                 velocity = wheel_vel + iter_val,\n",
    "                                 physics = False,\n",
    "                                 color = True,\n",
    "                                 min_vel = wheel_min_vel,\n",
    "                                 max_vel = wheel_max_vel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a04d5a02-6da4-4542-8da1-74c874fa0470",
   "metadata": {},
   "source": [
    "## Running the Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c7be63-a111-4c74-8523-ebee0a44afc1",
   "metadata": {},
   "source": [
    "Run the simulation in a loop. Begin by resetting the simulation using `condynsate.Simulator.reset` to reset the simulation to an initial state. Then use `condynsate.Simulator.await_keypress` to suspend further code execution until the user presses the desired key. Finally we create a simulation loop. The loop terminates when the flag `condynsate.Simulator.is_done` flips to True. In each time step of the simulation loop we do three things:\n",
    "1. Take a simulation time step.\n",
    "2. Apply torques to the `\"base_to_outer\"`, `\"outer_to_r2\"`, and `\"r2_to_r1\"` joint based on keyboard inputs if a successful simulation step is taken. Successful simulation steps are indicated by `condynsate.Simulator.step` having a `ret_code` greater than `0`.\n",
    "3. Apply wheel velocity iterations to `\"r1_to_core\"` based on keyboard inputs if a successful simulation step is taken.\n",
    "\n",
    "By setting the `real_time = True` flag and `max_time = None` in `condynsate.Simulator.step`, we tell the simulation to attempt to run in real time until the user presses the esc key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d1cfe1f9-a181-4c27-89e4-a7463aa59223",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PRESS ENTER TO START SIMULATION.\n",
      "PRESS ESC TO QUIT.\n",
      "PRESS SPACE TO PAUSE/RESUME SIMULATION.\n",
      "PRESS BACKSPACE TO RESET SIMULATION.\n",
      "CONTINUING...\n",
      "QUITTING...\n",
      "Termination command detected. Terminating keyboard listener. Goodbye\n"
     ]
    }
   ],
   "source": [
    "# Reset before running a simulation loop\n",
    "simulator.reset()\n",
    "\n",
    "# Suspend execution until user presses enter key\n",
    "simulator.await_keypress(key='enter')\n",
    "    \n",
    "# Run the simulation loop\n",
    "while(not simulator.is_done):\n",
    "    # Step the sim\n",
    "    ret_code = simulator.step(real_time = True,\n",
    "                              max_time = None) # Max time set to None means the sim runs until user presses ESC.\n",
    "\n",
    "    # If successful step was taken\n",
    "    if ret_code > 0:\n",
    "        # Apply torques based on key presses\n",
    "        torques = read_torque_keypresses(simulator)\n",
    "        apply_torques(simulator, gyro, torques)\n",
    "    \n",
    "        # Iterate wheel speed based on keypresses\n",
    "        iter_val = read_wheel_vel_keypresses(simulator)\n",
    "        apply_wheel_vel(simulator, gyro, iter_val, wheel_min_vel, wheel_max_vel)"
   ]
  }
 ],
 "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.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}