from pydoc import render_doc
import gymnasium as gym
import gym_envs
import numpy as np
import random

env = gym.make("adv_driving/metadrive_highway_env", render_mode="topdown")

print(env.unwrapped.metaenv.agents)
print(env.unwrapped.metaenv.num_agents)
env.reset()
print(env.unwrapped.metaenv.agents)
print(env.unwrapped.metaenv.num_agents)
frames = []
for i in range(500):
    action = [0, 0.0, random.uniform(0.5, 1)]
    action = np.array(action)
    #print(action)
    obs, ttc, terminated, _, info = env.step(action)
    #print("TTC: ", ttc)
    #print("obs: ", obs)
    
    if terminated:
        break
    frames.append(info["frame"])


env.close()
print("\nGenerate gif...")
import numpy as np
from PIL import Image

imgs = [frame for frame in frames]
imgs = [Image.fromarray(img) for img in imgs]
imgs[0].save("demo.gif", save_all=True, append_images=imgs[1:], duration=50, loop=0)
