"""
ImageWidget export images or video
==================================

This example is meant to be run in an interactive shell. To run in a script use the offscreen canvas.

imageio is required to export to images.

This example uses opencv to export to a video but you don't have to use opencv, you
can use any other video writer that can write numpy arrays to video frames.
"""

# test_example = true
# sphinx_gallery_pygfx_docs = 'code'

# %%
import fastplotlib as fpl
import imageio.v3 as iio
import cv2


movie = iio.imread("imageio:cockatoo.mp4")

iw = fpl.ImageWidget(
    movie,
    rgb=[True],
    figure_kwargs={"size": (700, 560), "canvas": "offscreen"}  # for offscreen rendering provide "canvas": "offscreen"
)

# show image widget
iw.show()

# %%

# wait for image widget to display

# export a snapshot of the figure as it is currently displayed
iw.figure.export("./image.png")

# obtain a numpy array
frame = iw.figure.export_numpy(rgb=True)

# %%
# create a video writer
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
writer = cv2.VideoWriter(
    "./iw_video.mp4",
    fourcc,
    20.0,
    (frame.shape[1], frame.shape[0]),
)

# write a video file of the first 100 frames
for i in range(100):
    frame = iw.figure.export_numpy(rgb=True)
    writer.write(frame)
    iw.current_index = {"t": i}

writer.release()

"""
Figure export images or video
=============================

This example is meant to be run in an interactive shell. To run in a script use the offscreen canvas.

imageio is required to export to images.

This example uses opencv to export to a video but you don't have to use opencv, you
can use any other video writer that can write numpy arrays to video frames.
"""

# test_example = false
# sphinx_gallery_pygfx_docs = 'code'

# %%
import fastplotlib as fpl
import numpy as np
import cv2

# create a random distribution of 10,000 xyz coordinates
n_points = 10_000

# dimensions always have to be [n_points, xyz]
dims = (n_points, 3)

clouds_offset = 15

# create some random clouds
normal = np.random.normal(size=dims, scale=5)
# stack the data into a single array
cloud = np.vstack(
    [
        normal - clouds_offset,
        normal,
        normal + clouds_offset,
    ]
)

# color each of them separately
colors = ["yellow"] * n_points + ["cyan"] * n_points + ["magenta"] * n_points

# create plot
figure = fpl.Figure(
    size=(700, 560),
    canvas="offscreen"  # for offscreen rendering, ex. in a script
)
subplot_scatter = figure[0, 0]
# use an alpha value since this will be a lot of points
scatter_graphic = subplot_scatter.add_scatter(data=cloud, sizes=3, colors=colors, alpha=0.6)


def update():
    # move every point by a small amount
    deltas = np.random.normal(size=scatter_graphic.data.value.shape, loc=0, scale=0.05)
    scatter_graphic.data = scatter_graphic.data.value + deltas


figure.show()

# %%

# wait for figure to display

# export a snapshot of the figure as it is currently displayed
figure.export("./image.png")

# obtain a numpy array
frame = figure.export_numpy(rgb=True)

# %%
# create a video writer
fourcc = cv2.VideoWriter_fourcc(*'XVID')
writer = cv2.VideoWriter(
    "./scatter_video.avi",
    fourcc,
    50.0,
    (frame.shape[1], frame.shape[0]),
)

# write a video file with 500 frames
for i in range(500):
    frame = figure.export_numpy(rgb=True)
    writer.write(frame)
    # update the graphic
    update()

writer.release()