import numpy as np
from agents import BayesianAgent, Ucb1Agent
from mab import Environment

def main():

    env = Environment([0.1, 0.2, 0.3])
    num_plays = 100000
    agent = BayesianAgent(env)
    agent_u = Ucb1Agent(env, verbose=False)

    data = np.empty(num_plays)
    data_u = np.empty(num_plays)
    for i in range(num_plays):
        data[i] = agent.play(i+1)
        data_u[i] = agent_u.play(i+1)

    print("Total Reward (Bayesian): {}".format(np.sum(data)))
    for i, bs in enumerate(agent.bandit_stats):
        print("Bandit {} (Bayesian): Attempts {}, Reward {}".format(i, bs.attempts, bs.sum_of_x))
    print("Total Reward (UCB1): {}".format(np.sum(data_u)))
    for i, bs in enumerate(agent_u.bandit_stats):
        print("Bandit {} (UCB1): Attempts {}, Reward {}".format(i, bs.attempts, bs.sum_of_x))


if __name__ == "__main__":
    main()
