# Widget for displaying movies interactively.

%matplotlib widget

from IPython.display import display
import matplotlib.pyplot as plt
import types
from ipywidgets import Dropdown, IntSlider, Layout, HBox, VBox, Checkbox, Label
import numpy as np
import h5py

plt.ioff()

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

# Define dimensions
nx, ny, nz = data.shape
X, Y, Z = np.mgrid[0:nx,0:ny,0:nz]

dpi = 72

# Create a figure with 3D ax
fig = plt.figure(figsize=(400/dpi, 400/dpi), dpi=dpi)

fig.canvas.resizable = False
fig.canvas.header_visible = False
fig.canvas.footer_visible = False
fig.canvas.toolbar_visible = True
fig.canvas.layout.width = '400px'
fig.canvas.toolbar_position = 'bottom'

ax = fig.add_subplot(111, projection='3d')

kw = {
    'vmin': data.min(),
    'vmax': data.max(),
    'alpha':0.875,
    'cmap': 'magma',
    'levels': np.linspace(data.min(), data.max(), 21),
}

# Plot contour surfaces
Cz = ax.contourf(
    X[:, :, nz//2], Y[:, :, nz//2], data[:, :, nz//2],
    zdir='z', offset=nz//2, **kw
)

Cy = ax.contourf(
    X[:, ny//2, :], data[:, ny//2, :], Z[:, ny//2, :],
    zdir='y', offset=ny//2, **kw
)

Cx = ax.contourf(
    data[nx//2, :, :], Y[nx//2, :, :], Z[nx//2, :, :],
    zdir='x', offset=nx//2, **kw
)

Cs = [Cx,Cy,Cz]

def setvisible(self,vis):
    for c in self.collections: c.set_visible(vis)
        
Cx.set_visible = types.MethodType(setvisible,Cx)
Cy.set_visible = types.MethodType(setvisible,Cy)
Cz.set_visible = types.MethodType(setvisible,Cz)

xmin, xmax = X.min(), X.max()
ymin, ymax = Y.min(), Y.max()
zmin, zmax = Z.min(), Z.max()
ax.set(xlim=[xmin, xmax], ylim=[ymin, ymax], zlim=[zmin, zmax])

ax.axis('off')
fig.set_frameon(False)
ax.set_box_aspect((1,1,1))

# Colorbar
cb = fig.colorbar(Cz, ax=ax,fraction=0.025,pad=0.025,location='right',format='%.3f')
cbar = [cb] # this is hacky, switch to classes?
# cb.yaxis.set_tick_params(labelcolor=(0.4,0.4,0.4))  % not sure what the right handle is here


fig.tight_layout()

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 = change['new']
    Cs[0].set_cmap(cmap)
    Cs[1].set_cmap(cmap)
    Cs[2].set_cmap(cmap)
    
    cbar[0].remove()
    cbar[0] = fig.colorbar(Cs[2], ax=ax,fraction=0.025,pad=0.025,location='right',format='%.3f')
    
    fig.canvas.draw_idle()
    return None

def update_x_slice(change):
    if cx_checkbox.value:
        offset = change['new']
        Cs[0].remove()
        Cs[0] = ax.contourf(
            data[offset, :, :], Y[offset, :, :], Z[offset, :, :],
            zdir='x', offset=offset, **(kw|{'cmap':cmap_widget.value})
        )
        Cs[0].set_visible = types.MethodType(setvisible,Cs[0])
        fig.canvas.draw_idle()
    return None

def update_y_slice(change):
    if cy_checkbox.value:
        offset = change['new']
        Cs[1].remove()
        Cs[1] = ax.contourf(
            X[:, offset, :], data[:, offset, :], Z[:, offset, :],
            zdir='y', offset=offset, **(kw|{'cmap':cmap_widget.value})
        )
        Cs[1].set_visible = types.MethodType(setvisible,Cs[1])
        fig.canvas.draw_idle()
    return None

def update_z_slice(change):
    if cz_checkbox.value:
        offset = change['new']
        Cs[2].remove()
        Cs[2] = ax.contourf(
            X[:, :, offset], Y[:, :, offset], data[:, :, offset],
            zdir='z', offset=offset, **(kw|{'cmap':cmap_widget.value})
        )
        Cs[2].set_visible = types.MethodType(setvisible,Cs[2])
        fig.canvas.draw_idle()
    return None
    
def toggle_x_slice(change):
    visible = change['new']
    Cs[0].set_visible(visible)
    fig.canvas.draw_idle()
    return None

def toggle_y_slice(change):
    visible = change['new']
    Cs[1].set_visible(visible)
    fig.canvas.draw_idle()
    return None

def toggle_z_slice(change):
    visible = change['new']
    Cs[2].set_visible(visible)
    fig.canvas.draw_idle()
    return None

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

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

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

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

slider_layout = Layout(width='175px')
cx_slider = IntSlider(
    value=nx//2,
    min=0,
    max=nx,
    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,
    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,
    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,
        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='600px'
)

display(
    HBox([
        fig.canvas,
        controls
    ],
        layout=visualization_layout
        )
)