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

# ## Volume Mesh
# 
# Mesh is an object which displays triangles in 3d.
# An scalar can be displayed on the mesh using color map.

# In[ ]:


import vtk
import k3d
import numpy as np
import math
import ipywidgets as widgets
from vtk.util import numpy_support
from k3d.helpers import quad, download

filename = download('https://github.com/marcomusy/vtkplotter-examples/raw/master/vtkplotter_examples/data/embryo.slc')
reader = vtk.vtkSLCReader()
reader.SetFileName(filename)
reader.Update()
vti = reader.GetOutput()

x, y, z = vti.GetDimensions()
volume_data = numpy_support.vtk_to_numpy(vti.GetPointData().GetArray(0)).reshape(-1, y, x).astype(np.float32)


# In[ ]:


basic_color_maps = [(attr, getattr(k3d.basic_color_maps, attr)) for attr in dir(k3d.basic_color_maps) if not attr.startswith('__')]
paraview_color_maps = [(attr, getattr(k3d.paraview_color_maps, attr)) for attr in dir(k3d.paraview_color_maps) if not attr.startswith('__')]
matplotlib_color_maps = [(attr, getattr(k3d.matplotlib_color_maps, attr)) for attr in dir(k3d.matplotlib_color_maps) if not attr.startswith('__')]
colormaps = basic_color_maps + paraview_color_maps + matplotlib_color_maps


# In[ ]:


vertices, indices = quad(20.0, 20.0)

plot = k3d.plot()

obj = k3d.mesh(vertices, indices, volume=volume_data, side='double', volume_bounds=[-10, 10, -10, 10, -10, 10])

model_matrix = {}
for t in np.linspace(0, 2* np.pi, 100):
    obj.transform.rotation = [t, math.sin(t), math.cos(t), 1]
    model_matrix[str(t)] = obj.model_matrix
    
obj.model_matrix = model_matrix    
plot += obj

plot.display()


# In[ ]:


tf_editor = k3d.transfer_function_editor()

@widgets.interact(x=widgets.Dropdown(options=colormaps, description='ColorMap:'))
def g(x):
    tf_editor.color_map = np.array(x, dtype=np.float32)

_ = widgets.link((tf_editor, 'color_map'), (obj, 'color_map'))
_ = widgets.link((tf_editor, 'opacity_function'), (obj, 'opacity_function'))

tf_editor.display()


# In[ ]:


tf_editor.opacity_function = [
    0, 0, 
    0.04, 0,
    0.1, 1, 
    1,1
]


# In[ ]:


plot.start_auto_play()


# In[ ]:


plot.stop_auto_play()


# In[ ]:


f = open('./volume.html', 'w', encoding='UTF-8')
f.write(plot.get_snapshot(9, 'K3DInstance.startAutoPlay();'))
f.close()


# In[ ]: