"""
在多臂赌博机的实际环境下测试贪心策略，玻尔兹曼策略，UCB策略。
第一步：要定义一些多臂赌博机系统的基本信息。

第二步：训练
    a,选动作，以怎么样的方式选策略。例如：贪心策略，玻尔兹曼策略，UCB策略。
    b,将动作发给环境，环境给出即时奖励。
第三步：做图，x：玩家玩游戏的总的次数，y：累加奖励。
            i,                 sum_reward(i)
"""
# 第一步：要定义一些多臂赌博机系统的基本信息。
import matplotlib.pyplot as plt
import numpy as np
class Gamble(object):
    def __init__(self):
        """
        多臂赌博机系统的基本信息。
            1.多臂赌博机的基本信息(假设只有三台赌博机)
            2.玩家玩的总场数。counts
            3.玩家发出的动作。action
            4.当前赌博机获得的即时奖励。current_reward
            5.第i台赌博机的总的奖励。sum_reward
            6.第i台赌博机的台t-1时刻的平均奖励。average_reward,当前时刻为t时刻。
            7.每个臂摇动的次数。
            8.总奖励total_reward
        """
        self.counts = 0# 玩家一共完了多少场游戏。
        self.action = np.array([1, 2, 3])# 玩家可以摇动那台赌博机。
        self.current_reward = 0.0
        self.sum_reward = np.array([0.0, 0.0, 0.0])# 三台赌博机累计奖励。
        self.average_reward = np.array([0.0, 0.0, 0.0])# 三台赌博机t-1时刻的平均reward。
        self.total_visits = np.array([0, 0, 0])# 每个赌博机被摇动的次数。
        self.total_reward = 0.0

        self.sum_reward_history = []# 记录第i的奖励累加值
        self.visit_history = [] # i
    def reset(self):
        self.counts = 0  # 玩家一共完了多少场游戏。
        self.action = np.array([1, 2, 3])  # 玩家可以摇动那台赌博机。
        self.current_reward = 0
        self.sum_reward = np.array([0.0, 0.0, 0.0])  # 三台赌博机累计奖励。
        self.average_reward = np.array([0.0, 0.0, 0.0])  # 三台赌博机t-1时刻的平均reward。
        self.total_visits = np.array([0, 0, 0])  # 每个赌博机摇动的次数。
        self.total_reward = 0

    def choose_action(self,policy, **kwargs):
        action = 0
        # 按照贪心算法，选择action
        if policy == "e_greedy":
            one = np.random.random()
            epsilon = kwargs["epsilon"]
            if one < epsilon:
                action = np.random.randint(1,4)
            else:
                action = np.argmax(self.average_reward) + 1
        if policy == "UCB":
            c_ratio = kwargs["c_ratio"]

            if 0 in self.total_visits:
                action = np.where(self.total_visits==0)[0][0]+1
            else:# 传播用的不赖
                value = self.average_reward + c_ratio*np.sqrt(np.log(self.counts)/self.total_visits)
                action = np.argmax(value) + 1

        return action
    def step(self, a):
        r = 0
        if a == 1:
            r = np.random.normal(1,1)
        elif a == 2:
            r = np.random.normal(2, 1)
        elif a == 3:
            r = np.random.normal(1.5, 1)
        return r

    def plot(self, colors, policy, style):

        print(policy, self.average_reward)
        print("三次动作的次数：", self.total_visits)
        plt.figure(1)

        plt.plot(self.visit_history, self.sum_reward_history, colors, label=policy, linestyle=style)
        plt.legend()
        plt.xlabel("visit_history", fontsize=20)
        plt.ylabel("sum_reward_history", fontsize=20)

    def train(self, play_total, policy, **kwargs):
        """
        训练：a,玩家选择action,b环境交互。
        :return:
        """
        for i in range(play_total):
            t_action = 0
            # 玩家以怎样的方式选策略。
            if policy == 'e_greedy':
                t_action = self.choose_action(policy, epsilon=kwargs['epsilon'])
            elif policy == "UCB":
                t_action = self.choose_action(policy, c_ratio=kwargs["c_ratio"])
            if policy == "boltzmann":
                tau = kwargs["temperature"]
                p = np.exp(self.average_reward/tau)/(np.sum(np.sum(np.exp(self.average_reward/tau))))
                t_action = np.random.choice([1,2,3], p = p.ravel())

            self.current_reward = self.step(t_action)
            self.counts += 1
            # 计算每个臂摇动几次, 不用累加了
            self.average_reward[t_action-1] = \
                (self.average_reward[t_action-1]*self.total_visits[t_action-1]+self.current_reward)/(self.total_visits[t_action-1]+1)

            self.sum_reward[t_action-1] += self.current_reward
            self.total_visits[t_action - 1] += 1

            self.total_reward += self.current_reward
            self.sum_reward_history.append(self.total_reward)
            self.visit_history.append(i)







if __name__ == "__main__":
    # 第一步：定义多臂赌博机系统的基本变量。
    np.random.seed(0)
    k_gamble = Gamble()
    print("k_gamble:", k_gamble)
    total = 2000

    # 第二步：train
    k_gamble.train(play_total=total, policy='e_greedy', epsilon=0.05)
    # 第三步：画图
    k_gamble.plot(colors="r", policy="e_greedy", style="-.")
    # 第四步：重置
    k_gamble.reset()

    k_gamble.train(play_total=total, policy="UCB", c_ratio=0.2)
    k_gamble.plot(colors="g", policy="UCB", style="-")
    k_gamble.reset()
    k_gamble.train(play_total=total, policy="boltzmann", temperature=1)
    k_gamble.plot(colors="b", policy="botzmann", style="--")
    k_gamble.reset()
    plt.show()