# 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 as iio
from PIL import Image
import matplotlib

#For big animations
matplotlib.rcParams['animation.embed_limit'] = 2**128

dpi = 70
plot_every_frame = 6
width = 700
image_name = "data/DoubleTilt_static.jpg"
movie_name = "data/DoubleTilt.mp4"
second_movie_name = "data/DoubleTiltWithSample.mp4"

plt.ioff()
fig, axs = plt.subplots(3, figsize=(width/dpi, width/dpi), dpi=dpi)

#Static plot
axs[0].axis('off')
im = Image.open(image_name)
im=axs[0].imshow(im)

#Movie
axs[1].axis('off')
vid = iio.get_reader(movie_name)
num_frames = vid.count_frames()
array = vid.get_data(0)
array = array[:,:,:] # grayscale
im=axs[1].imshow(array)

def display_frame(index):
    array = vid.get_data(index*plot_every_frame)[:,:,:]
    #array = 255 - array # reverse
    im.set_data(array)
    fig.canvas.draw_idle()

    return None
    
animation = FuncAnimation(
    fig,
    display_frame,
    frames=num_frames//plot_every_frame,
    repeat=False,
    blit=False,
    interval=50,
)


#Second Movie
axs[2].axis('off')
second_vid = iio.get_reader(second_movie_name)
second_num_frames = second_vid.count_frames()
second_array = second_vid.get_data(0)
second_array = second_array[:,:,:] # grayscale
im=axs[2].imshow(second_array)

def display_frame_second(index):
    second_array = second_vid.get_data(index*plot_every_frame)[:,:,:]
    #array = 255 - array # reverse
    im.set_data(second_array)
    fig.canvas.draw_idle()

    return None
    
second_animation = FuncAnimation(
    fig,
    display_frame_second,
    frames=second_num_frames//plot_every_frame,
    repeat=False,
    blit=False,
    interval=50,
)


HTML(animation.to_jshtml())

HTML(second_animation.to_jshtml())