import numpy as np

##################################################
# RL-lab
##################################################
#
# Find the way out of the labyrinth
#
##################################################

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.ones(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)))


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

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

def to_this_door(door_last_room):
	q, r = divmod(door_last_room + 2, 4)
	return r

def choose_door(P):
	Pn = P / np.sum(P)
	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 normalize(P):
	for x in range(0, lab_size_x):
		for y in range(0, lab_size_y):
			pos = (x,y)
			P0 = P[pos]
			P[pos] = P0/np.sum(P0)

	return P

def forget(forgetting_factor):
	for x in range(0, lab_size_x):
		for y in range(0, lab_size_y):
			pos = (x,y)
			P_door[pos,:] *= forgetting_factor

def run(pos, N_trials, learning_rate, penalty_factor=0.99, forgetting_factor=0.99):
	lab_visited = np.zeros((lab_size_x, lab_size_y))
	moves_needed = 0
	door_last = -1
	for n in range(0, N_trials):

		if lab[pos] == -1:
			break

#		forget(forgetting_factor)

		P_door[pos,:] *= forgetting_factor
		while True:
			fail = False
			door = choose_door(P_door[pos])
			if door_last >= 0:
				this_door = to_this_door(door_last)
				if this_door == door:
					P_door[pos][door] *= (1.0 - learning_rate)

			pos_new = tuple(pos + move[door + 1])

			if pos_new[0] < 0 or pos_new[1] < 0:
				fail = True

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

			try:
				if lab[pos_new] == 0:
					fail = True
			except:
				pass

			if fail:
				P_door[pos][door] *= penalty_factor
				continue
			else:
				P_door[pos][door] = P_door[pos][door] + learning_rate
			break


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

	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.5)
	print ('Visited map after {} moves:'.format(moves_needed))
#	print (lab_visited)
#	print ('P_door after {} moves:'.format(moves_needed))
#	print (P_door)

moves_needed, lab_visited = run(pos=(0,0), N_trials=1000, learning_rate=0.0)
print ('Finished after {} moves'.format(moves_needed))
print('Visited map after {} moves:'.format(moves_needed))
print(lab_visited)
print ('P_door after {} moves:'.format(moves_needed))

for i in range(0, num_doors_per_room):
	print (P_door[:,:,i])