import numpy as np
import matplotlib.pyplot as pl
from typing import Callable

float_formatter = "{:.3f}".format
np.set_printoptions(formatter={'float_kind': float_formatter})

K = 10


def bandit(a: int, bias: np.array):
	zn = np.random.uniform() + bias[a]
	return zn


def test(episode_len: int, epsilon: float, qu: np.array, nu: np.array, bandit: Callable) -> np.array:
	r_sum = 0
	r_mean = []
	n = 0
	for j in range(0, episode_len):
		# choose action
		z = np.random.uniform()
		if z <= epsilon:
			# Choose random action
			a = np.random.randint(low=0, high=K)
		else:
			# Choose best action
			a = np.argmax(qu)

		# get reward from bandit
		r = bandit(a)
		nu[a] = nu[a] + 1
		qu[a] = qu[a] + (r - qu[a]) / nu[a]

		# Reduce tendency to explore with number of steps (or with age for humans)
		epsilon = epsilon * (1 - rho)

		# Statistics
		r_sum += r
		n += 1
		r_mean.append(r_sum/n)

	return np.array(r_mean)


if __name__ == '__main__':

	epsilon = 0.5		# Anti-greediness
	rho = 0.005		# reduce epsilon with age

	# init bandits with different biases for shifting reward probability
	# -> expected reward q*(a)
	ql_star = np.array(np.linspace(-0.5, +0.5, K))

	num_realisations = 10
	episode_len = 10000
	r_mean = np.zeros(episode_len)
	for age in range(0, num_realisations):
		# Init Q and N
		qu = np.zeros(K)
		nu = np.zeros(K)

		r_mean += test(episode_len, epsilon, qu, nu, bandit=lambda a: bandit(a, ql_star))

	print(f"ql_star  = {ql_star}")
	print(f"qu       = {qu}")
	print(f"nu       = {nu}")

	pl.plot(r_mean/num_realisations)
	pl.grid()
	pl.title(f"Norm. E(R), num. realizations: Z={num_realisations}, num. bandits: K={K}")
	pl.xlabel("Episode")
	pl.ylabel("E(R)/Z")
	ax = pl.gca()
	ax.set_xlim([-episode_len/10, episode_len])
	ax.set_ylim([0, 1])

	pl.show()