# Widget for displaying images interactively.

%matplotlib widget

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.path import Path
from matplotlib.patches import PathPatch
from IPython.display import display
from ipywidgets import Checkbox, Dropdown, FloatRangeSlider, Layout, HBox, VBox

# Import stack of images

image_stack = np.load('data/im_graphene_EWR_small.npz')['im_graphene_EWR']

plt.close('all')
plt.ioff()

dpi = 72
hist_range_plot = (-4,4)
hist_range_init = (-1,2)
hist_num_bins = 100

# Histogram visualization
def compute_histogram_values(array):
    int_mean = np.mean(array)
    int_std = np.sqrt(np.mean((array - int_mean)**2))
    
    int_min = int_mean + hist_range_plot[0] * int_std
    int_max = int_mean + hist_range_plot[1] * int_std
    
    init_min = int_mean + hist_range_init[0] * int_std
    init_max = int_mean + hist_range_init[1] * int_std
    
    int_ranges = (int_mean, int_std, int_min, int_max,init_min,init_max)
    
    hist_bins = np.linspace(
        int_ranges[2],
        int_ranges[3],
        hist_num_bins+1,
        endpoint=True)
    hist_data, _ = np.histogram(
        array.ravel(),
        bins=hist_bins,
    )
    hist_data = hist_data.astype('float')
    hist_data /= np.max(hist_data)
    
    hist_bins = hist_bins[:-1] + (hist_bins[1] - hist_bins[0])/2
    
    return hist_bins, hist_data, int_ranges

hist_bins_all, hist_data_all, int_ranges_all, hist_slider_state = [], [], [], []
for array in image_stack:
    hist_bins, hist_data, int_ranges = compute_histogram_values(array)
    hist_bins_all.append(hist_bins)
    hist_data_all.append(hist_data)
    int_ranges_all.append(int_ranges)
    hist_slider_state.append(
        {
            'min':int_ranges[2],
            'max':int_ranges[3],
            'step':(int_ranges[3]-int_ranges[2])/100,
            'value':int_ranges[4:6]
        }
    )
    
array = image_stack[0]
hist_bins, hist_data, int_ranges = hist_bins_all[0], hist_data_all[0], int_ranges_all[0]
fig_hist, ax_hist = plt.subplots(figsize=(190/dpi, 160/dpi), dpi=dpi)

hist_path = Path(np.array([hist_bins,hist_data]).T)
hist_patch = PathPatch(hist_path,visible=False,transform=ax_hist.transData)

cmap = plt.colormaps.get_cmap('gray')

im_hist = ax_hist.imshow(
    hist_bins[None],
    vmin=int_ranges[4],
    vmax=int_ranges[5],
    cmap=cmap,
    origin='lower',
    aspect="auto",
    interpolation='bilinear',
    clip_path=hist_patch,
    extent=[int_ranges[2],int_ranges[3],0,np.quantile(hist_data,0.9)],
    clip_on=True
)

h_vlines = ax_hist.vlines(
    int_ranges[4:6],
    ymin = 0,
    ymax = 1.1,
    colors = [cmap(cmap.N),cmap(0)],
)

ax_hist.set_xlim((int_ranges[2], int_ranges[3]));
ax_hist.set_ylim((0, np.quantile(hist_data,0.9)));
ax_hist.set(yticks=[])
ax_hist.set(yticklabels=[])
ax_hist.set_facecolor((0.9,0.85,0.85))

fig_hist.canvas.toolbar_visible = False
fig_hist.canvas.header_visible = False
fig_hist.canvas.footer_visible = False
fig_hist.canvas.resizable = False
fig_hist.tight_layout()

