{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SLERP Visualization on the Unit Sphere\n",
    "\n",
    "This notebook demonstrates Spherical Linear Interpolation (SLERP) between quaternions with an animated visualization showing the interpolation path on the unit sphere."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from IPython.display import HTML\n",
    "from matplotlib.animation import FuncAnimation\n",
    "\n",
    "from fastquat import Quaternion\n",
    "\n",
    "# Configure matplotlib for better rendering\n",
    "plt.rcParams['figure.dpi'] = 100\n",
    "plt.rcParams['animation.html'] = 'html5'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup: Define Quaternions and Test Vector\n",
    "\n",
    "We'll interpolate between two quaternions that represent significantly different rotations to make the arc clearly visible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define two quaternions with a large angular separation for a visible arc\n",
    "# q1: Identity (no rotation)\n",
    "q1 = Quaternion.ones()\n",
    "\n",
    "# q2: 120 degree rotation around axis (1, 1, 1) for a large, visible arc\n",
    "axis = jnp.array([1.0, 1.0, 1.0])\n",
    "axis = axis / jnp.linalg.norm(axis)  # Normalize\n",
    "angle = 2 * jnp.pi / 3  # 120 degrees\n",
    "q2 = Quaternion.from_scalar_vector(jnp.cos(angle / 2), jnp.sin(angle / 2) * axis)\n",
    "\n",
    "print(f'Start quaternion q1: {q1}')\n",
    "print(f'End quaternion q2: {q2}')\n",
    "print(f'Dot product: {jnp.sum(q1.wxyz * q2.wxyz):.4f}')\n",
    "print(\n",
    "    f'Angular separation: {jnp.degrees(2 * jnp.arccos(jnp.abs(jnp.sum(q1.wxyz * q2.wxyz)))):.2f}'\n",
    "    f' radians'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate SLERP Interpolation\n",
    "\n",
    "Create a smooth interpolation between the quaternions and apply them to a test vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate interpolation parameters for smooth animation\n",
    "n_frames = 50\n",
    "t_values = jnp.linspace(0, 1, n_frames)\n",
    "\n",
    "# SLERP interpolation\n",
    "slerp_quaternions = q1.slerp(q2, t_values)\n",
    "\n",
    "# Test vector to rotate (unit vector along x-axis)\n",
    "test_vector = jnp.array([1.0, 0.0, 0.0])\n",
    "\n",
    "# Apply rotations to test vector\n",
    "slerp_rotated = slerp_quaternions.rotate_vector(test_vector)\n",
    "\n",
    "print(f'Generated {len(slerp_rotated)} frames for animation')\n",
    "print(f'All quaternions normalized: {jnp.allclose(slerp_quaternions.norm(), 1.0)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animated Display\n",
    "\n",
    "Now create an animation showing the interpolation progressing along the arc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create animated visualization\n",
    "fig = plt.figure(figsize=(10, 8))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "# Draw unit sphere (static)\n",
    "u = np.linspace(0, 2 * np.pi, 50)\n",
    "v = np.linspace(0, np.pi, 50)\n",
    "sphere_x = np.outer(np.cos(u), np.sin(v))\n",
    "sphere_y = np.outer(np.sin(u), np.sin(v))\n",
    "sphere_z = np.outer(np.ones(np.size(u)), np.cos(v))\n",
    "ax.plot_surface(sphere_x, sphere_y, sphere_z, alpha=0.1, color='lightgray')\n",
    "\n",
    "# Plot the complete path in light color\n",
    "ax.plot(\n",
    "    slerp_rotated[:, 0],\n",
    "    slerp_rotated[:, 1],\n",
    "    slerp_rotated[:, 2],\n",
    "    '--',\n",
    "    color='lightblue',\n",
    "    linewidth=1,\n",
    "    alpha=0.5,\n",
    "    label='Full Path',\n",
    ")\n",
    "\n",
    "# Start and end points\n",
    "ax.scatter(*slerp_rotated[0], color='green', s=100, label='Start', zorder=5)\n",
    "ax.scatter(*slerp_rotated[-1], color='red', s=100, label='End', zorder=5)\n",
    "\n",
    "# Initialize animated elements\n",
    "current_point = ax.scatter([], [], [], color='blue', s=150, zorder=6)\n",
    "(progress_line,) = ax.plot([], [], [], 'b-', linewidth=3, label='Progress')\n",
    "\n",
    "ax.set_xlabel('X')\n",
    "ax.set_ylabel('Y')\n",
    "ax.set_zlabel('Z')\n",
    "ax.set_title('SLERP Animation on Unit Sphere')\n",
    "ax.legend()\n",
    "ax.set_box_aspect([1, 1, 1])\n",
    "\n",
    "# Set fixed view limits\n",
    "ax.set_xlim([-1.2, 1.2])\n",
    "ax.set_ylim([-1.2, 1.2])\n",
    "ax.set_zlim([-1.2, 1.2])\n",
    "\n",
    "# Apply -90° rotation around k (z-axis) to orient towards viewer\n",
    "ax.view_init(elev=20, azim=45)\n",
    "\n",
    "# Initialize vector arrow variable\n",
    "origin = [0, 0, 0]\n",
    "vector_arrow = None\n",
    "\n",
    "\n",
    "def animate(frame):\n",
    "    global vector_arrow\n",
    "\n",
    "    # Update current point\n",
    "    current_pos = slerp_rotated[frame]\n",
    "    current_point._offsets3d = ([current_pos[0]], [current_pos[1]], [current_pos[2]])\n",
    "\n",
    "    # Update progress line (path so far)\n",
    "    progress_line.set_data_3d(\n",
    "        slerp_rotated[: frame + 1, 0],\n",
    "        slerp_rotated[: frame + 1, 1],\n",
    "        slerp_rotated[: frame + 1, 2],\n",
    "    )\n",
    "\n",
    "    # Update vector arrow from origin to current point\n",
    "    # Remove previous arrow if it exists\n",
    "    if vector_arrow is not None:\n",
    "        vector_arrow.remove()\n",
    "\n",
    "    # Create new arrow\n",
    "    vector_arrow = ax.quiver(\n",
    "        *origin,\n",
    "        *current_pos,\n",
    "        color='red',\n",
    "        arrow_length_ratio=0.1,\n",
    "        linewidth=3,\n",
    "        alpha=0.7,\n",
    "    )\n",
    "\n",
    "    return current_point, progress_line\n",
    "\n",
    "\n",
    "# Create animation - don't show the final static frame\n",
    "anim = FuncAnimation(fig, animate, frames=n_frames, interval=100, blit=False, repeat=True)\n",
    "\n",
    "# Close the static figure to prevent it from showing\n",
    "plt.close(fig)\n",
    "\n",
    "# Only display the HTML animation widget\n",
    "HTML(anim.to_jshtml())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}