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

# In[1]:


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


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


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


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


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


# In[6]:


tf_editor.color_map


# In[7]:


tf_editor.opacity_function


# In[8]:


volume.color_map


# In[9]:


volume.opacity_function


# In[10]:


k3d.colormaps.basic_color_maps.Jet


# In[11]:


tf_editor


# In[13]:


a = widgets.link((tf_editor, 'color_map'), (volume, 'color_map'))


# In[ ]: