{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7681195e-ca60-4242-a300-d5f85af30866",
   "metadata": {},
   "source": [
    "# Animator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8157195-cd68-4472-b896-0ce9d87c8ee2",
   "metadata": {},
   "source": [
    "## Setting up the Simulation Environment"
   ]
  },
  {
   "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": 1,
   "id": "a93a1982-107a-4a85-a825-99d1cbb7d443",
   "metadata": {},
   "outputs": [],
   "source": [
    "from condynsate import Simulator as conSim\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.simulator` that does not use the keyboard but does use the animator and the visualizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "11c00a7e-876d-4285-8a8c-c7ce62de283a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You can open the visualizer by visiting the following URL:\n",
      "http://127.0.0.1:7000/static/\n"
     ]
    }
   ],
   "source": [
    "simulator = conSim(keyboard = False,\n",
    "                   visualization = True,\n",
    "                   animation = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63405e52-239a-47e2-9b07-c7180600eed6",
   "metadata": {},
   "source": [
    "Load three large planes into the simulator to represent the ground and walls. Set `fixed = True` so that they are fixed in space, and set `update_vis = False`. Because the ground and walls will not move, there is no need to send updates to the visualizer. This saves resources. Further, adjust the position and rotation of the walls so they look good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "036c9a01-2b8c-43ca-9891-a560677d6805",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the ground\n",
    "ground = simulator.load_urdf(urdf_path = assets['plane_big'],\n",
    "                             fixed = True,\n",
    "                             update_vis = False)\n",
    "\n",
    "# Load the walls\n",
    "left_wall = simulator.load_urdf(urdf_path = assets['plane_big'],\n",
    "                                tex_path = assets['concrete_img'], # Choose a default texture for appearance\n",
    "                                position = [0., -5., 0.],\n",
    "                                roll = -1.5707,\n",
    "                                fixed = True,\n",
    "                                update_vis = False)\n",
    "right_wall = simulator.load_urdf(urdf_path = assets['plane_big'],\n",
    "                                 tex_path = assets['concrete_img'],  # Choose a default texture for appearance\n",
    "                                 position = [0., 5., 0.],\n",
    "                                 roll = 1.5707,\n",
    "                                 fixed = True,\n",
    "                                 update_vis = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59e2aa59-3a1a-42d0-afc2-5801ea5e6997",
   "metadata": {},
   "source": [
    "Load the cart into the simulator. Set `fixed = False` so that the cart is able to freely move. Set `update_vis = True` so that its position updates are reflected in the visualizer. Additionally, set the position to `[0.0, 0.0, 0.25]` so that its wheels rest on the ground plane, and set yaw to 90 degrees so that is moves in the Y plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a1c75c92-214f-4dd1-933c-1d67c50eaf5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "cart = simulator.load_urdf(urdf_path = assets['cart'],\n",
    "                           position = [0., 0., 0.25],\n",
    "                           yaw = 1.5707,\n",
    "                           fixed = False,\n",
    "                           update_vis = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "413ad256-1d76-47e5-a841-2fd6bc8c4a5c",
   "metadata": {},
   "source": [
    "Set the each wheel joint's and the pendulum arm joint's friction to some small value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2b261027-9578-415f-b6e5-ac289212c94d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set joint damping\n",
    "joint_names = ('chassis_to_wheel_1',\n",
    "               'chassis_to_wheel_2',\n",
    "               'chassis_to_wheel_3',\n",
    "               'chassis_to_wheel_4',\n",
    "               'chassis_to_arm',)\n",
    "for joint_name in joint_names:\n",
    "    simulator.set_joint_damping(urdf_obj = cart,\n",
    "                               joint_name = joint_name,\n",
    "                               damping = 0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d806decf-6831-4af2-973e-954e2a1aac94",
   "metadata": {},
   "source": [
    "Set the initial angle of the pendulum. 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ad5f4fa2-b267-44a3-a4ee-f76a37f8316f",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_angle = 0.1745\n",
    "simulator.set_joint_position(urdf_obj = cart,\n",
    "                             joint_name = 'chassis_to_arm',\n",
    "                             position = initial_angle, # For a rotational joint, this value is in radians\n",
    "                             initial_cond = True,\n",
    "                             physics = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60a8314c-e9de-4fb7-aac7-d7e1978a5e3d",
   "metadata": {},
   "source": [
    "## Setting up the Animator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cce92419-457c-4510-a6cd-2928c29c4723",
   "metadata": {},
   "source": [
    "Because we set `animator = True` in the simulator initialization, the animator is ready to add new plots to it. To do this we use the `condynsate.Simulator.add_plot` function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86f2e9b2-bec4-4bab-bd86-daebc64455b1",
   "metadata": {},
   "source": [
    "In the first plot, we want to show the mean wheel torque as a function of time. To do so, we create a simple line plot. This is done with the argument `plot_type = 'line'`. Further we add a title and axis labels with the `title = \"Torque vs Time\"`, `x_label = \"Time [Seconds]\"`, and `y_label = \"Torque [Nm]\"`, respectively. Next, to format the line itself, we set its line width to 2.5 with `line_width = 2.5`. Finally, we limit the y-axis view between -0.8 and 0.8 with `y_lim = [-0.80, 0.80]` and add a thin horizontal line across y=0 with `h_zero_line = True`. The return value of `condynsate.Simulator.add_plot` is the identifier of the plot just added. It will be used to later to tell the animator which plot to update when adding data points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "79952ec7-62cb-4f75-b407-a369deb73180",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_1 = simulator.add_plot(plot_type = 'line',\n",
    "                            title = \"Torque vs Time\",\n",
    "                            x_label = \"Time [Seconds]\",\n",
    "                            y_label = \"Torque [Nm]\",\n",
    "                            line_width = 2.5,\n",
    "                            y_lim = [-0.80, 0.80],\n",
    "                            h_zero_line = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57150eac-032f-43d5-b5fa-e46612629b10",
   "metadata": {},
   "source": [
    "Now we want another line plot that shows both the pendulum angle and mean wheel angle simultaneously. Because we want two lines on this plot, we call for 2 artists, one for each line. All other arguments are similar aside for `labels = [\"Pendulum\", \"Wheel\"]` which tells the animator what each arist should be called for the legend. Further, when multiple artists are used `color`, `line_width`, and `line_style` may either be a single argument (applied to all artists) or a list of arguemnts with length equal to number of artists. When `n_artists > 1`, the return value of `condynsate.Simulator.add_plot` is the identifier of the plot just added and a tuple of artist identifiers (one for each artist). These will also be used to later to tell the animator which plot adn artist to update when adding data points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c3e3f718-d919-4c26-b559-6eaa7fc205d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_2, plot_2_artists = simulator.add_plot(n_artists = 2,\n",
    "                                            plot_type = 'line',\n",
    "                                            title = \"Angles vs Time\",\n",
    "                                            x_label = \"Time [Seconds]\",\n",
    "                                            y_label = \"Angles [Rad]\",\n",
    "                                            color = [\"m\", \"c\"], # One color for each artist\n",
    "                                            line_width = 2.5,   # Single value to give all artists same line_width\n",
    "                                            line_style = \"-\",   # Single value to give all artists same line_style\n",
    "                                            label = [\"Pendulum\", \"Wheel\"], # One label for each artist\n",
    "                                            h_zero_line = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "426a9436-b21e-4976-9f55-d0a8bf2844c0",
   "metadata": {},
   "source": [
    "Finally, we would like a bar plot that shows each wheel's deviation from the mean wheel angle. To do so, we add a new plot of type `'bar'` with 4 artists. This time, we also add a thin verticle line along x=0 so we can see if the wheel deviation is leading or lagging the mean. Also, we fix the x axis limits so it is easier to read."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "06f8b9ab-c3e2-415b-8f32-0344178e1971",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_3, plot_3_artists = simulator.add_plot(n_artists = 4,\n",
    "                                            plot_type = 'bar',\n",
    "                                            x_lim = [-0.4, 0.4],\n",
    "                                            title = \"Wheel Deviation from Mean\",\n",
    "                                            x_label = \"Deviation [Rad]\",\n",
    "                                            y_label = \"Wheel Number\",\n",
    "                                            color = \"b\",\n",
    "                                            label = [\"W1\", \"W2\", \"W3\", \"W4\"],\n",
    "                                            v_zero_line = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22193edc-584b-4d63-a081-4fa6f85d0cd2",
   "metadata": {},
   "source": [
    "The plots are now set up how we want, so we can open the animator GUI. Once opened, no additional plots can be added. Note: `%%capture` is required to supress the animator drawing the first frame to this notebook. Though it is not required, it is reccomended. It must be the first line in the cell in which it's called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b1cb404a-fc5c-4e02-8075-f525128bad1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture \n",
    "simulator.start_animator()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81accdb1-5e09-4f2b-950b-0d5af934510e",
   "metadata": {},
   "source": [
    "As expected, a figure with our three plots opens."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba01c0c0-bc8a-453b-8bd7-b187c28e6fd4",
   "metadata": {},
   "source": [
    "## Running the Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4706a656-382b-4757-bb50-dc0dba8c3a27",
   "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 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 four things:\n",
    "1. Take a simulation time step.\n",
    "2. If a successful simulation step was taken, read the wheel angles, wheel angular velocities, pendulum angle, and pendulum angular velocity. Successful simulation steps are indicated by `condynsate.Simulator.step` having a `ret_code` greater than `0`.\n",
    "3. If a successful simulation step was taken, apply torques to wheels based on the mean wheel angle, mean wheel angular velocity, pendulum angle, and pendulum angular velocity.\n",
    "4. If a successful simulation step was taken, update the animator."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b029ec-66f5-453f-bc30-4baf7e767863",
   "metadata": {},
   "source": [
    "Let's start by making a function that reads each angle and angular velocity we want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1ca5591c-e575-4d6f-b6c5-fc46f37ae987",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_states(simulator, cart):\n",
    "    # Define which joints to read\n",
    "    joint_names = ('chassis_to_wheel_1',\n",
    "                   'chassis_to_wheel_2',\n",
    "                   'chassis_to_wheel_3',\n",
    "                   'chassis_to_wheel_4',\n",
    "                   'chassis_to_arm',)\n",
    "\n",
    "    # Loop over each joint and read the state extracting angle and angular vel.\n",
    "    angles = []\n",
    "    angular_vels = []\n",
    "    for joint_name in joint_names:\n",
    "        state = simulator.get_joint_state(urdf_obj = cart,\n",
    "                                          joint_name = joint_name)\n",
    "        angles.append(state['position'])\n",
    "        angular_vels.append(state['velocity'])\n",
    "    \n",
    "    # For the wheels, return the mean angle and angular vel.\n",
    "    mean_wheel_angle = np.mean(angles[:-1])\n",
    "    mean_wheel_ang_vel = np.mean(angular_vels[:-1])\n",
    "\n",
    "    # Additionally, calculate the wheel deviation from mean.\n",
    "    wheel_deviations = []\n",
    "    for wheel_angle in angles[:-1]:\n",
    "        wheel_deviations.append(wheel_angle - mean_wheel_angle)\n",
    "    \n",
    "    return (mean_wheel_angle, mean_wheel_ang_vel), (angles[-1], angular_vels[-1]), wheel_deviations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "118b4204-3a9a-4881-a49f-97539d2cf3eb",
   "metadata": {},
   "source": [
    "Next we make a function that applies torque based on the state we read."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "bb95a144-0f24-4971-b633-f1bcb9385e62",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_torque(simulator, cart, mean_wheel_state, pendulum_state):\n",
    "    # Calculate the torque via magic\n",
    "    K = np.array([[ -2. ,  -0.1, -10. ,  -0.1]])\n",
    "    m_e = np.zeros((4,1))\n",
    "    n_e = np.zeros((1,1))\n",
    "    m = np.array([[pendulum_state[1]],\n",
    "                  [mean_wheel_state[1]],\n",
    "                  [pendulum_state[0]],\n",
    "                  [mean_wheel_state[0]]])\n",
    "    torque = (-K@(m - m_e) + n_e).flatten()[0]\n",
    "    torque = np.clip(torque, -0.75, 0.75)\n",
    "    \n",
    "    # Define which joints we apply torque to\n",
    "    joint_names = ('chassis_to_wheel_1',\n",
    "                   'chassis_to_wheel_2',\n",
    "                   'chassis_to_wheel_3',\n",
    "                   'chassis_to_wheel_4',)\n",
    "\n",
    "    # Apply the torque we calculated\n",
    "    for joint_name in joint_names:\n",
    "        simulator.set_joint_torque(urdf_obj = cart,\n",
    "                                   joint_name = joint_name,\n",
    "                                   torque = torque,\n",
    "                                   show_arrow = True, # Define arrow params to visualize the torque\n",
    "                                   arrow_scale = 0.25,\n",
    "                                   arrow_offset = 0.025)\n",
    "\n",
    "    return torque"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07d0b23a-255b-41d7-af95-b62850311df5",
   "metadata": {},
   "source": [
    "Next we want to build a function that updates the first line plot. We want this function to add a single datum to the end of the data list being plotted as a line. To do this, we will use the `condynsate.Simulator.add_line_datum` function. The arguments to this function are the plot to which the point is added and the (x,y) coordinates of the point. The x coordinate will be simulation time which we extract with `condynsate.Simulator.time`, and the y coordinate is the torque we calculated in the previous function. After `condynsate.Simulator.add_line_datum` is called, the animator automatically updates to add the point at the specified frame rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "69bf2caa-d1dc-4d3f-96bd-fdcf1dd4228e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_torque_point(simulator, plot_id, torque):\n",
    "    simulator.add_line_datum(plot_id = plot_id,\n",
    "                             x = simulator.time,\n",
    "                             y = torque)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b98bac29-c3c3-4f40-8d28-2771d4608068",
   "metadata": {},
   "source": [
    "We build a similar function for the wheel and pendulum angles, but this time the plot has 2 artists. Therefore, we need to specify to which artist we are adding the data points. This is done with the `artist_id` argument. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a9becff9-fa2c-4edb-82fa-7fb6c90ff815",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_state_points(simulator, plot, plot_artists, mean_wheel_state, pendulum_state):\n",
    "    simulator.add_line_datum(plot_id = plot,\n",
    "                             artist_id = plot_artists[0], # Draw to the first artist (whose label is \"Pendulum\")\n",
    "                             x = simulator.time,\n",
    "                             y = pendulum_state[0])\n",
    "    simulator.add_line_datum(plot_id = plot,\n",
    "                             artist_id = plot_artists[1], # Draw to the second artist (whose label is \"Wheel\")\n",
    "                             x = simulator.time,           \n",
    "                             y = mean_wheel_state[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5ac5b3d-5117-44b4-9cae-2deef5f8dab7",
   "metadata": {},
   "source": [
    "Now we want a function to edit the bar chart. This is very similar to `add_state_points`, however we use the function 'condynsate.Simulator.set_bar_value' to edit a bar chart's values. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "64e9e5e5-cf10-424b-86c3-c83a4ab1fbc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_deviation_values(simulator, plot, plot_artists, wheel_deviations):\n",
    "    for plot_artist, wheel_deviation in zip(plot_artists, wheel_deviations):\n",
    "        simulator.set_bar_value(plot_id = plot,\n",
    "                                artist_id = plot_artist,\n",
    "                                value = wheel_deviation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ae80d89-d125-4182-8b77-2b54c4f70301",
   "metadata": {},
   "source": [
    "Finally, we build our simulation loop. By setting the `real_time = True` flag and `max_time = 10.0` in `condynsate.Simulator.step`, we tell the simulation to attempt to run in real time for 5 seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6fdf5bd7-4088-40f9-91dc-f72733c36019",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset before running a simulation loop\n",
    "simulator.reset()\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 = 10.0)\n",
    "\n",
    "    # If successful step was taken\n",
    "    if ret_code > 0:\n",
    "        # Read the states and apply the torque\n",
    "        mean_wheel_state, pendulum_state, wheel_deviations = read_states(simulator, cart)\n",
    "        torque = apply_torque(simulator, cart, mean_wheel_state, pendulum_state)\n",
    "\n",
    "        # Add the new data points to the animator\n",
    "        add_torque_point(simulator, plot_1, torque)\n",
    "        add_state_points(simulator, plot_2, plot_2_artists, mean_wheel_state, pendulum_state)\n",
    "        set_deviation_values(simulator, plot_3, plot_3_artists, wheel_deviations)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f4d704c-5612-4570-85c3-bbd13865e170",
   "metadata": {},
   "source": [
    "Watch as the pendulum magically balances while the torque and angles are plotted in real time on the animator GUI."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "952e5d13-f04f-41f2-bbd4-1b4e2bb248d2",
   "metadata": {},
   "source": [
    "Once the simulation is done, we terminate the animator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7ef8de02-1feb-4767-ac92-51bc83eb0b7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulator.terminate_animator()"
   ]
  }
 ],
 "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
}