import streamlit as st
import gymnasium
from gymnasium.wrappers.record_video import RecordVideo
from PIL import Image
import time

steps = st.number_input("Step Number", min_value=1, value=100)
sim_clicked = st.button("Simulate")

if sim_clicked:
    env = gymnasium.make("MountainCarContinuous-v0", render_mode="rgb_array")
    env = RecordVideo(env, video_folder="./videos")
    obs, info = env.reset()
    with st.empty():
        for i in range(steps):
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            rgb_array = env.render()

            image = Image.fromarray(rgb_array)
            st.image(image, caption=f"Step {i}")
            time.sleep(0.02)

            if terminated or truncated:
                break
        env.close()

    st.video("./videos/rl-video-episode-0.mp4")