#!/usr/bin/env python
# coding: utf-8

# # Jupyter meets MathBox2
# 
# This notebook contains a few examples of embedded [MathBox2](https://github.com/unconed/mathbox) plots.

# ## Boilerplate code

# In[1]:


import json
import numpy as np
from IPython.display import HTML, Javascript, display


# In[2]:


def json_numpy_serialzer(o):
    '''Helper function to serialize NumPy arrays.'''
    if isinstance(o, np.ndarray):
        return o.tolist()
    raise TypeError("{} of type {} is not JSON serializable".format(repr(o), type(o)))

def jsglobal(**params):
    '''Populate JS global namespace with provided Python obejcts.'''
    code = [];
    for name, value in params.items():
        jsdata = json.dumps(value, default=json_numpy_serialzer)
        code.append("window.{} = {};".format(name, jsdata))
    display(Javascript("\n".join(code)))


# In[3]:


get_ipython().run_cell_magic('javascript', '', '\n// Loading the compiled MathBox bundle.\nrequire.config({\n    paths: {\n        mathBox: \'//cdn.rawgit.com/unconed/mathbox/eaeb8e15/build/mathbox-bundle\'\n    }\n});\n\n// Helper function that setups WebGL context and initializes MathBox.\nwindow.with_mathbox = function(element, func) {\n    require([\'mathBox\'], function(){\n        var mathbox = mathBox({\n          plugins: [\'core\', \'controls\', \'cursor\', \'mathbox\'],\n          controls: { klass: THREE.OrbitControls },\n          mathbox: {inspect: false},\n          element: element[0],\n          loop: {start: false},\n            \n        });\n        var three = mathbox.three;\n        three.renderer.setClearColor(new THREE.Color(0xFFFFFF), 1.0);\n        three.camera.position.set(-1, 1, 2);\n        three.controls.noKeys = true;\n        \n        three.element.style.height = "400px";\n        three.element.style.width = "100%";\n        \n        function isInViewport(element) {\n          var rect = element.getBoundingClientRect();\n          var html = document.documentElement;\n          var w = window.innerWidth || html.clientWidth;\n          var h = window.innerHeight || html.clientHeight;\n          return rect.top < h && rect.left < w && rect.bottom > 0 && rect.right > 0;\n        }\n        \n        // Running update/render loop only for visible plots.\n        var intervalId = setInterval(function(){\n            if (three.element.offsetParent === null) {\n                clearInterval(intervalId);\n                three.destroy();\n                return;\n            }\n            var visible = isInViewport(three.canvas);\n            if (three.Loop.running != visible) {\n                visible? three.Loop.start() : three.Loop.stop();\n            }\n        }, 100);\n\n        func(mathbox);\n        \n        window.dispatchEvent(new Event(\'resize\'));\n    })\n}\n')


# ## Simple surface plot
# 
# This code snippet shows an 3d surface plot of a function, defined in JS callback.

# In[4]:


get_ipython().run_cell_magic('javascript', '', "with_mathbox(element, function(mathbox) {\n   mathbox.cartesian({},{rotation:(t)=>[0, t*0.1, 0]}) // Setup rotating the coordinate frame.\n     .grid({axes: [1, 3]})                             // Add a grid to it.\n     .area({width:50, height:50,                       // This defines 2D data source, sampled from JS callback \n            expr: function(emit, x, y, i, j, t){\n              var r = Math.sqrt(x*x+y*y);\n              var z = Math.sin(r*10-t*0.5)*0.2 + 0.3;\n              emit(x, z, y);\n           }})\n     .surface({color:'#AAA', shaded:true})             // Adding surface primitives, that draw data provided by \n     .surface({color:'#55A', lineX:true, lineY:true, fill:false, zBias:1}); // the last defined datasource.\n})\n")


# ## Feeding data from Python
# 
# Drawing JS-defined functions is nice, but what if we'd like to draw some data generated by Python code?

# In[5]:


# Make an array of 3d points.
np.random.seed(123)
t = np.linspace(0, 2*np.pi, 1000)
x, y, z = np.sin(t*10), np.sin(t*20), np.sin(t*30+0.5)
pos = np.vstack([x, y, z]).T
pos += np.random.normal(size=pos.shape)*0.02

jsglobal(POS=pos) # Pass this array to JS-world as a global variable "POS".


# In[6]:


get_ipython().run_cell_magic('javascript', '', 'with_mathbox(element, function(mathbox) {\n    mathbox.cartesian({},{rotation:(t)=>[0, t*0.1, 0]})\n      .grid({axes: [1, 3]})\n      // Now we can see the data on JS side!\n      .array({data:POS, channels:3, live:false})\n      .point({color:"#55a"})\n      .line({width:1.0})\n})\n')


# ## Tesseract rotation
# 
# Let's draw something more exotic!

# In[7]:


