{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rot_Z(alpha):\n",
    "    return array([[cos(alpha), sin(alpha), 0, 0], [-sin(alpha), cos(alpha), 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rot_Y(alpha, complementer = True):\n",
    "    if complementer:\n",
    "        alpha = pi - alpha\n",
    "    return array([[cos(alpha), 0, sin(alpha), 0], [0, 1, 0, 0], [-sin(alpha), 0, cos(alpha), 0], [0, 0, 0, 1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shift_Z(delta):\n",
    "    return array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, delta], [0, 0, 0, 1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class robGEOMETRY:\n",
    "    H_BODY = 1\n",
    "    L_ARM1 = 1\n",
    "    L_ARM2 = 1\n",
    "    L_GRIP = 1\n",
    "    L_GRIP_MID = .5\n",
    "    A_GRIP_MID = pi / 20\n",
    "    \n",
    "    Z_THRESH = 0\n",
    "    \n",
    "    ORIGO = array([0, 0, 0, 1])\n",
    "    P1 = array([0, 0, H_BODY, 1])\n",
    "    \n",
    "    @staticmethod\n",
    "    def DzOz(alpha, delta):\n",
    "        return dot(shift_Z(delta), rot_Z(alpha))\n",
    "    \n",
    "    @staticmethod\n",
    "    def OyDz(alpha, delta, complementer = True):\n",
    "        return dot(rot_Y(alpha, complementer), shift_Z(delta))\n",
    "\n",
    "    @staticmethod\n",
    "    def OyOz(alpha, beta):\n",
    "        return dot(rot_Y(alpha), rot_Z(beta))\n",
    "    \n",
    "    def __init__(self, angles = array([0, 3 * pi / 4, 3 * pi / 4, pi, 0, pi / 8])):\n",
    "        assert len(angles) == 6\n",
    "        self._angles = angles\n",
    "    \n",
    "    @property\n",
    "    def angles(self): return self._angles\n",
    "    @angles.setter\n",
    "    def angles(self, a):\n",
    "        self._angles = a\n",
    "        \n",
    "    @property\n",
    "    def T1(self): return self.DzOz(self.angles[0], self.H_BODY)\n",
    "    @property\n",
    "    def T2(self): return self.OyDz(self.angles[1], self.L_ARM1)\n",
    "    @property\n",
    "    def T3(self): return self.OyDz(self.angles[2], self.L_ARM2)\n",
    "    @property\n",
    "    def T4(self): return dot(rot_Y(self.angles[3], complementer = False), rot_Z(self.angles[4]))\n",
    "    @property\n",
    "    def T5s(self):\n",
    "        a = self.angles[5]\n",
    "        return [\n",
    "            self.OyDz(a, self.L_GRIP, complementer = False),\n",
    "            self.OyDz(-a, self.L_GRIP, complementer = False),\n",
    "            self.OyDz(a + self.A_GRIP_MID, self.L_GRIP_MID, complementer = False),\n",
    "            self.OyDz(-a - self.A_GRIP_MID, self.L_GRIP_MID, complementer = False)\n",
    "        ]\n",
    "    \n",
    "#    @property\n",
    "#    def P1(self): return dot(self.T1, self.ORIGO)\n",
    "    @property\n",
    "    def P2(self): return dot(self.T1, dot(self.T2, self.ORIGO))\n",
    "    @property\n",
    "    def P3(self): return dot(self.T1, dot(self.T2, dot(self.T3, self.ORIGO)))\n",
    "    @property\n",
    "    def PGs(self): return [ dot(self.T1, dot(self.T2, dot(self.T3, dot(self.T4, dot(T5, self.ORIGO))))) for T5 in self.T5s ]\n",
    "\n",
    "    def seg(self):\n",
    "        P5 = self.PGs\n",
    "        return [\n",
    "            (\n",
    "                self.ORIGO[i], self.P1[i], self.P2[i], self.P3[i], \n",
    "                P5[0][i], P5[2][i], self.P3[i], \n",
    "                P5[1][i], P5[3][i], self.P3[i]\n",
    "            ) for i in range(3)\n",
    "        ]\n",
    "    \n",
    "    def test(self):\n",
    "        for p in self.PGs:\n",
    "            if p[2] <= self.Z_THRESH:\n",
    "                return False\n",
    "        return True\n",
    "    \n",
    "    def inv(self, P3):\n",
    "        D = norm(P3 - self.P1)\n",
    "        assert D < self.L_ARM1 + self.L_ARM2, \"triangle inequality broken\"\n",
    "        x, y, z = P3[0:3]\n",
    "        dh = z - self.H_BODY\n",
    "        l_ = sqrt(x * x + y * y)\n",
    "        alpha = arctan2(y, x)\n",
    "        beta = pi / 2 + arctan2(dh, l_)\n",
    "        gamma = arccos((D * D - self.L_ARM1 * self.L_ARM1 - self.L_ARM2 * self.L_ARM2) / 2 / self.L_ARM1 / self.L_ARM2)\n",
    "        delta = sum(self.angles[:4]) - (alpha + beta + gamma)\n",
    "        return alpha, beta, gamma, delta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "rg = robGEOMETRY()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd93c9c536e84c87933f47ae3a5cc0d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=False, description='Hit floor', disabled=True)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e052373b1895425f8430e856d5047e56",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.0, description='a_body', max=1.5707963267948966, min=-1.570796326794…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wt = ipywidgets.FloatSlider\n",
    "c = rg.angles\n",
    "checker = ipywidgets.Checkbox(disabled = True, description = 'Hit floor', value = not rg.test())\n",
    "display(checker)\n",
    "@ipywidgets.interact(\n",
    "    a_body = wt(min = -pi / 2, max = pi / 2, value = c[0]),\n",
    "    a_shoulder = wt(min = 0, max = pi, value = c[1]),\n",
    "    a_elbow = wt(min = 0, max = pi, value = c[2]),\n",
    "    a_wrist = wt(min = -pi / 2, max = pi / 2, value = c[3]),\n",
    "    a_wrist_rot = wt(min = -pi, max = pi, value = c[4]),\n",
    "    a_grip = wt(min = 0, max = pi / 4, value = c[5])\n",
    ")\n",
    "def _p(a_body, a_shoulder, a_elbow, a_wrist, a_wrist_rot, a_grip):\n",
    "    rg.angles = array([a_body, a_shoulder, a_elbow, a_wrist, a_wrist_rot, a_grip])\n",
    "    checker.value = not rg.test()\n",
    "    P = rg.seg()\n",
    "    figure(figsize = (8, 8))\n",
    "    subplot(2, 2, 1)\n",
    "    plot(P[0], P[2])\n",
    "    xlabel('x')\n",
    "    ylabel('z')\n",
    "    xlim((-4, 4))\n",
    "    ylim((-4, 4))\n",
    "    grid()\n",
    "    subplot(2, 2, 2)\n",
    "    plot(P[1], P[2])\n",
    "    xlabel('y')\n",
    "    ylabel('z')\n",
    "    xlim((-4, 4))\n",
    "    ylim((-4, 4))\n",
    "    grid()\n",
    "    subplot(2, 2, 3)\n",
    "    plot(P[0], P[1])\n",
    "    xlabel('x')\n",
    "    ylabel('y')\n",
    "    xlim((-4, 4))\n",
    "    ylim((-4, 4))\n",
    "    grid()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.7853981633974483, 0.9553166181245093, 1.047197551196598, 3.495272974461031)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rg.inv(array([1, 1, 0, 1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}