# Widget combining k3d + ipywidgets for volume rendering.

from IPython.display import display
import matplotlib.pyplot as plt
import k3d

from ipywidgets import Checkbox, IntSlider, Layout, HBox, VBox, Label, Dropdown, FloatSlider, Output
import numpy as np
from scipy.ndimage import maximum_filter
import h5py

# colormap
cmap_default = 'magma'
cmap = np.array(getattr(k3d.matplotlib_color_maps,cmap_default),np.float32)

# Data
with h5py.File("data/CNT_overlap_tomo_missing.h5","r") as f:
    values = f['reconstruction'][:]
    p = np.logical_and(maximum_filter(
        values,
        size = 3,
    ) == values, values > 3e-3).transpose(2,1,0)
    
    values -= values.min()
    values /= values.max()
nx, ny, nz = values.shape

cnt = k3d.volume(
    values, 
    color_map=cmap,
    alpha_coef=20,
    color_range=[0.25,0.95],
    bounds=[0, nz, 0, ny, 0, nx]
)

xyz = np.argwhere(p).astype(np.float32)
points = k3d.points(
    xyz,
    point_size=2.5,
    color=175,
    opacity=0.5,
)

plot = k3d.plot(
    grid_visible=False,
    menu_visibility=False,
    camera_mode='orbit',
    colorbar_object_id=0,
    camera_fov=30,
)
plot += cnt
plot += points

points.visible = False

# Widgets 
sequential_cmaps = [
    'gray','viridis', 'plasma', 'inferno', 'magma', 'cividis','turbo',
    'Purples_r', 'Blues_r', 'Greens_r', 'Oranges_r', 'Reds_r',
    'YlOrBr_r', 'YlOrRd_r', 'OrRd_r', 'PuRd_r', 'RdPu_r', 'BuPu_r',
    'GnBu_r', 'PuBu_r', 'YlGnBu_r', 'PuBuGn_r', 'BuGn_r', 'YlGn_r'
]

def update_colormap(change):
    cmap_string = change['new']
    cmap = np.array(getattr(k3d.matplotlib_color_maps,cmap_string),np.float32)
    cnt.color_map = cmap
    return None

def toggle_vol(change):
    visible = change['new']
    cnt.visible = visible
    return None

def toggle_atoms(change):
    visible = change['new']
    points.visible =visible
    return None

def update_atoms_size(change):
    size = change['new']
    points.point_size = size
    return None

cmap_widget =Dropdown(options=sequential_cmaps,value=cmap_default,description="Colormap",indent=False,layout=Layout(width='175px'))
cmap_widget.observe(update_colormap,names='value')

layout = Layout(width='175px')

vol_checkbox = Checkbox(
    value=True,
    description='Volume',
    indent=True,
    layout=layout,
)
vol_checkbox.observe(toggle_vol,names='value')

atoms_checkbox = Checkbox(
    value=False,
    description='Atoms',
    indent=True,
    layout=layout,
)
atoms_checkbox.observe(toggle_atoms,names='value')

atoms_slider = FloatSlider(
    value=2.5,
    min=1,
    max=10,
    step=0.25,
    description='atoms size',
    continuous_update=False,
    layout=layout,
    readout=False,
)
atoms_slider.observe(update_atoms_size,names='value')

controls_layout = Layout(
    display='flex',
    flex_flow='column',
    align_items='center',
    width='200px'
)

controls = VBox(
    [
        Label("Visual Controls"),
        cmap_widget,
        vol_checkbox,
        atoms_checkbox,
        atoms_slider,
    ],
    layout=controls_layout
)

visualization_layout = Layout(
    display='flex',
    flex_flow='row',
    align_items='center',
    width='630px'
)

out = Output(layout=Layout(width="400px"))
with out:
    plot.display()

display(
    HBox([
        out,
        controls
    ],
        layout=visualization_layout
        )
)