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

# # pythreejs: 3D rendering in the browser 
# 
# ## A Jupyter - threejs bridge
# 
# ## https://github.com/jupyter-widgets/pythreejs
# 
# 
# Pythreejs is a jupyter interactive widget bringing fast WebGL 3d visualization to the Jupyter notebook.
# 
# - Originally authored by Jason Grout, currently maintained by Vidar Tonaas Fauske
# - BSD Licensed
# 
# Pythreejs is *not* a 3d plotting library, it only exposes the threejs scene objects to the Jupyter kernel.
# 
# **Installation:**
# 
# ```bash
# conda install -c conda-forge pythreejs
# ```

# ## This notebook based on [Examples.ipynb](https://github.com/jupyter-widgets/pythreejs/blob/master/examples/Examples.ipynb) 
# 
# One of the many example notebooks at https://github.com/jupyter-widgets/pythreejs/tree/master/examples

# In[ ]:


from pythreejs import *
import numpy as np
from IPython.display import display
from ipywidgets import HTML, Text, Output, VBox
from traitlets import link, dlink


# ## Simple sphere

# In[ ]:


ball = Mesh(geometry=SphereGeometry(radius=1), 
            material=MeshLambertMaterial(color='red'),
            position=[2, 1, 0])

c = PerspectiveCamera(position=[0, 5, 5], up=[0, 1, 0],
                      children=[DirectionalLight(color='white', position=[3, 5, 1], intensity=0.5)])

scene = Scene(children=[ball, c, AmbientLight(color='#777777')])

renderer = Renderer(camera=c, 
                    scene=scene, 
                    controls=[OrbitControls(controlling=c)])
display(renderer)


# Properties of the ball can be updated, with changes reflected immediately in the visualization.

# In[ ]:


ball.scale = (0.5,) * 3


# In[ ]:


import time, math
ball.material.color = '#4400dd'
for i in range(1, 150, 2):
    ball.scale = (i / 100.,) * 3
    ball.position = [math.cos(i / 10.), math.sin(i / 50.), i / 100.]
    time.sleep(.05)


# ## Clickable Surface

# In[ ]:


# Generate surface data:
view_width = 600
view_height = 400
nx, ny = (20, 20)
xmax=1
x = np.linspace(-xmax, xmax, nx)
y = np.linspace(-xmax, xmax, ny)
xx, yy = np.meshgrid(x, y)
z = xx ** 2 - yy ** 2
#z[6,1] = float('nan')


# Generate scene objects from data:
surf_g = SurfaceGeometry(z=list(z[::-1].flat), 
                         width=2 * xmax,
                         height=2 * xmax,
                         width_segments=nx - 1,
                         height_segments=ny - 1)

surf = Mesh(geometry=surf_g,
            material=MeshLambertMaterial(map=height_texture(z[::-1], 'YlGnBu_r')))

surfgrid = SurfaceGrid(geometry=surf_g, material=LineBasicMaterial(color='black'),
                       position=[0, 0, 1e-2])  # Avoid overlap by lifting grid slightly

# Set up picking bojects:
hover_point = Mesh(geometry=SphereGeometry(radius=0.05),
                   material=MeshLambertMaterial(color='green'))

click_picker = Picker(controlling=surf, event='dblclick')
hover_picker = Picker(controlling=surf, event='mousemove')

# Set up scene:
key_light = DirectionalLight(color='white', position=[3, 5, 1], intensity=0.4)
c = PerspectiveCamera(position=[0, 3, 3], up=[0, 0, 1], aspect=view_width / view_height,
                      children=[key_light])

scene = Scene(children=[surf, c, surfgrid, hover_point, AmbientLight(intensity=0.8)])

renderer = Renderer(camera=c, scene=scene,
                    width=view_width, height=view_height,
                    controls=[OrbitControls(controlling=c), click_picker, hover_picker])


# Set up picking responses:
# Add a new marker when double-clicking:
out = Output()
def f(change):
    value = change['new']
    with out:
        print('Clicked on %s' % (value,))
    point = Mesh(geometry=SphereGeometry(radius=0.05), 
                 material=MeshLambertMaterial(color='hotpink'),
                 position=value)
    scene.add(point)

click_picker.observe(f, names=['point'])

# Have marker follow picker point:
link((hover_point, 'position'), (hover_picker, 'point'))

# Show picker point coordinates as a label:
h = HTML()
def g(change):
    h.value = 'Green point at (%.3f, %.3f, %.3f)' % tuple(change['new'])
    h.value += '  Double-click to add marker'
g({'new': hover_point.position})
hover_picker.observe(g, names=['point'])

display(VBox([h, renderer, out]))


# In[ ]:


# when we change the z values of the geometry, we need to also change the height map
surf_g.z = list((-z[::-1]).flat)
surf.material.map = height_texture(-z[::-1])


# # Parametric Functions
# 
# 
# To use the ParametricGeometry class, you need to specify a javascript function as a string. The function should take two parameters that vary between 0 and 1, and return a `new THREE.Vector3(x,y,z)`.
# 
# If you want to build the surface in Python, you'll need to explicitly construct the vertices and faces and build a basic geometry from the vertices and faces.

# In[ ]:


f = """
function f(origu,origv) {
    // scale u and v to the ranges I want: [0, 2*pi]
    var u = 2*Math.PI*origu;
    var v = 2*Math.PI*origv;
    
    var x = Math.sin(u);
    var y = Math.cos(v);
    var z = Math.cos(u+v);
    
    return new THREE.Vector3(x,y,z)
}
"""
surf_g = ParametricGeometry(func=f, slices=16, stacks=16);

surf = Mesh(geometry=surf_g, material=MeshLambertMaterial(color='green', side='FrontSide'))
surf2 = Mesh(geometry=surf_g, material=MeshLambertMaterial(color='yellow', side='BackSide'))
c = PerspectiveCamera(position=[5, 5, 3], up=[0, 0, 1],
                      children=[DirectionalLight(color='white',
                                                 position=[3, 5, 1],
                                                 intensity=0.6)])
scene = Scene(children=[surf, surf2, c, AmbientLight(intensity=0.5)])
renderer = Renderer(camera=c, scene=scene, controls=[OrbitControls(controlling=c)], width=400, height=400)
display(renderer)