{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cadquery as cq \n",
    "from jupyter_cadquery import (\n",
    "    open_viewer, show,\n",
    "    set_defaults, get_defaults, \n",
    ")\n",
    "\n",
    "from cadquery_massembly import Mate, MAssembly, relocate\n",
    "\n",
    "# remove \"clean\" to avoid errors OCP kernel error\n",
    "cq.occ_impl.shapes.Shape.clean = lambda x: x\n",
    "\n",
    "cv = open_viewer(\"DiskArm\", cad_width=640, height=500, theme=\"light\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_defaults(axes=False, axes0=True, mate_scale=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import pi, sin, cos, sqrt, asin, degrees, radians\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Disk and Arm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![disk-arm](1-disk-arm.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_disk = 100\n",
    "dist_pivot = 200\n",
    "\n",
    "def angle_arm(angle_disk):\n",
    "    ra = np.deg2rad(angle_disk)\n",
    "    v = np.array((dist_pivot, 0)) - r_disk * np.array((cos(ra), sin(ra)))\n",
    "    return np.rad2deg(np.arctan2(*v[::-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def diagram(da, aa, i1, i2):\n",
    "    a1, a2 = radians(da), radians(aa)\n",
    "    r1, r2 = r_disk, dist_pivot + r_disk\n",
    "    x1, y1 = (0, r1 * cos(a1)), (0, r1 * sin(a1))\n",
    "    x2, y2 = (dist_pivot, dist_pivot - r2 * cos(a2)), (0, -r2 * sin(a2))\n",
    "\n",
    "    ax = fig.add_subplot(spec[i1, i2])\n",
    "    ax.set_ylim(-r_disk-10, r_disk+10)\n",
    "    ax.set_title(f\"a={da}\")\n",
    "\n",
    "    ax.plot((0, dist_pivot), (0,0), color=\"lightgrey\")\n",
    "    ax.plot(x1, y1)\n",
    "    ax.plot(x2, y2, \":\")\n",
    "    \n",
    "    circle = plt.Circle((0,0), 100, fill=False, color=\"lightgrey\", linestyle='--')\n",
    "    ax.add_patch(circle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "disk_angles = range(0, 360, 45)\n",
    "arm_angles = [angle_arm(d) for d in disk_angles]\n",
    "\n",
    "fig = plt.figure(constrained_layout=True)\n",
    "fig.set_size_inches(20, 1.75)\n",
    "spec = gridspec.GridSpec(ncols=8, nrows=1, figure=fig)\n",
    "\n",
    "for i, (da, aa) in enumerate(zip(disk_angles, arm_angles)):\n",
    "    diagram(da, aa, 0, i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Assembly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thickness = 5\n",
    "nr = 5\n",
    "\n",
    "disk = cq.Workplane().circle(r_disk + 2 * nr).extrude(thickness)\n",
    "nipple = cq.Workplane().circle(nr).extrude(thickness)\n",
    "disk = (\n",
    "    disk\n",
    "    .cut(nipple)\n",
    "    .union(nipple.translate((r_disk, 0, thickness)))\n",
    ")\n",
    "\n",
    "pivot_base = cq.Workplane().circle(2*nr).extrude(thickness)\n",
    "base = (\n",
    "    cq.Workplane()\n",
    "    .rect(6 * nr + dist_pivot, 6 * nr)\n",
    "    .extrude(thickness)\n",
    "    .translate((dist_pivot / 2, 0, 0))\n",
    "    .union(nipple.translate((dist_pivot, 0, thickness)))\n",
    "    .union(pivot_base.translate((0, 0, thickness)))\n",
    "    .union(nipple.translate((0, 0, 2*thickness)))\n",
    "    .edges(\"|Z\").fillet(3)\n",
    ")\n",
    "base.faces(\">Z[-2]\").wires(cq.NearestToPointSelector((dist_pivot + r_disk, 0))).tag(\"mate\")\n",
    "\n",
    "slot = (cq.Workplane()\n",
    "    .rect(2*r_disk, 2*nr)\n",
    "    .extrude(thickness)\n",
    "    .union(nipple.translate((-r_disk, 0, 0)))\n",
    "    .union(nipple.translate((r_disk, 0, 0)))\n",
    "    .translate((dist_pivot, 0, 0))\n",
    ")\n",
    "                                                                        \n",
    "arm = (\n",
    "    cq.Workplane()\n",
    "    .rect(4 * nr + (r_disk + dist_pivot), 4 * nr)\n",
    "    .extrude(thickness)\n",
    "    .edges(\"|Z\").fillet(3)    \n",
    "    .translate(((r_disk + dist_pivot) / 2, 0, 0))\n",
    "    .cut(nipple)\n",
    "    .cut(slot)\n",
    ")\n",
    "arm.faces(\">Z\").wires(cq.NearestToPointSelector((0,0))).tag(\"mate\")\n",
    "\n",
    "show(\n",
    "    disk,\n",
    "    base.translate((0, -1.5 * r_disk, 0)),\n",
    "    arm.translate((0, 1.5 * r_disk, 0)),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define assembly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import webcolors\n",
    "\n",
    "def create_disk_arm():\n",
    "    L = lambda *args: cq.Location(cq.Vector(*args))\n",
    "\n",
    "    return (MAssembly(base, name=\"base\", color=\"silver\", loc=L(-dist_pivot/2, 0, 0))\n",
    "        .add(disk, name=\"disk\", color=\"MediumAquaMarine\", loc=L(r_disk, -1.5 * r_disk, 0))\n",
    "        .add(arm,  name=\"arm\",  color=\"orange\", loc=L(0,  10*nr, 0))\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define mates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict as odict\n",
    "\n",
    "disk_arm = create_disk_arm()\n",
    "\n",
    "disk_arm.mate(\"base?mate\", name=\"disk_pivot\", origin=True, transforms=odict(rz=180))\n",
    "disk_arm.mate(\"base@faces@>Z\", name=\"arm_pivot\")\n",
    "disk_arm.mate(\"disk@faces@>Z[-2]\", name=\"disk\", origin=True)\n",
    "disk_arm.mate(\"arm?mate\", name=\"arm\", origin=True)\n",
    "\n",
    "show(disk_arm, render_mates=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Relocate and assemble"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ensure all parts are relocated so that the origin mates is the part origin \n",
    "disk_arm.relocate()\n",
    "\n",
    "# assemble each part\n",
    "disk_arm.assemble(\"arm\", \"arm_pivot\")\n",
    "disk_arm.assemble(\"disk\", \"disk_pivot\")\n",
    "\n",
    "d = show(disk_arm, render_mates=True, axes=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Animate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jupyter_cadquery import AnimationTrack\n",
    "\n",
    "times = np.linspace(0, 5, 181)\n",
    "disk_angles = np.linspace(0, 360, 181)\n",
    "arm_angles = [angle_arm(d) for d in disk_angles]\n",
    "\n",
    "# move disk\n",
    "# Note, the selector must follow the path in the CAD view navigation hierarchy\n",
    "d.add_track(AnimationTrack(f\"/base/disk\", \"rz\", times, disk_angles))\n",
    "\n",
    "# move arm\n",
    "d.add_track(AnimationTrack(f\"/base/arm\", \"rz\", times, arm_angles))\n",
    "\n",
    "d.animate(speed=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}