import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as stats
import seaborn as sns
import time

from bandits.agent import BetaAgent


class Environment(object):
    def __init__(self, bandits, agents, label='Multi-Armed Bandit'):
        self.bandits = bandits
        self.agents = agents
        self.label = label

    def reset(self):
        for bandit in self.bandits:
            bandit.reset()
        # self.bandits.reset()
        for agent in self.agents:
            agent.reset()

    def run(self, trials=100, experiments=1):
        scores = np.zeros((trials, len(self.agents)))
        optimal = np.zeros_like(scores)

        for _ in range(experiments):
            self.reset()
            for t in range(trials):
                for i, agent in enumerate(self.agents):
                    action = agent.choose()
                    sigma_ = None
                    # if i == 3 or i == 4:
                    #     sigma_ = 10
                    reward, is_optimal = self.bandits[i].pull(action, sigma_)
                    agent.observe(reward)

                    scores[t, i] += reward
                    if is_optimal:
                        optimal[t, i] += 1

        return scores / experiments, optimal / experiments

    def plot_results(self, scores, optimal):
        sns.set_style('white')
        sns.set_context('notebook')
        plt.figure(figsize=(20, 6.18))
        plt.subplot(1, 2, 1)
        plt.title(self.label)
        plt.plot(scores, alpha=0.7)
        plt.ylabel('Average Reward')
        plt.legend(self.agents, loc='best')
        plt.subplot(1, 2, 2)
        plt.plot(optimal * 100)
        plt.ylim(0, 100)
        plt.ylabel('% Optimal Action')
        plt.xlabel('Time Step')
        plt.legend(self.agents, loc='best')
        sns.despine()
        plt.tight_layout()
        # plt.subplots_adjust(wspace=0, hspace=1)  # 调整子图间距
        plt.savefig('../figures/' + time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime()) + '.pdf', format='pdf')
        plt.show()

    def plot_beliefs(self):
        # TODO 这个中的bandits还没改正确，别忘记了
        sns.set_context('talk')
        pal = sns.color_palette("cubehelix", n_colors=len(self.agents))
        plt.title(self.label + ' - Agent Beliefs')

        rows = 2
        cols = int(self.bandits.k / 2)

        axes = [plt.subplot(rows, cols, i + 1) for i in range(self.bandits.k)]
        for i, val in enumerate(self.bandits.action_values):
            color = 'r' if i == self.bandits.optimal else 'k'
            axes[i].vlines(val, 0, 1, colors=color)

        for i, agent in enumerate(self.agents):
            if type(agent) is not BetaAgent:
                for j, val in enumerate(agent.value_estimates):
                    axes[j].vlines(val, 0, 0.75, colors=pal[i], alpha=0.8)
            else:
                x = np.arange(0, 1, 0.001)
                y = np.array([stats.beta.pdf(x, a, b) for a, b in
                              zip(agent.alpha, agent.beta)])
                y /= np.max(y)
                for j, _y in enumerate(y):
                    axes[j].plot(x, _y, color=pal[i], alpha=0.8)

        min_p = np.argmin(self.bandits.action_values)
        for i, ax in enumerate(axes):
            ax.set_xlim(0, 1)
            if i % cols != 0:
                ax.set_yticklabels([])
            if i < cols:
                ax.set_xticklabels([])
            else:
                ax.set_xticks([0, 0.25, 0.5, 0.75, 1.0])
                ax.set_xticklabels(['0', '', '0.5', '', '1'])
            if i == int(cols / 2):
                title = '{}-arm Bandit - Agent Estimators'.format(self.bandits.k)
                ax.set_title(title)
            if i == min_p:
                ax.legend(self.agents)

        sns.despine()
        plt.show()
