# Widget combining pyvista + ipywidgets for volume rendering.

from IPython.display import display
import matplotlib.pyplot as plt
import pyvista as pv
import vtk
from ipywidgets import Checkbox, IntSlider, Layout, HBox, VBox, Label, Dropdown, FloatSlider
import numpy as np
from scipy.ndimage import maximum_filter
import h5py

cmap_default = 'magma'
cmap = pv.LookupTable(cmap_default)

# Data
with h5py.File("data/CNT_overlap_tomo_missing.h5","r") as f:
    values = f['reconstruction'][:]
nx, ny, nz = values.shape

# atomic coordinates
p = np.logical_and(maximum_filter(
    values,
    size = 3,
) == values, values > 3e-3)
xyz = np.argwhere(p)
sig = values[xyz[:,0],xyz[:,1],xyz[:,2]]
xyz = xyz.astype('float')

grid = pv.ImageData(
    dimensions=[nx+1,ny+1,nz+1],
    origin=(0,0,0),
    spacing=(1,1,1),
)
grid.cell_data["values"] = values.flatten(order="F")

pl = pv.Plotter()

vol = pl.add_volume(
    grid,
    name='volume',
    opacity="sigmoid",
    cmap=cmap_default,
    show_scalar_bar=False,
)

atoms = pl.add_points(
    xyz, 
    render_points_as_spheres=True,
    color=cmap(cmap.n_values*85//100)[:-1],
    ambient=0.125,
    point_size = 10.0,
)
atoms.SetVisibility(False)

grid_slice_x = grid.slice(
    normal='x',
    origin=(nx//2,ny//2,nz//2),
)

grid_slice_y = grid.slice(
    normal='y',
    origin=(nx//2,ny//2,nz//2),
)

grid_slice_z = grid.slice(
    normal='z',
    origin=(nx//2,ny//2,nz//2),
)

volume_clip_plane = pl.add_volume_clip_plane(
    vol,
    normal='x',
    origin=(0,ny//2,nz//2),
)
volume_clip_plane.Off()

slice_x = pl.add_mesh(
    grid_slice_x,
    show_scalar_bar=False,
    cmap=cmap_default,
    opacity=0.875,
    name="slice-x",
    render=False,
)

slice_y = pl.add_mesh(
    grid_slice_y,
    show_scalar_bar=False,
    cmap=cmap_default,
    opacity=0.875,
    name="slice-y",
    render=False,
)

slice_z = pl.add_mesh(
    grid_slice_z,
    show_scalar_bar=False,
    cmap=cmap_default,
    opacity=0.875,
    name="slice-z",
    render=False,
)

slices = [slice_x,slice_y,slice_z]

for slice in slices:
    slice.SetVisibility(False)

pl.camera.zoom(0.875)

viewer = pl.show(
    return_viewer = True,
    interactive_update=True,
    window_size=(440,440),
    jupyter_kwargs={'collapse_menu':True},
)

# 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 = pv.LookupTable(cmap_string)
    cf = vtk.vtkColorTransferFunction()
    
    for ii in range(cmap.n_values):
        cf.AddRGBPoint(ii, *cmap(ii/(cmap.n_values-1))[:-1])
    
    vol.prop.SetColor(0,cf)
    atoms.prop.SetColor(cmap(cmap.n_values*85//100)[:-1])

    cx = nx - cx_slider.value
    cy = cy_slider.value
    cz = cy_slider.value
    
    pl.remove_actor(slices[0])
    grid_slice_x = grid.slice(
        normal='x',
        origin=(cx,cy,cz),
    )
    
    slices[0] = pl.add_mesh(
        grid_slice_x,
        show_scalar_bar=False,
        cmap=cmap_string,
        opacity=0.875,
        name="slice-x",
        render=False,
    )
    slices[0].SetVisibility(cx_checkbox.value)

    pl.remove_actor(slices[1])   
    grid_slice_y = grid.slice(
        normal='y',
        origin=(cx,cy,cz),
    )

    slices[1] = pl.add_mesh(
        grid_slice_y,
        show_scalar_bar=False,
        cmap=cmap_string,
        opacity=0.875,
        name="slice-y",
        render=False,
    )
    slices[1].SetVisibility(cy_checkbox.value)

    pl.remove_actor(slices[2])   
    grid_slice_z = grid.slice(
        normal='z',
        origin=(cx,cy,cz),
    )

    slices[2] = pl.add_mesh(
        grid_slice_z,
        show_scalar_bar=False,
        cmap=cmap_string,
        opacity=0.875,
        name="slice-z",
        render=False,
    )
    slices[2].SetVisibility(cz_checkbox.value)

    pl.update(0.1)
    return None
    
def update_x_slice(change):
    if cx_checkbox.value:
        cmap = cmap_widget.value
        offset = nx-change['new']
        cy = cy_slider.value
        cz = cz_slider.value
        
        pl.remove_actor(slices[0])
    
        grid_slice_x = grid.slice(
            normal='x',
            origin=(offset,cy,cz),
        )
        
        slices[0] = pl.add_mesh(
            grid_slice_x,
            show_scalar_bar=False,
            cmap=cmap,
            opacity=0.875,
            name="slice-x",
            render=True,
        )
        pl.update(0.1)
    return None

def update_y_slice(change):
    if cy_checkbox.value:
        cmap = cmap_widget.value
        offset = change['new']
        cx = nx - cx_slider.value
        cz = cz_slider.value
        
        pl.remove_actor(slices[1])
    
        grid_slice_y = grid.slice(
            normal='y',
            origin=(cx,offset,cz),
        )
        
        slices[1] = pl.add_mesh(
            grid_slice_y,
            show_scalar_bar=False,
            cmap=cmap,
            opacity=0.875,
            name="slice-y",
            render=True,
        )
        pl.update(0.1)
    return None

def update_z_slice(change):
    if cz_checkbox.value:
        cmap = cmap_widget.value
        offset = change['new']
        cx = nx - cx_slider.value
        cy = cy_slider.value
        
        pl.remove_actor(slices[2])
    
        grid_slice_z = grid.slice(
            normal='z',
            origin=(cx,cy,offset),
        )
        
        slices[2] = pl.add_mesh(
            grid_slice_z,
            show_scalar_bar=False,
            cmap=cmap,
            opacity=0.875,
            name="slice-z",
            render=True,
        )
        pl.update(0.1)
    return None

def toggle_x_slice(change):
    visible = change['new']
    slices[0].SetVisibility(visible)
    pl.update(0.1)
    return None

def toggle_y_slice(change):
    visible = change['new']
    slices[1].SetVisibility(visible)
    pl.update(0.1)
    return None

def toggle_z_slice(change):
    visible = change['new']
    slices[2].SetVisibility(visible)
    pl.update(0.1)
    return None

def toggle_clip_plane(change):
    visible = change['new']
    if visible:
        volume_clip_plane.On()
    else:
        volume_clip_plane.Off()
    pl.update(0.1)
    return None

def toggle_vol(change):
    visible = change['new']
    vol.SetVisibility(visible)
    pl.update(0.1)
    return None

def toggle_atoms(change):
    visible = change['new']
    atoms.SetVisibility(visible)
    pl.update(0.1)
    return None

def update_atoms_size(change):
    size = change['new']
    atoms.prop.SetPointSize(size)
    pl.update(0.1)
    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')

checkbox_layout = Layout(width='175px')
slider_layout = Layout(width='175px')

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

clip_checkbox = Checkbox(
    value=False,
    description='Clip Plane',
    indent=True,
    layout=checkbox_layout,
)
clip_checkbox.observe(toggle_clip_plane,names='value')

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

atoms_slider = FloatSlider(
    value=10,
    min=5,
    max=20,
    step=0.125,
    description='atoms size',
    continuous_update=False,
    layout=slider_layout,
    readout=False,
)
atoms_slider.observe(update_atoms_size,names='value')

cx_checkbox = Checkbox(
    value=False,
    description='x axis slice',
    indent=True,
    layout=checkbox_layout,
)
cx_checkbox.observe(toggle_x_slice,names='value')

cy_checkbox = Checkbox(
    value=False,
    description='y axis slice',
    indent=True,
    layout=checkbox_layout,
)
cy_checkbox.observe(toggle_y_slice,names='value')

cz_checkbox = Checkbox(
    value=False,
    description='z axis slice',
    indent=True,
    layout=checkbox_layout,
)
cz_checkbox.observe(toggle_z_slice,names='value')

cx_slider = IntSlider(
    value=nx//2,
    min=0,
    max=nx-1,
    description='x index',
    continuous_update=False,
    layout=slider_layout,
    readout=False,
)
cx_slider.observe(update_x_slice,names='value')

cy_slider = IntSlider(
    value=ny//2,
    min=0,
    max=ny-1,
    description='y index',
    continuous_update=False,
    layout=slider_layout,
    readout=False,
)
cy_slider.observe(update_y_slice,names='value')

cz_slider = IntSlider(
    value=nz//2,
    min=0,
    max=nz-1,
    description='z index',
    continuous_update=False,
    layout=slider_layout,
    readout=False,
)
cz_slider.observe(update_z_slice,names='value')

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

controls = VBox(
    [
        Label("Visual Controls"),
        cmap_widget,
        vol_checkbox,
        clip_checkbox,
        atoms_checkbox,
        atoms_slider,
        Label("Slice Controls"),
        cx_checkbox,
        cy_checkbox,
        cz_checkbox,
        Label("Indices Controls"),
        cx_slider,
        cy_slider,
        cz_slider
    ],
    layout=controls_layout
)

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

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