{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jupyter backend demo: molecular viewer in the notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example shows how vispy's low-level gloo interface can be used to display a VisPy Canvas in a notebook. The higher level Scene and Plotting API are also supported. By default, vispy will detect that it is being run in a notebook and load the \"jupyter_rfb\" backend automatically. This does require that the \"jupyter_rfb\" package be installed.\n",
    "\n",
    "Due to the way widgets work in Jupyter, you can display the Canvas multiple times, but they will continue to represent the same object. In the case of timers, they will continue to run in the backgroudn until they are stopped, even if the output of your notebook cell is cleared and no canvas is visible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from vispy import gloo\n",
    "from vispy import app\n",
    "from vispy.util.transforms import perspective, translate, rotate\n",
    "from vispy.io import load_data_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vertex = \"\"\"\n",
    "#version 120\n",
    "\n",
    "uniform mat4 u_model;\n",
    "uniform mat4 u_view;\n",
    "uniform mat4 u_projection;\n",
    "uniform vec3 u_light_position;\n",
    "uniform vec3 u_light_spec_position;\n",
    "\n",
    "attribute vec3  a_position;\n",
    "attribute vec3  a_color;\n",
    "attribute float a_radius;\n",
    "\n",
    "varying vec3  v_color;\n",
    "varying vec4  v_eye_position;\n",
    "varying float v_radius;\n",
    "varying vec3  v_light_direction;\n",
    "\n",
    "void main (void) {\n",
    "    v_radius = a_radius;\n",
    "    v_color = a_color;\n",
    "\n",
    "    v_eye_position = u_view * u_model * vec4(a_position,1.0);\n",
    "    v_light_direction = normalize(u_light_position);\n",
    "    float dist = length(v_eye_position.xyz);\n",
    "\n",
    "    gl_Position = u_projection * v_eye_position;\n",
    "\n",
    "    // stackoverflow.com/questions/8608844/...\n",
    "    //  ... resizing-point-sprites-based-on-distance-from-the-camera\n",
    "    vec4  proj_corner = u_projection * vec4(a_radius, a_radius, v_eye_position.z, v_eye_position.w);  // # noqa\n",
    "    gl_PointSize = 512.0 * proj_corner.x / proj_corner.w;\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "fragment = \"\"\"\n",
    "#version 120\n",
    "\n",
    "uniform mat4 u_model;\n",
    "uniform mat4 u_view;\n",
    "uniform mat4 u_projection;\n",
    "uniform vec3 u_light_position;\n",
    "uniform vec3 u_light_spec_position;\n",
    "\n",
    "varying vec3  v_color;\n",
    "varying vec4  v_eye_position;\n",
    "varying float v_radius;\n",
    "varying vec3  v_light_direction;\n",
    "void main()\n",
    "{\n",
    "    // r^2 = (x - x0)^2 + (y - y0)^2 + (z - z0)^2\n",
    "    vec2 texcoord = gl_PointCoord* 2.0 - vec2(1.0);\n",
    "    float x = texcoord.x;\n",
    "    float y = texcoord.y;\n",
    "    float d = 1.0 - x*x - y*y;\n",
    "    if (d <= 0.0)\n",
    "        discard;\n",
    "\n",
    "    float z = sqrt(d);\n",
    "    vec4 pos = v_eye_position;\n",
    "    pos.z += v_radius*z;\n",
    "    vec3 pos2 = pos.xyz;\n",
    "    pos = u_projection * pos;\n",
    "    //gl_FragDepth = 0.5*(pos.z / pos.w)+0.5;\n",
    "    vec3 normal = vec3(x,y,z);\n",
    "    float diffuse = clamp(dot(normal, v_light_direction), 0.0, 1.0);\n",
    "\n",
    "    // Specular lighting.\n",
    "    vec3 M = pos2.xyz;\n",
    "    vec3 O = v_eye_position.xyz;\n",
    "    vec3 L = u_light_spec_position;\n",
    "    vec3 K = normalize(normalize(L - M) + normalize(O - M));\n",
    "    // WARNING: abs() is necessary, otherwise weird bugs may appear with some\n",
    "    // GPU drivers...\n",
    "    float specular = clamp(pow(abs(dot(normal, K)), 40.), 0.0, 1.0);\n",
    "    vec3 v_light = vec3(1., 1., 1.);\n",
    "    gl_FragColor.rgba = vec4(.15*v_color + .55*diffuse * v_color\n",
    "                        + .35*specular * v_light, 1.0);\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Canvas(app.Canvas):\n",
    "\n",
    "    def __init__(self):\n",
    "        app.Canvas.__init__(self, title='Molecular viewer',\n",
    "                            keys='interactive', size=(640, 480))\n",
    "        self.ps = self.pixel_scale\n",
    "\n",
    "        self.translate = 40\n",
    "        self.program = gloo.Program(vertex, fragment)\n",
    "        self.view = translate((0, 0, -self.translate))\n",
    "        self.model = np.eye(4, dtype=np.float32)\n",
    "        self.projection = np.eye(4, dtype=np.float32)\n",
    "\n",
    "        self.apply_zoom()\n",
    "\n",
    "        fname = load_data_file('molecular_viewer/micelle.npz')\n",
    "        self.load_molecule(fname)\n",
    "        self.load_data()\n",
    "\n",
    "        self.theta = 0\n",
    "        self.phi = 0\n",
    "\n",
    "        gloo.set_state(depth_test=True, clear_color='black')\n",
    "        self.timer = app.Timer('auto', connect=self.on_timer, start=True)\n",
    "\n",
    "        self.show()\n",
    "\n",
    "    def load_molecule(self, fname):\n",
    "        molecule = np.load(fname)['molecule']\n",
    "        self._nAtoms = molecule.shape[0]\n",
    "\n",
    "        # The x,y,z values store in one array\n",
    "        self.coords = molecule[:, :3]\n",
    "\n",
    "        # The array that will store the color and alpha scale for all the atoms\n",
    "        self.atomsColours = molecule[:, 3:6]\n",
    "\n",
    "        # The array that will store the scale for all the atoms.\n",
    "        self.atomsScales = molecule[:, 6]\n",
    "\n",
    "    def load_data(self):\n",
    "        n = self._nAtoms\n",
    "\n",
    "        data = np.zeros(n, [('a_position', np.float32, 3),\n",
    "                            ('a_color', np.float32, 3),\n",
    "                            ('a_radius', np.float32)])\n",
    "\n",
    "        data['a_position'] = self.coords\n",
    "        data['a_color'] = self.atomsColours\n",
    "        data['a_radius'] = self.atomsScales*self.ps\n",
    "\n",
    "        self.program.bind(gloo.VertexBuffer(data))\n",
    "\n",
    "        self.program['u_model'] = self.model\n",
    "        self.program['u_view'] = self.view\n",
    "        self.program['u_light_position'] = 0., 0., 2.\n",
    "        self.program['u_light_spec_position'] = -5., 5., -5.\n",
    "\n",
    "    def on_key_press(self, event):\n",
    "        if event.text == ' ':\n",
    "            if self.timer.running:\n",
    "                self.timer.stop()\n",
    "            else:\n",
    "                self.timer.start()\n",
    "\n",
    "    def on_timer(self, event):\n",
    "        self.theta += .25\n",
    "        self.phi += .25\n",
    "        self.model = np.dot(rotate(self.theta, (0, 0, 1)),\n",
    "                            rotate(self.phi, (0, 1, 0)))\n",
    "        self.program['u_model'] = self.model\n",
    "        self.update()\n",
    "\n",
    "    def on_resize(self, event):\n",
    "        width, height = event.physical_size\n",
    "        gloo.set_viewport(0, 0, width, height)\n",
    "        self.projection = perspective(25.0, width / float(height), 2.0, 100.0)\n",
    "        self.program['u_projection'] = self.projection\n",
    "\n",
    "    def apply_zoom(self):\n",
    "        width, height = self.physical_size\n",
    "        gloo.set_viewport(0, 0, width, height)\n",
    "        self.projection = perspective(25.0, width / float(height), 2.0, 100.0)\n",
    "        self.program['u_projection'] = self.projection\n",
    "\n",
    "    def on_mouse_wheel(self, event):\n",
    "        self.translate -= event.delta[1]\n",
    "        self.translate = max(-1, self.translate)\n",
    "        self.view = translate((0, 0, -self.translate))\n",
    "\n",
    "        self.program['u_view'] = self.view\n",
    "        self.update()\n",
    "\n",
    "    def on_draw(self, event):\n",
    "        gloo.clear()\n",
    "        self.program.draw('points')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = Canvas()\n",
    "c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To pause the rotation you can either hit the Spacebar or manually call the `stop(`)/`start()` methods on the timer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c.timer.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c.timer.start()"
   ]
  }
 ],
 "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}