import gymnasium as gym
from stable_baselines3 import DQN
import time

# Create the CartPole environment with rendering
env = gym.make("CartPole-v1", render_mode="human")

# Initialize a DQN model
model = DQN("MlpPolicy", env, verbose=1)

# Train the model for a few episodes (e.g., 5000 timesteps)
print("Training the agent...")
model.learn(total_timesteps=5000)

# Test the trained agent and visualize
print("Testing the trained agent...")
obs, _ = env.reset()
total_reward = 0
for _ in range(1000):  # Run for 1000 timesteps
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, done, truncated, info = env.step(action)
    total_reward += reward
    env.render()
    time.sleep(0.02)  # Slow down the rendering for visibility
    if done or truncated:
        print(f"Episode finished with total reward: {total_reward}")
        total_reward = 0
        obs, _ = env.reset()

env.close()