from brain import *
class Agent():
    def __init__(self,num_states,num_actions,reload=None,load_mem=True,batch_size=256):
        self.brain=Brain(num_states=num_states*Seq_max,num_actions=num_actions,reload=reload,load_mem=load_mem,batchsize=batch_size)
    
    def update_q(self):
        self.brain.replay()

    def get_action(self,state,episode):
        action=self.brain.decide(state,episode)
        return action
    
    def memorize(self,state,action,state_next,reward):
        if reward>=0.5:
            self.brain.mem.push_s(state,action,state_next,reward)
        else:
            self.brain.mem.push(state,action,state_next,reward)
            

    def update(self):
        self.brain.update()
    
    def load(self,path=None):
        self.brain.load(path)
    
    def save(self,path=None):
        self.brain.save(path)