import numpy as np

class environment:
    def __init__(self):
        self.envData = -1*np.ones( (5,4) )
        self.envData[1,3] = self.envData[2,3] = self.envData[3,3] = -1000
        self.envData[4,3] = 1000
        self.dogX = 0; self.dogY = 3
        
    def setDog(self, x=0, y=3):
        self.dogX = x; self.dogY = y 
        
    def move(self,a):  
        if a == 0 : self.dogY -= 1 # a = 0 east
        if a == 1 : self.dogX += 1 # a = 1 south, 
        if a == 2 : self.dogY += 1 # a = 2 west
        if a == 3 : self.dogX -= 1 # a = 3 north
        if self.dogY >3 or self.dogY <0 or self.dogX >4 or self.dogX <0:
            r = -10
        else: 
            r = self.envData[self.dogX,self.dogY]
        if self.dogY <0: self.dogY = 0
        if self.dogX <0: self.dogX = 0
        if self.dogX >4: self.dogX = 4
        if self.dogY >3: self.dogY = 3
        return(r)
        
    def sensors(self):
        return(self.dogX,self.dogY)
    
class learningDogSARSA:
    def __init__(self, env,vareps):
        self.hatQ = np.zeros( (5,4,4) ) # initialisiere hatQ(s,a)
        self.env = env
        self.vareps = vareps
        
    def _chooseAction(self,x,y):
        RonnieSoak = np.random.rand(1)
        if RonnieSoak < self.vareps:
            a = np.random.randint(4)
        else:
            localQ = self.hatQ[x,y,:]
            a = np.argmax(localQ)
        return(a)
        
    def learn(self, gamma, trainCycles):
        for i in range(trainCycles):
            (x,y) = self.env.sensors()
            a = self._chooseAction(x,y)
            r = self.env.move(a)
            (xN,yN) = self.env.sensors()
            rNext = self.hatQ[xN,yN,self._chooseAction(xN,yN)]
            self.hatQ[x,y,a] = r + gamma * rNext
            if r == 1000: self.env.setDog()
            
    def show(self, x=0, y=3):
        self.env.setDog(x,y)
        r = -1
        counter = 0 
        while r != 1000 and counter < 100:
            counter += 1
            (x,y) = self.env.sensors()
            a = self._chooseAction(x,y)
            r = self.env.move(a)
            print(a,r)

if __name__ == '__main__':
    np.random.seed(42)
    myenv = environment()
    robohund = learningDogSARSA(myenv,0.1)
    robohund.learn(0.4,2000)
    print('Start bei 0,3')
    robohund.show()
    print('Start bei 3,0')
    robohund.show(x=3,y=0)
                
                
    