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

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)))

%%javascript

// Loading the compiled MathBox bundle.
require.config({
    paths: {
        mathBox: '//cdn.rawgit.com/unconed/mathbox/eaeb8e15/build/mathbox-bundle'
    }
});

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

        func(mathbox);
        
        window.dispatchEvent(new Event('resize'));
    })
}

"""
Explanation: Jupyter meets MathBox2
This notebook contains a few examples of embedded MathBox2 plots.
NOTE that these examples are currently pegged to an old version of Mathbox! It would be lovely to update this to the modern stack.
Boilerplate code
End of explanation
"""


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

"""
Explanation: Simple surface plot
This code snippet shows an 3d surface plot of a function, defined in JS callback.
End of explanation
"""


# 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".

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

"""
Explanation: Feeding data from Python
Drawing JS-defined functions is nice, but what if we'd like to draw some data generated by Python code?
End of explanation
"""


%%javascript
with_mathbox(element, function(mathbox) {
  mathbox.three.element.style.height = '600px';
  mathbox.cartesian().grid({width: 2, opacity: 0.5, axes: [1, 3], origin: [0, -1, 0]});

  // Create a view that uses Stereographic 4d->3d projection, instead of 3d cartesian we used before.
  var view = mathbox.stereographic4({position:[0, 0, 0], scale:[0.5, 0.5, 0.5]});

  // Define Tesseract vertices and edges.
  var edges = [];
  var points = []
  for (var e=-1; e<2; e+=2)
  for (var i=-1; i<2; i+=2)
  for (var j=-1; j<2; j+=2)
  for (var k=-1; k<2; k+=2) {
    edges.push([i, j, k, e])
    edges.push([i, j, e, k])
    edges.push([i, e, j, k])
    edges.push([e, i, j, k])
    points.push([i, j, k, e])
  }

  view.matrix({width:edges.length/2, height:2, data:edges, live: false})
  .transpose({order:"yx", id:"edges"})
  .array({data:points, id:"points"})
  .clock({speed:0.25})
  // Animate rotation in 4d space.
  .transform4({}, {matrix:function(t) {
     var c = Math.cos(t), s = Math.sin(t);
     return [c, 0, 0,-s,
             0, 1, 0, 0,
             0, 0, 1, 0,
             s, 0, 0, c];
     }})  
    // Draw points. 
    .point({size:8, points:"#points"})
    // Label them. 
    .format({live:false, expr:(x, y, z, w)=>{
      return x+", "+y+", "+z+", "+w;
    }}).label({size:16, depth:0.5, outline:0.5})
    // This line linearly interpolates our edges in 4d space before doing projection,
    // which gives us nice curved edges.
    .lerp({width:32, source:"#edges"})
    .line({color:0x3090FF, depth:1.0, width:4});
})

"""
Explanation: Tesseract rotation
Let's draw something more exotic!
End of explanation
"""


%%javascript
with_mathbox(element, function(mathbox) {

mathbox.three.camera.position.set(-0.1, 1, 1.5);
var W = 512, H = 256;
mathbox
  .rtt({width:W, height:H, type:"float", id:"rtt"}) // offscreen rendering
    // main simulation shader code
    .shader({code:`
      uniform vec2 dataSize;
      uniform vec2 spot;
      uniform vec2 fk;
      vec4 getsample(vec2 p);
      vec2 sample(vec2 p) {
        return getsample(mod(p, dataSize)).xy;
      }
      vec4 main(vec2 p) {
        if (length(spot-p)<2.0) {
          return vec4(0.0, 0.5, 0.0, 0.0);
        }
        float f = fk.x, k = fk.y;
        const vec2 dx=vec2(1.,0.0), dy=vec2(0.0,1.);
        vec2 v = sample(p);
        vec2 lap = sample(p+dx)+sample(p-dx)+sample(p+dy)+sample(p-dy)-4.0*v;
        float rate = v.x * v.y * v.y;
        vec2 dv = vec2(0.2, 0.1)*lap + vec2(-rate, rate);
        dv += vec2(f * (1.0 - v.x), -(f + k) * v.y);
        v = clamp(v+dv, 0.0, 1.0);
        return vec4(v, 0.0, 0.0);
      }`, fk:[0.034, 0.056]}, {spot:(t)=>[(t*0.02+0.75)%1*W, (t*0.12+0.5)%1*H]})
      .play({  // animate Gray-Scott reaction-diffusion parameters
        loop: true, to:4, pace:3.0,
        script:[
          {fk:[0.034, 0.056]}, 
          {fk:[0.029, 0.057]},
          {fk:[0.014, 0.054]},
          {fk:[0.025, 0.060]},
          {fk:[0.034, 0.056]}]})
      .resample({indices:2}).compose() // this triggers Render-to-Texture pass
   .end() // back from offscreen rendering

   .cartesian({}, {rotation:(t)=>[t*0.1+1.5, 0, 0]})
       // shader to compute surface colors
       .shader({code:`
          vec4 sample(vec4 p);
          vec4 main(vec4 p) {
            float v = sample(p).y;
            return vec4(0.5+v, 0.5, 0.5, 1.0);
          }
         `}).resample()
       // shader to compute 3d positions of torus vertices
       .shader({code:`
         uniform vec4 dataSize;
         const float pi = 3.141593;
         vec4 sample(vec4 p);
         vec4 main(vec4 p) {
           float v = sample(p).y;
           vec2 pq = p.xy/(dataSize.xy-1.0)*2.0*pi;
           float r = v*0.2 + 0.3;
           float a = 0.7 + r*cos(pq.y);
           return vec4(a*cos(pq.x), a*sin(pq.x), r*sin(pq.y), 0.0);
         }`}).resample({source:'#rtt'})
       // draw the torus!
       .surface({shaded:true, closedX:true, colors:'<<', color:"#fff"});

   // this hack triggers RTT pass multiple times per frame for faster simulation
   mathbox.three.on('update', function(){
     var rtt = mathbox.select('rtt')[0].controller.rtt;
     for (var i=0; i<5; ++i)
       rtt.render()
   })
})

"""
Explanation: Gray-Scott Reaction-Diffusion 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.
End of explanation
"""


