import gym 
import random





class Agent():
    def __init__(self, env):
        self.action_size = env.action_space.n
        print("Action Size: ", self.action_size)

    
    def get_action(self, state):
        # action = random.choice(range(self.action_size))
        pole_angles = state[2]
        action = 0 if pole_angles < 0 else 1
        return action


if __name__ == "__main__":

    game_name = "CartPole-v1"

    env = gym.make(game_name)

    agent = Agent(env)
    state = env.reset()

    for _ in range(200):
        # action = env.action_space.sample()
        action = agent.get_action(state)
        state, reward, done, info = env.step(action)
        env.render()    