import gym
from stable_baselines3 import DQN
from tetris_env import TetrisEnv
import pygame
import time

def main():
    env = TetrisEnv()

    model = DQN.load("models/dqn_tetris")

    episodes = 5

    for ep in range(1, episodes + 1):
        obs = env.reset()
        done = False
        total_reward = 0
        while not done:
            env.render(mode='human')

            action, _states = model.predict(obs, deterministic=True)

            obs, reward, done, info = env.step(action)

            total_reward += reward
            pygame.time.wait(100)  

        print(f"Episode {ep}: Total Reward = {total_reward}")

    env.close()

if __name__ == "__main__":
    main()