import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

# Alternate implementation from:
# https://www.kaggle.com/code/parsasam/reinforcement-learning-notes-multi-armed-bandits

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

k = 10
num_problems = 200
num_steps = 2000

q_star = np.random.normal(0, 1, (num_problems, k))
arms = [0] * k

for i in range(10):
	arms[i] = np.random.normal(q_star[0, i], 1, 2000)  # first problem as a sample

plt.figure(figsize=(12, 8))
plt.ylabel('Rewards distribution')
plt.xlabel('Actions')
plt.xticks(range(1, 11))
plt.yticks(np.arange(-5, 5, 0.5))

plt.violinplot(arms, positions=range(1, 11), showmedians=True)
plt.show()

TIE_BREAK = 0.01

def bandit(action, problem):
	return np.random.normal(q_star[problem, action], 1)


def simple_max(Q, N, t):
	am = np.argmax(Q)
	fm = Q == Q.max()
	ffm = np.flatnonzero(fm)
	return np.random.choice(ffm)  # breaking ties randomly


def simple_bandit(k, epsilon, steps, initial_Q, alpha=0, argmax_func=simple_max):
	rewards = np.zeros(steps)
	actions = np.zeros(steps)

	for i in tqdm(range(num_problems)):
		Q = np.ones(k) * initial_Q  # initial Q
		N = np.zeros(k)  # initalize number of rewards given
		best_action = np.argmax(q_star[i])
		for t in range(steps):
			if np.random.rand() < epsilon:  # explore
				a = np.random.randint(k)
			else:  # exploit
				a = argmax_func(Q, N, t)

			reward = bandit(a, i)

			N[a] += 1
			if alpha > 0:
				Q[a] = Q[a] + (reward - Q[a]) * alpha
			else:
				Q[a] = Q[a] + (reward - Q[a]) / N[a]

			rewards[t] += reward

			if a == best_action:
				actions[t] += 1

	return np.divide(rewards, num_problems), np.divide(actions, num_problems)


def main():
	ep_0, ac_0 = simple_bandit(k=10, epsilon=0, steps=num_steps, initial_Q=0)
	ep_01, ac_01 = simple_bandit(k=10, epsilon=0.01, steps=num_steps, initial_Q=0)
	ep_1, ac_1 = simple_bandit(k=10, epsilon=0.1, steps=num_steps, initial_Q=0)

	plt.figure(figsize=(12,6))
	plt.plot(ep_0, 'g', label='epsilon = 0')
	plt.plot(ep_01, 'r', label='epsilon = 0.01')
	plt.plot(ep_1, 'b', label='epsilon = 0.1')
	plt.legend()
	plt.show()


if __name__ == '__main__':
	main()