{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93688312",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Graphics demo\n",
    "\n",
    "Requirements:\n",
    "- [2-link Pendulum](../example_2_link_pendulum/notebook.ipynb)\n",
    "\n",
    "In this tutorial we will:\n",
    "- Create a two-link pendulum multibody\n",
    "- Attach a primitive geometry to a body\n",
    "- Attach a 3D mesh geometry to a body\n",
    "- Setup web-based 3d visualization\n",
    "- Render to a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce9260c1",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import atexit\n",
    "import numpy as np\n",
    "from typing import cast\n",
    "from pathlib import Path\n",
    "\n",
    "from IPython.display import display, Image\n",
    "\n",
    "import Karana.Core as kc\n",
    "import Karana.Math as km\n",
    "import Karana.Integrators as ki\n",
    "import Karana.Frame as kf\n",
    "import Karana.Dynamics as kd\n",
    "import Karana.Dynamics.SOADyn_types as kdt\n",
    "import Karana.Scene as ks\n",
    "import Karana.KUtils as ku\n",
    "import Karana.Models as kmdl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6531f439",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "resources = Path().absolute().parent / \"resources\"\n",
    "dodecahedron_file = str(resources / \"dodecahedron.obj\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf3ac55b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Setup a two-link pendulum multibody"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c24f60d7",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def createMbody() -> kd.Multibody:\n",
    "    \"\"\"Create the Multibody.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Multibody\n",
    "        The newly created multibody.\n",
    "    \"\"\"\n",
    "    fc = kf.FrameContainer(\"root\")\n",
    "    mb = kd.Multibody(\"mb\", fc)\n",
    "\n",
    "    # inertia used for both bodies\n",
    "    spI = km.SpatialInertia(2.0, np.zeros(3), np.diag([3, 2, 1]))\n",
    "\n",
    "    # inboard to joint transform used for both bodies\n",
    "    ib2j = km.HomTran()\n",
    "\n",
    "    # body to joint transform used for both bodies\n",
    "    b2j = km.HomTran(np.array([0, 0, 1]))\n",
    "\n",
    "    # setup body 1\n",
    "    bd1 = kd.PhysicalBody(\"bd1\", mb)\n",
    "    bd1.setSpatialInertia(spI)\n",
    "    bd1_hge = kd.PhysicalHinge(cast(kd.PhysicalBody, mb.virtualRoot()), bd1, kd.HingeType.REVOLUTE)\n",
    "    bd1.setBodyToJointTransform(b2j)\n",
    "    bd1_subhge = cast(kd.PinSubhinge, bd1_hge.subhinge(0))\n",
    "    bd1_subhge.setUnitAxis(np.array([0, 1, 0]))\n",
    "\n",
    "    # setup body 2\n",
    "    bd2 = kd.PhysicalBody(\"bd2\", mb)\n",
    "    bd2.setSpatialInertia(spI)\n",
    "    bd2_hge = kd.PhysicalHinge(bd1, bd2, kd.HingeType.REVOLUTE)\n",
    "    bd2_hge.onode().setBodyToNodeTransform(ib2j)\n",
    "    bd2.setBodyToJointTransform(b2j)\n",
    "    bd2_subhge = cast(kd.PinSubhinge, bd2_hge.subhinge(0))\n",
    "    bd2_subhge.setUnitAxis(np.array([0, 1, 0]))\n",
    "\n",
    "    mb.ensureHealthy()\n",
    "\n",
    "    # zero out hinge data (coords, velocities, accelerations, forces)\n",
    "    dark = np.zeros(mb.subhingeCoordData().nU())\n",
    "    mb.subhingeCoordData().setQ(dark)\n",
    "    mb.subhingeCoordData().setU(dark)\n",
    "    mb.subhingeCoordData().setUdot(dark)\n",
    "    mb.subhingeCoordData().setT(dark)\n",
    "\n",
    "    mb.setUniformGravAccel(np.zeros(3))\n",
    "    assert kc.allReady()\n",
    "\n",
    "    return mb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18d72d1a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Create the multibody and lookup important objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "faf0393e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "mb = createMbody()\n",
    "\n",
    "# get the first body\n",
    "bd1 = mb.getBody(\"bd1\")\n",
    "\n",
    "# get the second body\n",
    "bd2 = mb.getBody(\"bd2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47dcf2d1",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Geometry setup\n",
    "\n",
    "We create 3d geometries with associated material properties and attach them to bodies. This is commonly useful for collision detection and 3d visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "debf4014",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "cleanup_graphics, web_scene = mb.setupGraphics(port=0, wait_for_clients=1)\n",
    "proxy_scene = mb.getScene()\n",
    "\n",
    "# create a ball geometry to be used for both bodies\n",
    "ball_geom = ks.SphereGeometry(0.1)\n",
    "\n",
    "# load a dodecahedron geometry from file\n",
    "importer = ks.AssimpImporter()\n",
    "d12_geom = importer.importFrom(dodecahedron_file).geometries[0]\n",
    "\n",
    "# create red and blue materials to color the bodies\n",
    "mat_info = ks.PhysicalMaterialInfo()\n",
    "mat_info.color = ks.Color.BLUE\n",
    "blue = ks.PhysicalMaterial(mat_info)\n",
    "mat_info.color = ks.Color.RED\n",
    "red = ks.PhysicalMaterial(mat_info)\n",
    "\n",
    "# create a ks.ProxyScenePart and attach it to the body\n",
    "bd1_part = ks.ProxyScenePart(\"bd1_part\", scene=proxy_scene, geometry=ball_geom, material=blue)\n",
    "bd1_part.attachTo(bd1)\n",
    "\n",
    "bd2_part = ks.ProxyScenePart(\"bd2_part\", scene=proxy_scene, geometry=d12_geom, material=red)\n",
    "bd2_part.attachTo(bd2)\n",
    "bd2_part.setScale(0.2)\n",
    "\n",
    "# create some unattached objects (implicitly attached to the root frame)\n",
    "root_scene_node = ks.ProxySceneNode(\"root_scene_node\", scene=proxy_scene)\n",
    "\n",
    "del red, blue, mat_info, d12_geom, ball_geom"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f9e5412",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Setup camera angle and add ornamental parts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5238ac7a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# position the viewpoint camera of the visualization\n",
    "web_scene.defaultCamera().pointCameraAt([0.4, 4, -0.9], [0, 0, -1], [0, 0, 1])\n",
    "\n",
    "# visualize axes for important frames\n",
    "root_scene_node.graphics().showAxes(0.5)\n",
    "bd1_part.graphics().showAxes(0.5)\n",
    "bd2_part.graphics().showAxes(0.5)\n",
    "\n",
    "# add a trail to track the motion of body 2\n",
    "bd2_part.graphics().trail()\n",
    "\n",
    "proxy_scene.update()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a6cf2e2",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# set up state propagator\n",
    "sp = kd.StatePropagator(mb, ki.IntegratorType.CVODE)\n",
    "integ = sp.getIntegrator()\n",
    "\n",
    "# Create a UniformGravity model, and set it's gravitational acceleration.\n",
    "ug = kmdl.Gravity(\"grav_model\", sp, kmdl.UniformGravity(\"uniform_gravity\"), mb)\n",
    "ug.getGravityInterface().setGravity(np.array([0, 0, -9.81]), 0.0, kmdl.OutputUpdateType.PRE_HOP)\n",
    "del ug\n",
    "\n",
    "# Makes sure the visualization scene is updated after each state change.\n",
    "kmdl.UpdateProxyScene(\"update_proxy_scene\", sp, proxy_scene)\n",
    "\n",
    "# sync the simulation time with real-time.\n",
    "kmdl.SyncRealTime(\"sync_real_time\", sp, 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62d261df",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Register callback for end of each step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d7c172a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def post_step_fn(t, x):\n",
    "    \"\"\"Print out the time and state.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    t : float\n",
    "        The current time.\n",
    "    x : NDArray[np.float64]\n",
    "       The current state.\n",
    "    \"\"\"\n",
    "    print(f\"t = {float(integ.getTime()) / 1e9}s; x = {integ.getX()}\")\n",
    "\n",
    "\n",
    "sp.fns.post_hop_fns[\"update_and_info\"] = post_step_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5de4d87",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Initialize state and time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a569ae5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "x_init = np.array([0.5, -0.5, 0.0, 0.0])\n",
    "t_init = np.timedelta64(0, \"ns\")\n",
    "sp.setTime(t_init)\n",
    "sp.setState(x_init)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c71f917",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Register event to limit step size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88188c58",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "h = np.timedelta64(int(1e7), \"ns\")\n",
    "t = kd.TimedEvent(\"hop_size\", h, lambda _: None, False)\n",
    "t.period = h\n",
    "\n",
    "# register after time has been initialized\n",
    "sp.registerTimedEvent(t)\n",
    "del t\n",
    "\n",
    "print(f\"t = {float(integ.getTime()) / 1e9}s; x = {integ.getX()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5406f59",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Run the simulation\n",
    "\n",
    "The `advanceTo` time below can be increased to lengthen the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dba2fbc",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "sp.advanceTo(2.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ac906a7",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "web_scene.renderToFile(\"render.png\")\n",
    "display(Image(\"render.png\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b6f1e51",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4b6ed25",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def cleanup():\n",
    "    \"\"\"Cleanup the simulation.\"\"\"\n",
    "    global integ, bd1, bd2, bd1_part, bd2_part, root_scene_node, web_scene, proxy_scene\n",
    "    del integ, bd1, bd2, bd1_part, bd2_part, root_scene_node, web_scene, proxy_scene\n",
    "\n",
    "    kc.discard(sp)\n",
    "    cleanup_graphics()\n",
    "\n",
    "    fc = mb.frameContainer()\n",
    "    kc.discard(mb)\n",
    "    kc.discard(fc)\n",
    "\n",
    "\n",
    "atexit.register(cleanup)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75694af9",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Summary\n",
    "You now know how to add visual geometries to a simulation and render to a file."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  },
  "name": "notebook.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
