"""Basic reinforcement learning functions.
"""

from env import reward
import numpy as np
import env

actions = ["left", "right", "stay"]
gamma = 0.9
table = np.zeros((28, 28, 3))


def all_zero(list):
    """Return true if all item in an array are 0."""
    for item in list:
            for x in item:
                if x != 0:
                    return False
    return True


def choose_action(state):
    """Choose action based on q table."""
    if np.random.uniform() > gamma or all_zero(state):
        action = np.random.choice(actions)
        print("zero")
        return action
    else:
        action = np.argmax(state)
        action += 1
        index_1 = int(round(action / 3))
        index_2 = int(action - index_1 * 3)
        index_1 -= 1
        index_2 -= 1
        action = state[index_1][index_2]
        return action

def q(state, action):
    """Update the q table."""
    index = -1
    for item in actions:
        index += 1
        if item == action:
            action_n = index

    # q learning
    new_value = reward(action) + gamma * np.amax(table[state])

    print(new_value)
    table[env.object[0]][state][action_n] = new_value
    print("special")
    print(table[env.object[0]][state])

@@ -145,9 +145,40 @@
 = 0.
-9
+8 # make the agent explore a bit
 %0Atab
@@ -525,24 +525,49 @@
 ero(state):%0A
+    #if all_zero(state):%0A
         acti
@@ -823,21 +823,29 @@
     
+if 
 index_
-2 -= 1%0A
+1 == 0:%0A    
     
@@ -861,31 +861,112 @@
 n = 
-state%5Bindex_1%5D%5Bindex_2%5D
+%22left%22%0A        elif index_1 == 1:%0A            action = %22right%22%0A        else:%0A            action = %22stay%22
 %0A   
@@ -1033,24 +1033,92 @@
 q table.%22%22%22%0A
+    print(%22%5B+%5D INFO object:%22, env.object%5B0%5D, %22player:%22, env.player)%0A
     index = 
@@ -1293,21 +1293,29 @@
 x(table%5B
-state
+env.object%5B0%5D
 %5D)%0A%0A    
@@ -1333,95 +1333,73 @@
 alue
-)%0A    table%5Benv.object%5B0%5D%5D%5Bstate%5D%5Baction_n%5D = new_value%0A    print(%22special%22)%0A    print(
+, state, action)%0A    print(table%5Benv.object%5B0%5D%5D%5Bint(state)%5D)%0A    
 tabl
@@ -1419,12 +1419,38 @@
 0%5D%5D%5B
+int(
 state
-%5D
 )
+%5D%5Baction_n%5D = new_value
 %0A

Update choose_action function.