import numpy as np
import matplotlib.pyplot as pl
from pyparsing import alphas
from tqdm import tqdm

# For alternate implementation, see:
# https://www.kaggle.com/code/parsasam/reinforcement-learning-notes-multi-armed-bandits

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

# Init parameters
k_arms = 10
num_realisations = 2000
episode_len = 1000

REWARD_VARIANCE = 1.0

def simple_max(Q, N, t, _tie_break):
	am = np.argmax(Q + _tie_break[t])
	return am

class Param:
	def __init__(self, epsilon:float=0, rho:float=0, alpha:float =0, q_ic:float=0, _argmax_func=simple_max):

		# Anti-greediness (ability to explore)
		self.epsilon = epsilon

		# Reduce epsilon with age
		self.rho = rho

		# Constant step size
		self.alpha = alpha

		# Initial condition Q
		self.q_ic = q_ic

		# argmax function
		self.argmax_func=_argmax_func

	def __repr__(self):
		return f"epsilon={self.epsilon}, rho={self.rho}, alpha={self.alpha}, q_ic={self.q_ic}\nargmax={self.argmax_func.__repr__()}"

def test(_k_arms: int, _episode_len: int, ql_star, _tie_break, _param: Param) -> np.array:
	rewards = np.zeros(_episode_len)
	actions = np.zeros(_episode_len)

	# Init Q and N
	_qu = np.zeros(_k_arms) + _param.q_ic
	_nu = np.zeros(_k_arms)

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

	# Calc a in advance
	_a_expl = np.random.randint(_k_arms, size=_episode_len)

	# Calc rewards in advance
	_reward_z = REWARD_VARIANCE * np.random.normal(size=_episode_len)

	best_action = np.argmax(ql_star)
	for j in range(0, _episode_len):
		# choose action
		if z[j] < _param.epsilon:
			_a = _a_expl[j]
		else:
			_a = _param.argmax_func(_qu, _nu, j, _tie_break)

		# get reward from bandit
		_reward = _reward_z[j] + ql_star[_a]

		# calc
		if _param.alpha > 0:
			_qu[_a] = _qu[_a] + _param.alpha * (_reward - _qu[_a])
		else:
			_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 = _param.epsilon * (1 - _param.rho)

		# Statistics
		rewards[j] += _reward

		if _a == best_action:
			actions[j] += 1

	return rewards, actions

def batch(k_arms, num_episode, num_problems, _param: Param):
	# Init stats
	r_mean = np.zeros(episode_len)
	a_mean = np.zeros(episode_len)
	for k in tqdm(range(0, num_problems)):
		# init bandits with different biases for shifting reward probability
		# -> expected reward q*(a)
		tie_break = 0.001 * np.random.normal(size=(num_episode, k_arms))
		r, a = test(_k_arms=k_arms, _episode_len=num_episode, ql_star=q_star[k],
					_tie_break=tie_break, _param=_param)
		r_mean += r
		a_mean += a

	return r_mean, a_mean

if __name__ == '__main__':
	q_star = np.random.normal(0, 1, (num_realisations, k_arms))
	arms = np.zeros((num_realisations, k_arms))
	for i in range(k_arms):
		arms[:,i] = np.random.normal(q_star[0, i], 1, num_realisations)  # first problem as a sample

	pl.figure(figsize=(12, 8))
	pl.ylabel('Rewards distribution')
	pl.xlabel('Actions')
	pl.xticks(range(1, k_arms+1))
	pl.yticks(np.arange(-5, 5, 0.5))
	pl.violinplot(arms, positions=range(1, k_arms+1), showmedians=True)

	pl.figure(figsize=(12, 8))
#	params = [Param(epsilon=0.0), Param(epsilon=0.01), Param(epsilon=0.1)]
	params = [Param(epsilon=0.1), Param(alpha=0.2, q_ic=5)]
	legend = []
	for param in params:
		res_r, res_a = batch(k_arms, episode_len, num_realisations, param)
		legend.append(param)

		pl.subplot(2, 1, 1)
		pl.plot(res_r/num_realisations)

		pl.subplot(2, 1, 2)
		pl.plot(100*res_a/num_realisations)

	pl.subplot(2, 1, 1)
	pl.title(f"E(R) over {num_realisations} realizations, num. bandit arms: K={k_arms}")
	pl.legend(legend)
	pl.grid()
	ax = pl.gca()
	ax.set_xlim([-episode_len / 10, episode_len])
	ax.set_ylim([0, 1.6])
	pl.ylabel("E(R)")

	pl.subplot(2, 1, 2)
	pl.legend(legend)
	pl.grid()
	ax = pl.gca()
	ax.set_xlim([-episode_len / 10, episode_len])
	ax.set_ylim([0, 100])
	pl.ylabel("Percentage of optimal actions")
	pl.ylabel("%")

	pl.xlabel("Episode")

	pl.show()