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})

REWARD_VARIANCE = 0.1

def test(_k_arms: int, _episode_len: int, _bias: np.array, _param: tuple[float, float]) -> np.array:
	_epsilon = _param[0]		# Anti-greediness (ability to explore)
	_rho = _param[1]			# reduce epsilon with age
	_r_list = []
	_a_list = []
	_opt_a_list = []
	# Init Q and N
	_qu = np.zeros(_k_arms)
	_nu = np.zeros(_k_arms)

	# Calc z in advance
	z = np.random.uniform(size=_episode_len)

	# Calc a_expl in advance
	a_expl = np.random.randint(low=0, high=_k_arms, size=_episode_len)
	for j in range(0, _episode_len):

		# choose action
		_a = np.argmax(_qu) if z[j] > _epsilon else a_expl[j]

		# get reward from bandit
		_reward_list = REWARD_VARIANCE*np.random.normal(size=_k_arms) + _bias

		# get reward from action
		_reward = _reward_list[_a]

		# calc
		_nu[_a] = _nu[_a] + 1
		_qu[_a] = _qu[_a] + (_reward - _qu[_a]) / _nu[_a]

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

		# Statistics
		_a_list.append(_a)
		_r_list.append(_reward)
		opt_a = 1 if _a == np.argmax(_reward_list) else 0
		_opt_a_list.append(opt_a)

	return np.array(_r_list), np.array(_a_list), np.array(_opt_a_list)


if __name__ == '__main__':
	# Init parameters
	k_arms = 10
	highest_reward = 1.5
	num_realisations = 2000
	episode_len = 1000

	params = [(0.0, 0.0), (0.01, 0.002), (0.1, 0.002)]
	legend = []
	for param in params:
		# Init stats
		r_mean = np.zeros(episode_len)
		o_mean = np.zeros(episode_len)
		for realization in range(0, num_realisations):
			# init bandits with different biases for shifting reward probability
			# -> expected reward q*(a)
			ql_star = np.random.uniform(-highest_reward, +highest_reward, k_arms)
			r_list, a_list, o_list = test(_k_arms=k_arms, _episode_len=episode_len, _bias=ql_star, _param=param)
			r_mean += r_list
			o_mean += o_list

		pl.plot(r_mean/num_realisations)
		legend.append(f"Param {param}")

	pl.grid()
	pl.title(f"Norm. E(R), num. realizations: Z={num_realisations}, num. bandit arms: K={k_arms}")
	pl.xlabel("Episode")
	pl.ylabel("E(R)/Z")
	pl.legend(legend)
	ax = pl.gca()
	ax.set_xlim([-episode_len/10, episode_len])
	ax.set_ylim([0, highest_reward])

	pl.show()