import numpy as np
from ple import PLE
from ple.games.snake import Snake


class NaiveAgent():
    """
            This is our naive agent. It picks actions at random!
    """

    def __init__(self, p):
        #self.actions = actions
        #self.states = states
        self.instanta = p

    def pickAction(self, reward):
        print(self.instanta.getActionSet())
        print(self.instanta.getGameState())
        self.actions = self.instanta.getActionSet()
        return self.actions[np.random.randint(0, len(self.actions))]
        #return self.actions[np.argmax(np.array(list(filter(lambda x: type(x) is int,self.actions))))]



game = Snake(width=200,height=200)


reward = 0.0
#max_noops = 20
#nb_frames = 15000


p = PLE(game,fps=30,force_fps=False, display_screen=True,frame_skip=10)

agent = NaiveAgent(p)

p.init()

# lets do a random number of NOOP's
#for i in range(np.random.randint(0, max_noops)):
#    reward = p.act(p.NOOP)

# start our training loop
for f in range(1000):
    # if the game is over
    if p.game_over():
        p.reset_game()

    #print(p.getScreenDims())
    #obs = p.getScreenRGB()
    action = agent.pickAction(reward)
    #print("am ajuns aici")
    reward = p.act(action)
    #print(reward)