{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Example 03: Simple Robot Arm Kinematics\n",
    "\n",
    "Modeling a simple 3-joint robot arm using hierarchical coordinate frames. This example demonstrates:\n",
    "\n",
    "1. **Complex frame hierarchies**: Multi-level parent-child relationships (base → joint1 → joint2 → joint3 → hand)\n",
    "2. **Animation**: Updating frames dynamically to simulate robot motion\n",
    "3. **Data structures with frames**: Building a `RobotArm` class that manages multiple coordinate frames"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from hazy import Frame"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Build the robot arm hierarchy\n",
    "\n",
    "A robot arm is naturally modeled as a hierarchy of frames:\n",
    "- **Base**: The fixed mounting point (on the ground)\n",
    "- **Joints**: Rotation points that move relative to their parent\n",
    "- **Hand**: The end effector attached to the last joint\n",
    "\n",
    "Each joint is a child of the previous one, so rotating a joint affects all joints further down the chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "root = Frame.make_root(\"root\")\n",
    "\n",
    "robot_base = root.make_child(\"base\")\n",
    "first_joint = robot_base.make_child(\"joint1\").translate(z=1.0)\n",
    "second_joint = first_joint.make_child(\"joint2\").translate(x=1.0)\n",
    "third_joint = second_joint.make_child(\"joint3\").translate(x=1.0)\n",
    "robot_hand = third_joint.make_child(\"hand\").translate(z=-1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "**Transformation order**: `hazy` uses `S->R->T` (scale-rotate-translate) order. Rotations happen before translations, which is essential for joint behavior - joints rotate around their pivot points, not around translated positions.\n",
    "\n",
    "**The hierarchy:** Our robot arm has the following layout:\n",
    "```\n",
    "root → base → joint1 (+1 in z) → joint2 (+1 in x) → joint3 (+1 in x) → hand (-1 in z)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Create a RobotArm data structure\n",
    "\n",
    "To manage the robot's joints and visualize them, we create a `RobotArm` class that stores all joint frames and provides plotting functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class RobotArm:\n",
    "    joints: list[Frame]\n",
    "\n",
    "    def plot_robot(self, axes, color=\"black\", alpha=1.0, show_gismo: bool = True):\n",
    "        \"\"\"Plot the robot arm by connecting joint origins.\"\"\"\n",
    "        origins = np.array([joint.origin_global for joint in self.joints])\n",
    "        axes.plot(*origins.T, marker=\"o\", color=color, alpha=alpha, ms=10, lw=3)\n",
    "\n",
    "        if show_gismo:\n",
    "            for joint in self.joints:\n",
    "                plot_gismo(joint, axes, alpha=alpha)\n",
    "        return axes\n",
    "\n",
    "    def __getitem__(self, index: int) -> Frame:\n",
    "        return self.joints[index]\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.joints)\n",
    "\n",
    "\n",
    "def plot_gismo(frame: Frame, axes, scale: float = 0.2, alpha: float = 1.0):\n",
    "    \"\"\"Visualize coordinate system axes (x=blue, y=green, z=red) at frame origin.\"\"\"\n",
    "    x_axis = np.vstack(\n",
    "        [frame.origin_global, frame.origin_global + frame.x_axis_global * scale]\n",
    "    )\n",
    "    y_axis = np.vstack(\n",
    "        [frame.origin_global, frame.origin_global + frame.y_axis_global * scale]\n",
    "    )\n",
    "    z_axis = np.vstack(\n",
    "        [frame.origin_global, frame.origin_global + frame.z_axis_global * scale]\n",
    "    )\n",
    "\n",
    "    axes.plot(*x_axis.T, c=\"blue\", alpha=alpha)\n",
    "    axes.plot(*y_axis.T, c=\"green\", alpha=alpha)\n",
    "    axes.plot(*z_axis.T, c=\"red\", alpha=alpha)\n",
    "    return axes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "The `RobotArm` class:\n",
    "- Stores joint frames in a list\n",
    "- Provides `plot_robot()` to visualize the arm by connecting joint origins\n",
    "- Supports indexing (`robot[i]`) to access individual joints\n",
    "- Uses `joint.origin_global` to get the global position of each joint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Build the robot instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "robot = RobotArm(\n",
    "    joints=[robot_base, first_joint, second_joint, third_joint, robot_hand]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Visualize the initial configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 8))\n",
    "ax3d = fig.add_subplot(projection=\"3d\")\n",
    "\n",
    "ax3d.set_xlabel(\"x\", fontsize=10)\n",
    "ax3d.set_ylabel(\"y\", fontsize=10)\n",
    "ax3d.set_zlabel(\"z\", fontsize=10)\n",
    "ax3d.set_xlim(-2, 2)\n",
    "ax3d.set_ylim(-2, 2)\n",
    "ax3d.set_zlim(0, 4)\n",
    "\n",
    "robot.plot_robot(ax3d)\n",
    "\n",
    "ax3d.set_title(\"Simple Robot Arm\", fontsize=14, fontweight=\"bold\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Animate the robot arm\n",
    "\n",
    "Now we'll create an animation by rotating different joints over time. We'll:\n",
    "1. Rotate the base (rotates entire arm)\n",
    "2. Rotate joint3 alternately up and down\n",
    "\n",
    "Each configuration is drawn in a different color to show the motion sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = 8\n",
    "angle_per_step = 360 / steps\n",
    "\n",
    "transforms = []\n",
    "\n",
    "# move first joint up\n",
    "transforms.append([(1, {\"y\": -45, \"degrees\": True})])\n",
    "\n",
    "for i in range(1, steps):\n",
    "    swing_direction = 1 - 2 * (i % 2)  # alternates between +1 and -1\n",
    "    transforms.append(\n",
    "        [\n",
    "            (3, {\"y\": swing_direction * 45, \"degrees\": True}),\n",
    "            (0, {\"z\": angle_per_step, \"degrees\": True}),\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "Each transform is a list of `(joint_index, rotation_params)` tuples:\n",
    "- Index 0 = base, 1 = joint1, 3 = joint3\n",
    "- First step: rotate joint1 down by 45°\n",
    "- Subsequent steps: swing joint3 up/down, rotate base around z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "## Apply transformations and visualize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 10))\n",
    "ax3d = fig.add_subplot(projection=\"3d\")\n",
    "\n",
    "ax3d.set_xlabel(\"X\", fontsize=10)\n",
    "ax3d.set_ylabel(\"Y\", fontsize=10)\n",
    "ax3d.set_zlabel(\"Z\", fontsize=10)\n",
    "ax3d.set_xlim(-2, 2)\n",
    "ax3d.set_ylim(-2, 2)\n",
    "ax3d.set_zlim(0, 4)\n",
    "\n",
    "colors = plt.cm.viridis(np.linspace(0, 1, len(transforms)))\n",
    "alphas = np.linspace(0.3, 1.0, len(transforms))\n",
    "\n",
    "for color, alpha, transform_step in zip(colors, alphas, transforms):\n",
    "    for joint_index, rotation_kwargs in transform_step:\n",
    "        robot[joint_index].rotate_euler(**rotation_kwargs)\n",
    "\n",
    "    robot.plot_robot(ax3d, color=color, alpha=alpha)\n",
    "\n",
    "ax3d.set_title(\"Animated Robot Arm\", fontsize=14, fontweight=\"bold\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "**Key observation**: The robot arm frames were created once at the beginning. We modify them in the loop with `.rotate_euler()`, and `.origin_global` automatically returns the updated positions. No need to recreate the robot or manually track transformations - the frame hierarchy does it for us.\n",
    "\n",
    "## Transforming points through the hierarchy\n",
    "\n",
    "The real power of frame hierarchies emerges when you need to track objects attached to moving parts. Let's attach a tool to the robot hand and track its position in world coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tool tip is 0.5 units below the hand in the hand's local coordinate system\n",
    "tool_tip = robot_hand.point(0, 0, -0.5)\n",
    "\n",
    "print(f\"Tool tip in hand frame: {tool_tip:.2f}\")\n",
    "print(f\"Tool tip in world frame: {tool_tip.to_global():.2f}\")\n",
    "\n",
    "# Rotate the third joint - the tool tip position updates automatically\n",
    "robot[3].rotate_euler(y=90, degrees=True)\n",
    "\n",
    "print(f\"\\nAfter rotating joint3 by 90°:\")\n",
    "print(f\"Tool tip in hand frame: {tool_tip:.2f}\")  # unchanged\n",
    "print(f\"Tool tip in world frame: {tool_tip.to_global():.2f}\")  # updated"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "**What happened here:**\n",
    "- We defined `tool_tip` once in the hand's local frame\n",
    "- `.to_global()` automatically computed its world position through the entire kinematic chain\n",
    "- After rotating a joint, the same `tool_tip` object gives updated world coordinates\n",
    "- The local coordinates remained unchanged (still `[0, 0, -0.5]` in hand frame)\n",
    "\n",
    "This is the key insight: **define geometry once in the natural reference frame, query it anywhere.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "## Key Takeaways\n",
    "\n",
    "**Multi-level hierarchies**: Robot arms, camera rigs, and mechanical assemblies naturally map to frame hierarchies. Each joint/component is a child of its parent, creating a kinematic chain.\n",
    "\n",
    "**Persistent frames with dynamic updates**: Create frames once, modify them later with `.rotate_euler()`, `.translate()`, etc. All primitives and child frames automatically reflect the changes through cache invalidation.\n",
    "\n",
    "**Data structures with frames**: You can build higher-level abstractions (like `RobotArm`) that store and manage multiple frames. The frames remain live - calling `.origin_global` always gives current positions.\n",
    "\n",
    "**Automatic propagation**: Define points and vectors in their natural reference frame, then query them in any other frame with `.to_global()` or `.to_frame()`. The transformation chain is computed automatically.\n",
    "\n",
    "**Without hazy**: Track transformation matrices for each joint, manually compose them in the right order, update dependent transformations when a parent changes, manually transform tool positions through the chain. For a 5-joint arm with a tool, that's matrix multiplication through 6+ transforms for every query.\n",
    "\n",
    "**With hazy**: Define the hierarchy once, modify individual joints freely, query positions anywhere. The library handles composition and propagation automatically."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hazy-frames",
   "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
}