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

# In[ ]:


import k3d
import numpy as np
from math import sqrt, sin, cos
from skimage.measure import label

width = height = length = 100

def r(x, y, z):
    r = sqrt((x - width / 2) * (x - width / 2) + (y - height / 2) * (y - height / 2) + (z - length / 2) * (z - length / 2))
    r += sin(x / 2) * 3
    r += cos(y / 10) * 5
    
    return r

def f(x, y, z):
    return 0 if r(x, y, z) > width / 2 else (1 if y + sin(x / 20) * 10 > height / 2 else 2)

def on_click(x, y, z):
    coords = [x, y, z]
    missing = None
    missing_count = 0
    
    for idx, v in enumerate([np.array([0,0,1]), np.array([0,1,0]), np.array([1,0,0])]):
        axis = coords[idx]        
        
        if axis == 0 or (obj.voxels[tuple(np.array([z, y, x]) - v)] == 0):
            missing = -v
            missing_count += 1

        if axis == obj.voxels.shape[2 - idx] - 1 or obj.voxels[tuple(np.array([z, y, x]) + v)] == 0:
            missing = v
            missing_count += 1 
    
    if missing_count == 1:        
        if missing[0] != 0:
            slice = obj.voxels[z, :, :]
            mask = label(slice, connectivity = 2)            
            slice[mask == mask[y, x]] = plot.voxel_paint_color
        if missing[1] != 0:
            slice = obj.voxels[:, y, :]
            mask = label(slice, connectivity = 2)
            slice[mask == mask[z, x]] = plot.voxel_paint_color
        if missing[2] != 0:
            slice = obj.voxels[:, :, x]
            mask = label(slice, connectivity = 2)
            slice[mask == mask[z, y]] = plot.voxel_paint_color
            
        obj.push_data('voxels')
        
color_map = (0xffff00, 0xff0000, 0x00ff00)
voxels = np.array([[[f(x, y, z) for x in range(width)] for y in range(height)] for z in range(length)])

obj = k3d.voxels(voxels.astype(np.uint8), color_map)
plot = k3d.plot()
plot += obj

obj.click_callback = on_click
plot.voxel_paint_color = 3


# In[ ]:


# select Controls->Mode->Callback and click on object wall
plot.display()


# In[ ]:


plot.camera_auto_fit = False


# In[ ]: