from __future__ import print_function, absolute_import, division

import numpy as np
import matplotlib.pyplot as plt
from bandit.comparing_epsilons import Bandit, run_experiment as run_experiment_eps
from bandit.optimistic_epsilons import run_experiment as run_experiment_opt
from bandit.ucb1 import run_experiment as run_experiment_ucb

xrange = range


class BayesianBandit:
    def __init__(self, m):
        """
        :param m:  mean of the Bandit
        """
        self.m = m
        self.m0 = 0
        self.lambda0 = 1
        self.sum_x = 0
        self.tau = 1

    def pull(self):
        return np.random.randn() + self.m

    def sample(self):
        return np.random.randn() / np.sqrt(self.lambda0) + self.m0

    def update(self, x):
        # assume tau is 1
        self.lambda0 += 1
        self.sum_x += x
        self.m0 = self.tau * self.sum_x / self.lambda0  # (1 - 1.0 / self.N) * self.mean + 1.0 / self.N * x


def run_experiment(m1, m2, m3, N):
    # create bandit
    bandits = [BayesianBandit(m1), BayesianBandit(m2), BayesianBandit(m3)]

    data = np.empty(N)

    for i in xrange(N):
        # use b.sample() to decide which to use
        j = np.argmax([b.sample() for b in bandits])

        x = bandits[j].pull()
        bandits[j].update(x)

        # for the plot
        data[i] = x

    cumulative_average = np.cumsum(data) / (np.arange(N) + 1)

    # plot moving average ctr
    plt.plot(cumulative_average)
    plt.plot(np.ones(N) * m1)
    plt.plot(np.ones(N) * m2)
    plt.plot(np.ones(N) * m3)
    plt.xscale('log')
    plt.show()

    for i, b in enumerate(bandits):
        print(i, b.m0)

    return cumulative_average

def run_experiment_decaying_epsilon(m1, m2, m3, N):
    bandits = [Bandit(m1), Bandit(m2), Bandit(m3)]
    data = np.empty(N)

    for i in xrange(N):
        # epsilon greedy, decaying epsilon
        p = np.random.random()
        if p < 1.0/(i + 1):
            j = np.random.choice(3)
        else:
            j = np.argmax([b.mean for b in bandits])

        x = bandits[j].pull()
        bandits[j].update(x)
        # for the plot
        data[i] = x

    cumulative_average = np.cumsum(data) / (np.arange(N) + 1)

    # plot moving average ctr
    plt.plot(cumulative_average)
    plt.plot(np.ones(N) * m1)
    plt.plot(np.ones(N) * m2)
    plt.plot(np.ones(N) * m3)
    plt.xscale('log')
    plt.show()

    for i, bdt in enumerate(bandits):
        print(i, bdt.mean)

    return cumulative_average


def main():
    # fixed epsilon 0.1
    c_1 = run_experiment_eps(1.0, 2.0, 3.0, 0.1, 100000)
    # decaying epsilon = 1/(i+1)
    eps = run_experiment_decaying_epsilon(1.0, 2.0, 3.0, 100000)
    # ?
    oiv = run_experiment_opt(1.0, 2.0, 3.0, 100000)
    # ucb, upper confidence bound
    ucb = run_experiment_ucb(1.0, 2.0, 3.0, 100000)
    # BayesianBandit:
    bayes = run_experiment(1.0, 2.0, 3.0, 100000)

    # log scale plot
    plt.plot(c_1, label='eps=0.1')
    plt.plot(eps, label='decaying-epsilon-greedy')
    plt.plot(oiv, label='optimistic')
    plt.plot(ucb, label='ucb')
    plt.plot(bayes, label='bayes')
    plt.legend()
    plt.xscale('log')
    plt.show()


if __name__ == '__main__':
    main()
