import numpy as np


lab_size_x = 5
lab_size_y = 5
num_doors_per_room = 4


P_door = np.ndarray((lab_size_x, lab_size_y, num_doors_per_room))
P_door = np.random.uniform(0,1,P_door.shape)

print('P_door.shape = ', P_door.shape)
lab = np.rint(np.random.uniform(0, num_doors_per_room-1, size=(lab_size_x, lab_size_y)))


#    1
#    |
# 4 -0- 2
#    |
#    3

move = np.array([(0,0), (-1,0), (0,1), (1,0), (0,-1)])
lab = np.array([
	[3, 0, 0, 0, -1],
	[3, 0, 0, 2, 1],
	[2, 2, 3, 1, 4],
	[3, 4, 4, 2, 1],
	[2, 2, 2, 1, 0]
])

def choose_door(Pn):
	Pnorm = np.cumsum(Pn)
	size = len(Pn)
	Z = np.random.uniform(0,1)

	result = None
	for n in range(0, size):
		if Z <= Pnorm[n]:
			result = n
			break

	return result

def run(pos, N_trials, learning_rate):
	lab_visited = np.zeros(P_door.shape)
	moves_needed = 0
	for n in range(0, N_trials):

#		print ("Pos={}".format(pos))
		if lab[pos] == -1:
			break

		P = P_door[pos]
		Pn = P/np.sum(P)

		while True:
			door = choose_door(Pn)
			pos_new = tuple(pos + move[door + 1])

			if pos_new[0] < 0 or pos_new[1] < 0:
				continue

			if pos_new[0] >= lab_size_x or pos_new[1] >= lab_size_y:
				continue

			if lab[pos_new] == 0:
				continue

			break

		moves_needed += 1
		k = lab_visited[pos][door]
		P[door] = max(0, P[door] - learning_rate)
		P_door[pos] = P

		pos = pos_new
		lab_visited[pos][door] += 1

	return moves_needed, lab_visited

for i in range(0, 100):
	moves_needed, lab_visited = run(pos=(0,0), N_trials=1000, learning_rate=0.001)
	print ('Visited map after {} trials: {}'.format(moves_needed, lab_visited))
	print ('P_door after {} trials: {}'.format(moves_needed, P_door))

moves_needed, lab_visited = run(pos=(0,0), N_trials=1000, learning_rate=0.0)
print ('Finished after {} trials'.format(moves_needed))

#	P_door[pos, door] += learning_rate