# 导入必要的库
%matplotlib widget
import os
import matplotlib.pyplot as plt
import numpy as np
from ipywidgets import Play, Layout, VBox, HBox, IntSlider
from matplotlib.widgets import Slider
from IPython.display import display
import imageio as iio

# 加载视频数据
vid_mask = iio.get_reader(os.path.join("video", "microscopy_video_segmentation.mp4"))
vid_orig = iio.get_reader(os.path.join("video", "microscopy_video.mp4"))
num_frames = vid_mask.count_frames()

# 设置基本参数
dpi = 15  # DPI for rendering
plot_every_frame = 1
width = 300

# 关闭交互式绘图，初始化图形
plt.ioff()
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(2 * width / dpi, width / 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 = '620px'
fig.canvas.layout.height = "350px"
fig.canvas.toolbar_position = 'bottom'

# 添加标题
ax1.set_title("With Mask", fontsize=60, pad=10)
ax2.set_title("Microscopy Video", fontsize=60, pad=10)

# 隐藏轴
ax1.axis('off')
ax2.axis('off')

# 加载第一帧的数据并显示
array_mask = vid_mask.get_data(0)
array_orig = vid_orig.get_data(0)

# 显示视频的第一帧
im1 = ax1.imshow(array_mask)
im2 = ax2.imshow(array_orig)

# 调整布局
fig.tight_layout()
fig.subplots_adjust(top=0.93)

# 定义函数来显示当前帧
def display_frame(index):
    # 更新每一帧的数据
    array_mask = vid_mask.get_data(index * plot_every_frame)
    array_orig = vid_orig.get_data(index * plot_every_frame)

    im1.set_data(array_mask)
    im2.set_data(array_orig)

    fig.canvas.draw_idle()

# 创建播放控件
index_widget = Play(value=0, min=0, max=num_frames // plot_every_frame - 1, step=1, interval=50, layout=Layout(width='110px'))

# 创建滑动条控件
slider = IntSlider(value=0, min=0, max=num_frames // plot_every_frame - 1, step=1, description='', layout=Layout(width='500px'))

# 当滑块或播放控件的值变化时更新帧
def on_value_change(change):
    index_widget.value = change['new']  # 同步Play控件的值
    display_frame(change['new'])

# 同步播放控件与滑块
def on_play_change(change):
    slider.value = change['new']  # 同步Slider的值
    display_frame(change['new'])

# 绑定事件
slider.observe(on_value_change, names='value')
index_widget.observe(on_play_change, names='value')

# 实现同步缩放的功能
def on_zoom(event):
    if event.inaxes == ax1:  # 如果在第一个子图中操作
        ax2.set_xlim(ax1.get_xlim())  # 同步X轴缩放
        ax2.set_ylim(ax1.get_ylim())  # 同步Y轴缩放
    elif event.inaxes == ax2:  # 如果在第二个子图中操作
        ax1.set_xlim(ax2.get_xlim())  # 同步X轴缩放
        ax1.set_ylim(ax2.get_ylim())  # 同步Y轴缩放
    fig.canvas.draw_idle()

# 连接缩放同步事件
fig.canvas.mpl_connect('motion_notify_event', on_zoom)

# 控制播放的布局
controls_box = HBox([index_widget, slider])

# 将视频和控件一起显示
display(VBox([controls_box, fig.canvas]))