import numpy as np
import skvideo.io

from apple.envs.discrete_apple import get_apple_env

env = get_apple_env("full", time_limit=10, start_x=0, c=0.5, goal_x=8, render_mode="rgb_array")


imgs = []
env.reset()
for i in range(8):
    imgs.append(env.unwrapped.render())
    env.step(1)

for i in range(9):
    imgs.append(env.unwrapped.render())
    env.step(0)


skvideo.io.vwrite(
    "example_rollout.mp4",
    np.stack(imgs),
    inputdict={
        "-r": str(int(4)),
    },
    outputdict={
        "-f": "mp4",
        "-pix_fmt": "yuv420p",  # '-pix_fmt=yuv420p' needed for osx https://github.com/scikit-video/scikit-video/issues/74
    },
)