if __name__ == "__main__":
    true_means = [0.3, 0.5, 0.55, 0.4, 0.6]  # arm 4 is best
    n_rounds = 5000

    bandit_eg = BernoulliBandit(true_means, seed=0)
    regret_eg = epsilon_greedy_bandit(bandit_eg, n_rounds, epsilon=0.1)

    bandit_ucb = BernoulliBandit(true_means, seed=0)
    regret_ucb = ucb1_bandit(bandit_ucb, n_rounds)

    bandit_ts = BernoulliBandit(true_means, seed=0)
    regret_ts = thompson_sampling_bandit(bandit_ts, n_rounds)

    print("Cumulative regret after", n_rounds, "rounds:")
    print(f"  epsilon-greedy:    {regret_eg[-1]:.1f}")
    print(f"  UCB1:              {regret_ucb[-1]:.1f}")
    print(f"  Thompson sampling: {regret_ts[-1]:.1f}")

    print("\nRegret growth rate (last 1000 vs first 1000 rounds):")
    for name, regret in [
        ("epsilon-greedy", regret_eg),
        ("UCB1", regret_ucb),
        ("Thompson sampling", regret_ts),
    ]:
        early_rate = regret[999] - regret[0]
        late_rate = regret[-1] - regret[-1001]
        print(f"  {name}: early {early_rate:.1f}, late {late_rate:.1f}")