# Main visualization
fig, ax = plt.subplots(figsize=(480/dpi, 460/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 = '480px'
fig.canvas.toolbar_position = 'bottom'

ax.axis('off')
fig.set_frameon(False)
fig_hist.set_frameon(False)
ax_hist.xaxis.set_tick_params(labelcolor=(0.4,0.4,0.4))
# fig.patch.set_visible(False)
# fig_hist.patch.set_visible(False)

im=ax.imshow(array,cmap=cmap,vmin=int_ranges[4],vmax=int_ranges[5])
divider = make_axes_locatable(ax)
ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
ax_cb.yaxis.set_tick_params(labelcolor=(0.4,0.4,0.4))
fig.add_axes(ax_cb)
fig.colorbar(im, cax=ax_cb)
fig.tight_layout()

# Define pixel size to make scalebar
#pixel_size = 0.0151
#pixel_units = 'nm'

# 130 pixels ~ 2nm
# image is 1400x1400
scalebar_plot=ax.plot((70,70+130),(1400-70,1400-70),lw=5,color=cmap(cmap.N))
scalebar_text=ax.text((130+70*2)/2,1400-70-20,"2 nm",color=cmap(cmap.N),ha='center')

option_list_image = [
    'exit wave phase', 
    'exit wave amplitude', 
]

def display_frame(change):
    
    # main
    index_name = change['new']
    index = option_list_image.index(index_name)
    array = image_stack[index]
    
    im.set_data(array)

    # hist
    hist_bins, hist_data, int_ranges = hist_bins_all[index], hist_data_all[index], int_ranges_all[index]

    path = Path(np.array([hist_bins,hist_data]).T)
    patch = PathPatch(path,visible=False,transform=ax_hist.transData)
    
    im_hist.set_data(hist_bins[None])
    im_hist.set_clip_path(patch)
    im_hist.set_extent([int_ranges[2],int_ranges[3],0,np.quantile(hist_data,0.9)])
    ax_hist.set_xlim((int_ranges[2], int_ranges[3]));
    ax_hist.set_ylim((0, np.quantile(hist_data,0.9)));

    histogram_state = histogram_range_slider.get_state()
    histogram_state = histogram_state | hist_slider_state[index]

    histogram_range_slider.set_state(histogram_state)
    update_vlines({'new':hist_slider_state[index]['value']})

    fig.canvas.draw_idle()
    fig_hist.canvas.draw_idle()
    
    return None

def toggle_scalebar(change):
    
    scalebar = change['new']
    scalebar_plot[0].set_visible(scalebar)
    scalebar_text.set_visible(scalebar)
    fig.canvas.draw_idle()
    
    return None

def update_colormap(change):
    
    cmap_string = change['new']
    cmap = plt.colormaps.get_cmap(cmap_string)
    
    im.set_cmap(cmap)
    im_hist.set_cmap(cmap)
    h_vlines.set_colors([cmap(cmap.N),cmap(0)])

    scalebar_plot[0].set_color(cmap(cmap.N))
    scalebar_text.set_color(cmap(cmap.N))
    
    fig.canvas.draw_idle()
    fig_hist.canvas.draw_idle()
    return None

def update_vlines(change):
    
    index_name = index_widget.value
    index = option_list_image.index(index_name)

    min, max = change['new']
    p = np.array([
        [
            [min, 0],
            [min, 1.1],
        ],
        [
            [max, 0],
            [max, 1.1],
        ]
    ])
    
    h_vlines.set_segments(p)
    fig_hist.canvas.draw_idle()
    
    im.set_clim([min,max])
    im_hist.set_clim([min,max])
    
    hist_slider_state[index]['value']=[min,max]
    
    fig.canvas.draw_idle()
    fig_hist.canvas.draw_idle()
    
    return None

index_widget = Dropdown(options=option_list_image,index=0,layout=Layout(width='180px'))
index_widget.observe(display_frame,names='value')

scalebar_widget =Checkbox(value=True,description="Show scale bar",indent=False,layout=Layout(width='180px'))
scalebar_widget.observe(toggle_scalebar,names='value')

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'
]

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

histogram_range_slider = FloatRangeSlider(
    value=int_ranges[4:6],
    min=int_ranges[2],
    max=int_ranges[3],
    step=(int_ranges[3]-int_ranges[2])/100,
    continuous_update=False,
    orientation='horizontal',
    readout=False,
    indent=True,
    layout=Layout(width='190px')
)
histogram_range_slider.observe(update_vlines,names='value')

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

histogram_box_layout = Layout(
    display='flex',
    flex_flow='column',
    align_items='center',
    width='200px'
)
histogram_vbox = VBox([fig_hist.canvas,histogram_range_slider],layout=histogram_box_layout)

controls_vbox = VBox([histogram_vbox,index_widget,cmap_widget,scalebar_widget],layout=histogram_box_layout)

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