import numpy as np


class BernoulliBandit:
    def __init__(self, true_means, seed=0):
        self.true_means = np.array(true_means)
        self.rng = np.random.default_rng(seed)
        self.best_mean = self.true_means.max()

    def pull(self, arm):
        return float(self.rng.random() < self.true_means[arm])


def epsilon_greedy_bandit(bandit, n_rounds, epsilon=0.1, seed=1):
    rng = np.random.default_rng(seed)
    k = len(bandit.true_means)
    counts = np.zeros(k)
    means = np.zeros(k)
    regrets = []
    for t in range(1, n_rounds + 1):
        if rng.random() < epsilon:
            arm = rng.integers(k)
        else:
            arm = int(np.argmax(means))
        reward = bandit.pull(arm)
        counts[arm] += 1
        means[arm] += (reward - means[arm]) / counts[arm]
        regrets.append(bandit.best_mean - bandit.true_means[arm])
    return np.cumsum(regrets)


def ucb1_bandit(bandit, n_rounds):
    k = len(bandit.true_means)
    counts = np.zeros(k)
    means = np.zeros(k)
    regrets = []
    for t in range(1, n_rounds + 1):
        if t <= k:
            arm = t - 1  # pull each arm once to initialize
        else:
            bonus = np.sqrt(2 * np.log(t) / counts)
            arm = int(np.argmax(means + bonus))
        reward = bandit.pull(arm)
        counts[arm] += 1
        means[arm] += (reward - means[arm]) / counts[arm]
        regrets.append(bandit.best_mean - bandit.true_means[arm])
    return np.cumsum(regrets)


def thompson_sampling_bandit(bandit, n_rounds, seed=2):
    rng = np.random.default_rng(seed)
    k = len(bandit.true_means)
    alpha = np.ones(k)  # Beta(1,1) uniform prior per arm
    beta = np.ones(k)
    regrets = []
    for t in range(n_rounds):
        samples = rng.beta(alpha, beta)
        arm = int(np.argmax(samples))
        reward = bandit.pull(arm)
        if reward > 0.5:
            alpha[arm] += 1
        else:
            beta[arm] += 1
        regrets.append(bandit.best_mean - bandit.true_means[arm])
    return np.cumsum(regrets)