get_ipython().run_cell_magic('javascript', '', 'with_mathbox(element, function(mathbox) {\n  mathbox.three.element.style.height = \'600px\';\n  mathbox.cartesian().grid({width: 2, opacity: 0.5, axes: [1, 3], origin: [0, -1, 0]});\n\n  // Create a view that uses Stereographic 4d->3d projection, instead of 3d cartesian we used before.\n  var view = mathbox.stereographic4({position:[0, 0, 0], scale:[0.5, 0.5, 0.5]});\n\n  // Define Tesseract vertices and edges.\n  var edges = [];\n  var points = []\n  for (var e=-1; e<2; e+=2)\n  for (var i=-1; i<2; i+=2)\n  for (var j=-1; j<2; j+=2)\n  for (var k=-1; k<2; k+=2) {\n    edges.push([i, j, k, e])\n    edges.push([i, j, e, k])\n    edges.push([i, e, j, k])\n    edges.push([e, i, j, k])\n    points.push([i, j, k, e])\n  }\n\n  view.matrix({width:edges.length/2, height:2, data:edges, live: false})\n  .transpose({order:"yx", id:"edges"})\n  .array({data:points, id:"points"})\n  .clock({speed:0.25})\n  // Animate rotation in 4d space.\n  .transform4({}, {matrix:function(t) {\n     var c = Math.cos(t), s = Math.sin(t);\n     return [c, 0, 0,-s,\n             0, 1, 0, 0,\n             0, 0, 1, 0,\n             s, 0, 0, c];\n     }})  \n    // Draw points. \n    .point({size:8, points:"#points"})\n    // Label them. \n    .format({live:false, expr:(x, y, z, w)=>{\n      return x+", "+y+", "+z+", "+w;\n    }}).label({size:16, depth:0.5, outline:0.5})\n    // This line linearly interpolates our edges in 4d space before doing projection,\n    // which gives us nice curved edges.\n    .lerp({width:32, source:"#edges"})\n    .line({color:0x3090FF, depth:1.0, width:4});\n})\n')


# ## Gray-Scott [Reaction-Diffusion](http://mrob.com/pub/comp/xmorphia/) on Torus (GLSL)
# 
# Mathbox allows to inject custom GLSL functions into its dataflow pipelines. It also exposes Render-to-Texture functionality to make pre-CUDA style GPU computing possible with minimal amounts of boilerplate code. This sample computers reaction-diffusion simulation in an offscreen texture, and then uses this texture to displace and colorize torus surface.

# In[11]:


get_ipython().run_cell_magic('javascript', '', 'with_mathbox(element, function(mathbox) {\n\nmathbox.three.camera.position.set(-0.1, 1, 1.5);\nvar W = 512, H = 256;\nmathbox\n  .rtt({width:W, height:H, type:"float", id:"rtt"}) // offscreen rendering\n    // main simulation shader code\n    .shader({code:`\n      uniform vec2 dataSize;\n      uniform vec2 spot;\n      uniform vec2 fk;\n      vec4 getsample(vec2 p);\n      vec2 sample(vec2 p) {\n        return getsample(mod(p, dataSize)).xy;\n      }\n      vec4 main(vec2 p) {\n        if (length(spot-p)<2.0) {\n          return vec4(0.0, 0.5, 0.0, 0.0);\n        }\n        float f = fk.x, k = fk.y;\n        const vec2 dx=vec2(1.,0.0), dy=vec2(0.0,1.);\n        vec2 v = sample(p);\n        vec2 lap = sample(p+dx)+sample(p-dx)+sample(p+dy)+sample(p-dy)-4.0*v;\n        float rate = v.x * v.y * v.y;\n        vec2 dv = vec2(0.2, 0.1)*lap + vec2(-rate, rate);\n        dv += vec2(f * (1.0 - v.x), -(f + k) * v.y);\n        v = clamp(v+dv, 0.0, 1.0);\n        return vec4(v, 0.0, 0.0);\n      }`, fk:[0.034, 0.056]}, {spot:(t)=>[(t*0.02+0.75)%1*W, (t*0.12+0.5)%1*H]})\n      .play({  // animate Gray-Scott reaction-diffusion parameters\n        loop: true, to:4, pace:3.0,\n        script:[\n          {fk:[0.034, 0.056]}, \n          {fk:[0.029, 0.057]},\n          {fk:[0.014, 0.054]},\n          {fk:[0.025, 0.060]},\n          {fk:[0.034, 0.056]}]})\n      .resample({indices:2}).compose() // this triggers Render-to-Texture pass\n   .end() // back from offscreen rendering\n\n   .cartesian({}, {rotation:(t)=>[t*0.1+1.5, 0, 0]})\n       // shader to compute surface colors\n       .shader({code:`\n          vec4 sample(vec4 p);\n          vec4 main(vec4 p) {\n            float v = sample(p).y;\n            return vec4(0.5+v, 0.5, 0.5, 1.0);\n          }\n         `}).resample()\n       // shader to compute 3d positions of torus vertices\n       .shader({code:`\n         uniform vec4 dataSize;\n         const float pi = 3.141593;\n         vec4 sample(vec4 p);\n         vec4 main(vec4 p) {\n           float v = sample(p).y;\n           vec2 pq = p.xy/(dataSize.xy-1.0)*2.0*pi;\n           float r = v*0.2 + 0.3;\n           float a = 0.7 + r*cos(pq.y);\n           return vec4(a*cos(pq.x), a*sin(pq.x), r*sin(pq.y), 0.0);\n         }`}).resample({source:\'#rtt\'})\n       // draw the torus!\n       .surface({shaded:true, closedX:true, colors:\'<<\', color:"#fff"});\n\n   // this hack triggers RTT pass multiple times per frame for faster simulation\n   mathbox.three.on(\'update\', function(){\n     var rtt = mathbox.select(\'rtt\')[0].controller.rtt;\n     for (var i=0; i<5; ++i)\n       rtt.render()\n   })\n})\n')