# Widget for displaying movie animations.

%matplotlib widget

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.axes_grid1 import make_axes_locatable
from IPython.display import HTML
import numpy as np
import imageio.v3 as iio

plot_every_frame = 6
file_name = "data/45grains_compressed.mp4"
video_shape = iio.improps(file_name,plugin="pyav").shape
num_frames = video_shape[0] // plot_every_frame
frames = np.empty((num_frames,)+video_shape[1:-1],dtype=np.uint8)

with iio.imopen(file_name,"r",plugin="pyav") as vid:
    for index in range(num_frames):
        frames[index] = 255 - vid.read(index=index*plot_every_frame)[:,:,0]

dpi = 70
width = 400

plt.ioff()
fig, ax = plt.subplots(figsize=(width/dpi, width/dpi), dpi=dpi)
ax.axis('off')

array = frames[0]
nx, ny = array.shape

int_mean = np.mean(array)
int_std = np.sqrt(np.mean((array - int_mean)**2))

init_min = int_mean -1 * int_std
init_max = int_mean + 2 * int_std

im=ax.imshow(array,cmap='Greys',vmin=init_min,vmax=init_max)
divider = make_axes_locatable(ax)
ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
fig.add_axes(ax_cb)
fig.colorbar(im, cax=ax_cb)
fig.tight_layout()

ax.plot((20,70),(ny-20,ny-20),lw=5,color='k')
ax.text((70+20)/2,ny-20-10,"2 nm",color='k',ha='center')

def display_frame(index):
    im.set_data(frames[index])
    return [im]
    
animation = FuncAnimation(
    fig,
    display_frame,
    frames=num_frames,
    repeat=False,
    blit=True,
    interval=50,
)

HTML(animation.to_jshtml())