{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inverse Kinematics Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import fsolve\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Polar coordinate functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cartesian_to_polar(x, y):\n",
    "    r = np.sqrt(x**2 + y**2)\n",
    "    theta = np.arctan2(y, x)\n",
    "    return r, theta\n",
    "\n",
    "def polar_to_cartesian(r, theta):\n",
    "    x = r * np.cos(theta)\n",
    "    y = r * np.sin(theta)\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate end-joint values from xyz position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[82.5301003420473,\n",
       " -83.54918182551367,\n",
       " 111.24503912126411,\n",
       " -117.69585729575046,\n",
       " -90.0,\n",
       " 90.0]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def normalize_degree(theta):\n",
    "    # Normalizes degree theta from -1.5pi to 1.5pi\n",
    "    multiplier = 1.5\n",
    "    normalized_theta = theta % (math.pi * multiplier)\n",
    "    # Maintain the negative sign if the original angle is negative\n",
    "    if theta < 0:\n",
    "        normalized_theta -= math.pi * multiplier\n",
    "\n",
    "    # Return angle\n",
    "    return normalized_theta\n",
    "\n",
    "def get_joints_from_xyz_rel(x, y, z, rx=0, ry=-math.pi/2, rz=0, initial_guess = (math.pi/2, math.pi/2, 0)):\n",
    "    # Get limbs and offsets\n",
    "    offset_x, offset_y, offset_z = (0, 0, 0.14)     # Tool offset\n",
    "    l_bs, l1, l2, l3, l_wt = (0.1333, .425, .39225, .1267, .0997)    # Limb lengths\n",
    "    \n",
    "    # Calculate base angle and r relative to shoulder joint\n",
    "    def calculate_theta(x, y, a):\n",
    "        # Calculate if we need the + or - in our equations\n",
    "        if (x>a and y>=0) or (x>-a and y<0):\n",
    "            flip = 1\n",
    "        elif (x<a and y>=0) or (x<-a and y<0):\n",
    "            flip = -1\n",
    "        else: \n",
    "            # Critical section (x=a, or x=-a). Infinite slope\n",
    "            # Return 0 or 180 depending on sign\n",
    "            return math.atan2(y, 0) - math.pi/2\n",
    "        \n",
    "        # Calculate tangent line y = mx + b\n",
    "        if abs(a) != abs(x): # If there is no division by 0\n",
    "            m = (x*y + math.sqrt(x*x*y*y-(x*x-a*a)*(y*y-a*a)))/(x*x-a*a)\n",
    "        else: # Deal with edge case when x^2=a^2\n",
    "            m = flip*(-a*a+y*y)/(a*y-flip*abs(a*y))\n",
    "        b = flip * a * math.sqrt(1+m*m)\n",
    "\n",
    "        # Calculate equivalent tangent point on circle\n",
    "        cx = (-flip*m*b)/(1+m*m)\n",
    "        cy = m*cx + flip*b\n",
    "\n",
    "        # Calculate base angle, make angle negative if flip=1\n",
    "        theta = math.atan2(cy, cx) + (-math.pi if flip==1 else 0)\n",
    "\n",
    "        return theta \n",
    "    \n",
    "    base_theta = calculate_theta(x, y, l_bs)\n",
    "    cx, cy = l_bs*math.cos(base_theta), l_bs*math.sin(base_theta)\n",
    "    r = math.sqrt((x-cx)**2 + (y-cy)**2) \n",
    "\n",
    "\n",
    "    # Formulas to find out joint positions for (r, z)\n",
    "    def inv_kin_r_z(p):\n",
    "        a, b, c = p            \n",
    "\n",
    "        return (l1*math.cos(a) + l2*math.cos(a-b) + l3*math.cos(a-b-c) - r, # r\n",
    "                l1*math.sin(a) + l2*math.sin(a-b) - l3*math.sin(a-b-c) - (l3*math.sin(a-b-c)) - (z + offset_z),  # z\n",
    "                a-b-c) # wrist angle\n",
    "\n",
    "\n",
    "    # Normalize angles\n",
    "    base, shoulder, elbow, wrist1 = [normalize_degree(deg) for deg in [base_theta, *fsolve(inv_kin_r_z, initial_guess)]]\n",
    "\n",
    "    # Return result\n",
    "    return base, shoulder, elbow, wrist1, ry, rz\n",
    "\n",
    "def get_joints_from_xyz_abs(x, y, z, rx=0, ry=-math.pi/2, rz=math.pi/2):\n",
    "    joints = get_joints_from_xyz_rel(x, y, z, rx, ry, rz)\n",
    "\n",
    "    # Joint offsets\n",
    "    # Base, Shoulder, Elbow, Wrist\n",
    "    inverse = [1, -1, 1, 1, 1, 1]\n",
    "    offsets = [-math.pi/2, 0, 0, -math.pi/2, 0, 0]\n",
    "\n",
    "    # Return adjusted joint positions\n",
    "    robot_angles = [o+j*i for j, o, i in zip(joints, offsets, inverse)]\n",
    "\n",
    "    return robot_angles\n",
    "\n",
    "# Print degree rotation for each joint (robot angles)\n",
    "[math.degrees(deg) for deg in get_joints_from_xyz_abs(-0.2, -0.5, 0.1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulate Arm and Joint Angles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_arm_side_view(x, y, z, details=False):\n",
    "\n",
    "    # Get joint angles\n",
    "    l1, l2, l3 = .425, .39225, .1267\n",
    "    offset_x, offset_y, offset_z = (0, 0, 0.14)\n",
    "    r, theta = cartesian_to_polar(x, y)\n",
    "    base, shoulder, elbow, wrist, _, _ = get_joints_from_xyz_rel(x, y, z)\n",
    "\n",
    "    # Print angles\n",
    "    if details:\n",
    "        print('Target position (x,y,z):', x, y, z)\n",
    "        print('R: ', round(math.sqrt(x**2+y**2),4))\n",
    "        print('Angles (base, shoulder, elbow, wrist):', [round(math.degrees(i), 4) for i in [base, shoulder, elbow, wrist]])\n",
    "        print('Robot Angles:', [round(math.degrees(i), 4) for i in get_joints_from_xyz_abs(x, y, z)])\n",
    "\n",
    "    # Calculate each joint's endpoint position\n",
    "    x1, y1 = polar_to_cartesian(l1, shoulder)\n",
    "    x2, y2 = polar_to_cartesian(l2, shoulder-elbow)\n",
    "    x2 += x1\n",
    "    y2 += y1\n",
    "    x3, y3 = polar_to_cartesian(l3, shoulder-elbow-wrist)\n",
    "    x3 += x2\n",
    "    y3 += y2 \n",
    "    \n",
    "    tx = x3\n",
    "    ty = y3 - offset_z\n",
    "\n",
    "    # Print each joint's endpoint position\n",
    "    if details:\n",
    "        print('elbow (x,y):', round(x1,3), round(y1,3))\n",
    "        print('wrist (x,y):', round(x2,3), round(y2,3))\n",
    "        print('tool  (x,y):', round(x3,3), round(y3,3))\n",
    "\n",
    "    # Draw limbs\n",
    "    plt.plot([0, x1], [0, y1], color='cyan', linewidth=7)\n",
    "    plt.plot([x1, x2], [y1, y2], color='orange', linewidth=7)\n",
    "    plt.plot([x2, x2+l3], [y2, y2], color='red', linewidth=7)\n",
    "    plt.plot()\n",
    "\n",
    "    # Draw toolpoint\n",
    "    plt.plot([x3, tx], [y3, ty], color='black', linewidth=7)\n",
    "\n",
    "    # Display angles\n",
    "    plt.text(0, 0.02, f'{round(math.degrees(shoulder), 2)}°')\n",
    "    plt.text(x1, y1+0.02, f'{round(math.degrees(elbow), 2)}°')\n",
    "    plt.text(x2, y2+0.02, f'{round(math.degrees(wrist), 2)}°')\n",
    "\n",
    "    # Display r arrow\n",
    "    plt.annotate(f'', xy=(0, 0), xycoords='data', xytext=(x3, 0), textcoords='data', arrowprops={'arrowstyle': '<->'})\n",
    "    plt.annotate(f'r={round(r,4)}', xy=(x2/2, 0.01), xycoords='data', xytext=(x2/2, 0), textcoords='offset points')\n",
    "\n",
    "    # Display z arrow\n",
    "    plt.annotate(f'', xy=(x3, 0), xycoords='data', xytext=(tx, ty), textcoords='data', arrowprops={'arrowstyle': '<->'})\n",
    "    plt.annotate(f'z={round(ty,4)}', xy=(tx+0.01, ty/2), xycoords='data', xytext=(x3/2, 0), textcoords='offset points')\n",
    "   \n",
    "    # Display plot\n",
    "    ax = plt.subplot(111)\n",
    "    ax.spines[['right', 'top']].set_visible(False)\n",
    "    plt.axis('equal')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_arm_top_view(x, y, z, details=False):\n",
    "    # Get joint and position information\n",
    "    angles = get_joints_from_xyz_rel(x, y, z)\n",
    "    offset_x, offset_y, offset_z = (0, 0, 0.14)     # Tool offset\n",
    "    l_bs, l1, l2, l3, l_wt = (0.1333, .425, .39225, .1267, .0997)    # Limb lengths\n",
    "    cx, cy = l_bs*math.cos(angles[0]), l_bs*math.sin(angles[0]) # Base tangent point\n",
    "    line_angle = math.pi/2+angles[0]\n",
    "\n",
    "    if details:\n",
    "        print('Angles:', [round(math.degrees(angle),3) for angle in angles])\n",
    "        print(f'Circle position (cx, cy): ({round(cx, 3)}, {round(cy, 3)})')\n",
    "\n",
    "\n",
    "    # Plot coordinate system\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(1, 1, 1)\n",
    "    ax.spines['left'].set_position('center')\n",
    "    ax.spines['bottom'].set_position('center')\n",
    "    ax.spines['right'].set_color('none')\n",
    "    ax.spines['top'].set_color('none')\n",
    "    ax.set_yticklabels([])\n",
    "    ax.set_xticklabels([])\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "\n",
    "    # Target point\n",
    "    # plt.plot(x, y, 'go')\n",
    "\n",
    "    # Circle\n",
    "    circle = plt.Circle((0, 0), l_bs, color='b', fill=False, linewidth=1)\n",
    "    plt.plot([0, cx], [0, cy], color='b', linewidth=1)\n",
    "    ax.add_patch(circle)\n",
    "\n",
    "    # Draw limbs\n",
    "    # Shoulder\n",
    "    x1, y1 = polar_to_cartesian(l1*math.cos(angles[1]), line_angle)\n",
    "    x1, y1 = cx+x1, cy+y1\n",
    "    z1 = l1*math.sin(angles[1])\n",
    "    plt.plot([cx, x1], [cy, y1], color='cyan', linewidth=3)\n",
    "\n",
    "    # Elbow\n",
    "    x2, y2 = polar_to_cartesian(l2*math.cos(angles[1]-angles[2]), line_angle)\n",
    "    x2 += x1\n",
    "    y2 += y1\n",
    "    z2 = z1 + l2*math.sin(angles[1]-angles[2])\n",
    "    plt.plot([x1, x2], [y1, y2], color='orange', linewidth=2)\n",
    "\n",
    "    # Wrist\n",
    "    x3, y3 = polar_to_cartesian(l3*math.cos(angles[1]-angles[2]-angles[3]), line_angle)\n",
    "    x3 += x2\n",
    "    y3 += y2\n",
    "    z3 = z2 + l3*math.sin(angles[1]-angles[2]-angles[3]) \n",
    "    plt.plot([x2, x3], [y2, y3], color='red', linewidth=2)\n",
    "\n",
    "    # Print joint positions\n",
    "    if details:\n",
    "        print(f'Shoulder (x, y, z): ({round(x1,3)}, {round(y1,3)}, {round(z1,3)})')\n",
    "        print(f'Elbow    (x, y, z): ({round(x2,3)}, {round(y2,3)}, {round(z2,3)})')\n",
    "        print(f'Wrist    (x, y, z): ({round(x3,3)}, {round(y3,3)}, {round(z3,3)})')\n",
    "\n",
    "    # Display angle\n",
    "    plt.text(0.01, -0.01-cy/abs(cy+0.00001)*0.02, f'{round(math.degrees(angles[0]), 2)}°', fontsize=7)\n",
    "\n",
    "    # Display x arrow\n",
    "    sign = x3/abs(x3)\n",
    "    plt.annotate(f'', xy=(0, y3), xycoords='data', xytext=(x3, y3), textcoords='data', arrowprops={'arrowstyle': '<->'})\n",
    "    plt.annotate(f'x={round(x3,3)}', xy=(-0.1-sign*0.13, y3-0.015), xycoords='data', xytext=(x2/2, 0), textcoords='offset points')\n",
    "\n",
    "    # Display y arrow\n",
    "    sign = y3/abs(y3)\n",
    "    plt.annotate(f'', xy=(x3, 0), xycoords='data', xytext=(x3, y3), textcoords='data', arrowprops={'arrowstyle': '<->'})\n",
    "    plt.annotate(f'y={round(y3,3)}', xy=((x3-0.1), -0.015-sign*0.03), xycoords='data', xytext=(x3/2, 0), textcoords='offset points')\n",
    "\n",
    "\n",
    "    # Set axis limits and labels\n",
    "    axis_limit = math.hypot(x, y)+.1\n",
    "    plt.axis('square')\n",
    "    plt.xlim(-axis_limit, axis_limit)\n",
    "    plt.ylim(-axis_limit, axis_limit)\n",
    "\n",
    "    # Adjust the position of axis labels\n",
    "    plt.xlabel('+x', horizontalalignment='right', x=1.05)\n",
    "    plt.ylabel('+y', verticalalignment='top', rotation=0, y=1.05)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Angles: [-53.867, 89.583, 87.39, 2.193, -90.0, 0.0]\n",
      "Circle position (cx, cy): (0.079, -0.108)\n",
      "Shoulder (x, y, z): (0.081, -0.106, 0.425)\n",
      "Elbow    (x, y, z): (0.398, 0.125, 0.44)\n",
      "Wrist    (x, y, z): (0.5, 0.2, 0.44)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 14004 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Target position (x,y,z): 0.5 0.2 0.3\n",
      "R:  0.5385\n",
      "Angles (base, shoulder, elbow, wrist): [-53.8671, 89.5827, 87.3895, 2.1932]\n",
      "Robot Angles: [-143.8671, -89.5827, 87.3895, -87.8068, -90.0, 90.0]\n",
      "elbow (x,y): 0.003 0.425\n",
      "wrist (x,y): 0.395 0.44\n",
      "tool  (x,y): 0.522 0.44\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 18820 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def draw_arm(x, y, z, details=False):\n",
    "    draw_arm_top_view(x,y,z, details=details)\n",
    "    draw_arm_side_view(x,y,z, details=details)\n",
    "\n",
    "draw_arm(0.5,0.2,0.3, details=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'elbow': (0.06441407671759522, -0.08297834710053949, 0.42498873000714044),\n",
       " 'wrist': (0.3809831975755794, 0.14814654552767514, 0.44000000000002354),\n",
       " 'wrist2': (0.4617483203289997, 0.20711259462904327, 0.44000000000002354),\n",
       " 'tool': (0.4617483203289997, 0.20711259462904327, 0.5800000000000236)}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def joint_positions(x, y, z):\n",
    "    # Get joint and position information\n",
    "    angles = get_joints_from_xyz_abs(x, y, z)\n",
    "    offset_x, offset_y, offset_z = (0, 0, 0.14)     # Tool offset\n",
    "    l_bs, l1, l2, l3, l_wt = (0.1333, .425, .39225, .1267, .0997)    # Limb lengths\n",
    "    cx, cy = l_bs*math.cos(angles[0]), l_bs*math.sin(angles[0]) # Base tangent point\n",
    "    line_angle = math.pi/2+angles[0]\n",
    "\n",
    "    # Elbow\n",
    "    x1, y1 = polar_to_cartesian(l1*math.cos(angles[1]), line_angle)\n",
    "    x1, y1 = cx+x1, cy+y1\n",
    "    z1 = l1*math.sin(angles[1])\n",
    "\n",
    "    # Wrist\n",
    "    x2, y2 = polar_to_cartesian(l2*math.cos(angles[1]-angles[2]), line_angle)\n",
    "    x2 += x1\n",
    "    y2 += y1\n",
    "    z2 = z1 + l2*math.sin(angles[1]-angles[2])\n",
    "\n",
    "    # Wrist 2\n",
    "    x3, y3 = polar_to_cartesian(l3*math.cos(angles[1]-angles[2]-angles[3]), line_angle)\n",
    "    x3 += x2\n",
    "    y3 += y2\n",
    "    z3 = z2 + l3*math.sin(angles[1]-angles[2]-angles[3]) \n",
    "\n",
    "    # Tool\n",
    "    x4, y4 = polar_to_cartesian(offset_z*math.cos(angles[1]-angles[2]-angles[3]-angles[4]), line_angle)\n",
    "    x4 += x3\n",
    "    y4 += y3\n",
    "    z4 = z3 + offset_z*math.sin(angles[1]-angles[2]-angles[3]-angles[4])\n",
    "\n",
    "\n",
    "    positions = {\n",
    "        'elbow': (x1, y1, z1),\n",
    "        'wrist': (x2, y2, z2),\n",
    "        'wrist2': (x3, y3, z3),\n",
    "        'tool': (x4, y4, z4)\n",
    "    }\n",
    "\n",
    "    return positions\n",
    "\n",
    "joint_positions(0.5,0.2,0.3)"
   ]
  }
 ],
 "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
