# imports 
%matplotlib widget

from ipywidgets import Play, Checkbox, Dropdown, FloatRangeSlider, Layout, HBox, VBox
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
import imageio as iio
import numpy as np
# from scipy.ndimage import gaussian_filter

# inputs and pixel size
file_data = "data/AgAu_no_overlay.mp4"
pixel_size = 0.83
pixel_units = 'nm'

scale_bar_length = 100
scale_bar_length_pixels = scale_bar_length / pixel_size

# Get data, compute high pass filtered version of movie
# Import video frames
vid = iio.get_reader(file_data)
num_frames = vid.count_frames()

array = vid.get_data(0)
# array = array[:,:,0] # grayscale - take first channel of movie

# high pass filtered array
# array_hp = array.astype('float') - gaussian_filter(
#     array.astype('float'),
#     sigma = (11,11,0.1),
#     mode = 'nearest')
# array_hp -= np.mean(array_hp)
# array_hp /= np.sqrt(np.mean(array_hp**2))
# array_hp = 255*np.clip((array_hp - 4) / (4 + 4), 0, 1)
# array = array_hp.astype('int')

dpi = 72
plt.close('all')
plt.ioff()

hist_range_plot = (-4,4)
hist_range_init = (-2.5,2.5)
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, hist_data, int_ranges = compute_histogram_values(array)
fig_hist, ax_hist = plt.subplots(figsize=(180/dpi, 150/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=(470/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 = '470px'
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))

# image plotting
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()

# overlay
color_overlay = (1,1,1)
fontsize_overlay = 16
lw_overlay = 4
hw_overlay = 12
hl_overlay = 12
overlay_plot = (
    ax.arrow(
    140,70,
    -30,30,
    lw = lw_overlay,
    color = color_overlay,
    head_width = hw_overlay,
    head_length = hl_overlay,
    ),
    ax.arrow(
    300,210,
    -30,30,
    lw = lw_overlay,
    color = color_overlay,
    head_width = hw_overlay,
    head_length = hl_overlay,
    ),
    ax.arrow(
    530,400,
    -30,30,
    lw = lw_overlay,
    color = color_overlay,
    head_width = hw_overlay,
    head_length = hl_overlay,
    ),
    ax.arrow(
    760,590,
    -30,30,
    lw = lw_overlay,
    color = color_overlay,
    head_width = hw_overlay,
    head_length = hl_overlay,
    ),
)
                
overlay_text = (
    ax.text(
    200,50,
    'Leading crack tip',
    color = color_overlay,
    fontsize = fontsize_overlay,
    ha='center',
    ),
    ax.text(
    360,190,
    'Coalesced voids',
    color = color_overlay,
    fontsize = fontsize_overlay,
    ha='center',
    ),
    ax.text(
    600,380,
    'Large voids',
    color = color_overlay,
    fontsize = fontsize_overlay,
    ha='center',
    ),
    ax.text(
    830,570,
    'Small voids',
    color = color_overlay,
    fontsize = fontsize_overlay,
    ha='center',
    ),
)
    
"""
# scale bar
color_scale_bar = (1,1,1)
scalebar_plot = ax.plot(
    (20,20 + scale_bar_length_pixels),
    (array.shape[0]-20,array.shape[0]-20),
    lw=5,color=color_scale_bar)
scalebar_text = ax.text(
    20 + scale_bar_length/2,
    array.shape[0]-20-10,
    str(scale_bar_length) + ' ' + pixel_units,
    color=color_scale_bar,
    ha='center')
"""

def display_frame(change):

    # main
    index = change['new']
    array = vid.get_data(index)[:,:,0] # grayscale image
    # array = 255 - array # reverse array
    im.set_data(array)
    fig.canvas.draw_idle()

    # hist
    hist_bins, hist_data, int_ranges = compute_histogram_values(array)

    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)

    fig_hist.canvas.draw_idle()
    
    return None



def toggle_overlay(change):
    overlay = change['new']
    for a0 in range(len(overlay_plot)):
        overlay_plot[a0].set_visible(overlay)
        overlay_text[a0].set_visible(overlay)
    fig.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(0))
    scalebar_text.set_color(cmap(0))
    
    fig.canvas.draw_idle()
    fig_hist.canvas.draw_idle()
    return None

def update_vlines(change):
    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])
    
    fig.canvas.draw_idle()
    fig_hist.canvas.draw_idle()
    
    return None

index_widget = Play(value=0,min=0,max=num_frames-1,step=1,interval=50,layout=Layout(width='180px'))
index_widget.observe(display_frame,names='value')


overlay_widget = Checkbox(
    value=True,
    description="Show overlays",
    indent=False,
    layout=Layout(width='180px'))
overlay_widget.observe(toggle_overlay,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],
    continuous_update=False,
    orientation='horizontal',
    readout=False,
    indent=True,
    layout=Layout(width='180px')
)
histogram_range_slider.observe(update_vlines,names='value')

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

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,
    overlay_widget,
],layout=histogram_box_layout)

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