%matplotlib widget

import numpy as np
import matplotlib.pyplot as plt

from IPython.display import display
from ipywidgets import HBox, VBox, widgets, interact, Dropdown, Label, Layout
import h5py

hf = h5py.File('data/nanoparticles.h5', 'r')
DPC = hf.get('DPC')[:]
array_3D = hf.get('data3D')[:]
hf.close()
qx = np.fft.fftfreq(array_3D.shape[2] * 2, 1/array_3D.shape[2]/2)
qy = np.fft.fftfreq(array_3D.shape[2] * 2, 1/array_3D.shape[2]/2)
qya, qxa = np.meshgrid(qx, qy)
qra = np.fft.fftshift(np.sqrt(qxa**2 + qya **2))
q_pixel_size = 15/6.933

with plt.ioff():
    dpi = 72
    fig = plt.figure(figsize=(675/dpi, 225/dpi), dpi=dpi)

ax0 = fig.add_axes([0.04,  0.05,  0.28, 0.75])
ax1 = fig.add_axes([0.37,  0.05,  0.28, 0.75])
ax2 = fig.add_axes([0.70,  0.05,  0.28, 0.75])
ax3 = fig.add_axes([0.048,  0.05,  0.082, 0.2])
ax4 = fig.add_axes([0.378,  0.05,  0.082, 0.2])

cmap = 'gray'

im0 = ax0.imshow(
    array_3D[:,:,0:int(15/q_pixel_size)].sum(2), 
    cmap = cmap
)

im1 = ax1.imshow(
    array_3D[:,:,int(22/q_pixel_size):int(100/q_pixel_size)].sum(2), 
    cmap = cmap
)

im2 = ax2.imshow(
    DPC,
    cmap = cmap)

cmap2 = "YlGn_r"
im3 = ax3.imshow(
    qra < 15/q_pixel_size,
    cmap = cmap2,
    vmax = 1.5, 
    vmin = -0.5
)


circle = plt.Circle(
    (qra.shape[0]/2, qra.shape[1]/2), 15/q_pixel_size,
    color='black', fill=False, lw =1
)
ax3.add_patch(circle)

im4 = ax4.imshow(
    np.logical_and(qra > 22/q_pixel_size, qra < 100/q_pixel_size),
    cmap = cmap2,
    vmax = 1.5, 
    vmin = -0.5
)

circle = plt.Circle(
    (qra.shape[0]/2, qra.shape[1]/2), 15/q_pixel_size,
    color='black', fill=False, lw =1
)
ax4.add_patch(circle)
ax0.set_xticks([])  
ax0.set_yticks([]) 
ax1.set_xticks([])  
ax1.set_yticks([]) 
ax2.set_xticks([])  
ax2.set_yticks([]) 
ax3.set_xticks([])  
ax3.set_yticks([]) 
ax4.set_xticks([])  
ax4.set_yticks([]) 

ax0.set_title('Circular detector virtual image');
ax1.set_title('Annular detector virtual image');
ax2.set_title('DPC reconstruction');

def update_ims(BF_radius, DF_radius_1, DF_radius_2):
    BF_radius = int(BF_radius/q_pixel_size)
    DF_radius_1 = int(DF_radius_1/q_pixel_size)
    DF_radius_2 = int(DF_radius_2/q_pixel_size)

    BF = array_3D[:,:,0:BF_radius].sum(2)
    
    im0.set_data(BF)
    im0.set_clim(vmax = BF.max(), vmin = BF.min())
    DF = array_3D[:,:,DF_radius_1:DF_radius_2].sum(2)
    im1.set_data(DF)
    im1.set_clim(vmax = DF.max(), vmin = DF.min())
    im3.set_data(qra < BF_radius)
    im4.set_data(np.logical_and(qra > DF_radius_1, qra < DF_radius_2 ))
    
    fig.canvas.draw_idle()

style = {
    'description_width': 'initial',
}

layout = Layout(width='250px',height='30px')

BF_radius = widgets.IntSlider(
    value = 16, min = 0, max =170, 
    step = 0.5,
    description = "BF radius (mrad)",
    style = style,
    layout = layout,
)


DF_radius_1 = widgets.IntSlider(
    value = 50, min = 0, max = 170, 
    step = 1,
    description = "Inner radius (mrad)",
    style = style,
    layout = layout,
)


DF_radius_2 = widgets.IntSlider(
    value = 120, min = 0, max = 170, 
    step = 1,
    description = "Outer radius (mrad)",
    style = style,
    layout = layout,
)

widgets.interactive_output(
    update_ims, 
    {
        'BF_radius':BF_radius,
        'DF_radius_1':DF_radius_1,
        'DF_radius_2':DF_radius_2,
    },
)



option_list = (
    'select an option',
    'BF',
    'ABF',
    'ADF',
    'HAADF',
)

# update the plots with a pre-selected function
def select_preset_eventhandler(change):
    if change.new == option_list[1]: #BF
        BF_radius.value = 16
    
    if change.new == option_list[2]: #ABF
        DF_radius_1.value = 9
        DF_radius_2.value = 16
        
    if change.new == option_list[3]: #ADF
        DF_radius_1.value = 50
        DF_radius_2.value = 120
    
    if change.new == option_list[4]: #HAADF
        DF_radius_1.value = 90
        DF_radius_2.value = 170
            
# Widgets
dropdown = Dropdown(
    options = option_list,
    layout = Layout(width='125px',height='30px'),
)
dropdown.observe(select_preset_eventhandler, names='value')

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

widget = widgets.VBox(
    [
        fig.canvas,
        HBox([
            BF_radius, 
            VBox([DF_radius_1,DF_radius_2]),
            dropdown
        ]),
    ],
)

display(widget);