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

# In[ ]:


import k3d
import math
import numpy as np
import nibabel as nib
from k3d.helpers import download
import ipywidgets as widgets
import vtk
from vtk.util import numpy_support


# 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[ ]:


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

bounds = vti.GetBounds()
x, y, z = vti.GetDimensions()
img = numpy_support.vtk_to_numpy(vti.GetPointData().GetArray(0)).reshape(-1, y, x)


# In[ ]:


tf_editor = k3d.transfer_function_editor()
volume = k3d.volume(img.astype(np.float16))

@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'), (volume, 'color_map'))
_ = widgets.link((tf_editor, 'opacity_function'), (volume, 'opacity_function'))


# In[ ]:


plot = k3d.plot()
plot += volume
tf_editor.display();
plot.display()


# In[ ]:


tf_editor.color_map


# In[ ]:


tf_editor.opacity_function


# In[ ]:


volume.color_map


# In[ ]:


volume.opacity_function


# In[ ]:


k3d.colormaps.basic_color_maps.Jet


# In[ ]: