from simple_grid import MultiRoomGrid
from PIL import Image


def main():
    env = MultiRoomGrid(
        [["wwow", "owwo"], ["wwow", "ooww"]], start_rooms=[[0, 0]], goal_rooms=[[1, 0]], room_size = 3,
        render_mode="rgb_array", highlight=False
    )
    
    observation, info = env.reset()
    
    frames = []
    for _ in range(1000):
        action = env.action_space.sample()
        observation, reward, terminated, truncated, info = env.step(action)
        
        frame = Image.fromarray(env.get_frame(), 'RGB')
        frames.append(frame)

        if terminated or truncated:
            observation, info = env.reset()
    
    env.close()
    
    frames[0].save('./random_agent.gif', save_all=True, append_images=frames[1:], duration=100, loop=0)
    
    
        
if __name__ == "__main__":
    main()