# Widget for displaying movies interactively.

%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 numpy as np
import imageio.v3 as iio

plot_every_frame = 2
file_name = "data/45grains_compressed.mp4"
video_shape = iio.improps(file_name,plugin="pyav").shape
num_frames = video_shape[0] // plot_every_frame

dpi = 72
array = iio.imread(file_name,index=0,plugin='pyav')[:,:,0]

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

hist_range_plot = (-4,4)
hist_range_init = (-2,1)
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=(420/dpi, 430/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 = '430px'
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))

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()

scalebar_plot=ax.plot((20,70),(484-20,484-20),lw=5,color=cmap(0))
scalebar_text=ax.text((70+20)/2,484-20-10,"2 nm",color=cmap(0),ha='center')

def display_frame(change):

    # main
    index = change['new']
    array = iio.imread(file_name,index=index*plot_every_frame,plugin='pyav')[:,:,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_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')

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='630px'
)

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
        )
)