import gym
env = gym.make('Blackjack-v0')

"""
Explanation: 迷你项目：蒙特卡洛方法
在此 notebook 中，你将自己编写很多蒙特卡洛 (MC) 算法的实现。
虽然我们提供了一些起始代码，但是你可以删掉这些提示并从头编写代码。
第 0 部分：探索 BlackjackEnv
请使用以下代码单元格创建 Blackjack 环境的实例。
End of explanation
"""


    STICK = 0
    HIT = 1

"""
Explanation: 每个状态都是包含以下三个元素的 3 元组：
- 玩家的当前点数之和 $\in {0, 1, \ldots, 31}$，
- 庄家朝上的牌点数之和  $\in {1, \ldots, 10}$，及
- 玩家是否有能使用的王牌（no $=0$、yes $=1$）。
智能体可以执行两个潜在动作：
End of explanation
"""


print(env.observation_space)
print(env.action_space)

"""
Explanation: 通过运行以下代码单元格进行验证。
End of explanation
"""


for i_episode in range(3):
    state = env.reset()
    while True:
        print(state)
        action = env.action_space.sample()
        state, reward, done, info = env.step(action)
        if done:
            print('End game! Reward: ', reward)
            print('You won :)\n') if reward > 0 else print('You lost :(\n')
            break

"""
Explanation: 执行以下代码单元格以按照随机策略玩二十一点。
（代码当前会玩三次二十一点——你可以随意修改该数字，或者多次运行该单元格。该单元格旨在让你体验当智能体与环境互动时返回的输出结果。）
End of explanation
"""


def generate_episode_from_limit(bj_env):
    episode = []
    state = bj_env.reset()
    while True:
        action = 0 if state[0] > 18 else 1
        next_state, reward, done, info = bj_env.step(action)
        episode.append((state, action, reward))
        state = next_state
        if done:
            break
    return episode

"""
Explanation: 第 1 部分：MC 预测 - 状态值
在此部分，你将自己编写 MC 预测的实现（用于估算状态值函数）。
我们首先将研究以下策略：如果点数之和超过 18，玩家将始终停止出牌。函数  generate_episode_from_limit 会根据该策略抽取一个阶段。 
该函数会接收以下输入：
- bj_env：这是 OpenAI Gym 的 Blackjack 环境的实例。
它会返回以下输出：
- episode：这是一个（状态、动作、奖励）元组列表，对应的是 $(S_0, A_0, R_1, \ldots, S_{T-1}, A_{T-1}, R_{T})$， 其中 $T$ 是最终时间步。具体而言，episode[i] 返回 $(S_i, A_i, R_{i+1})$， episode[i][0]、episode[i][1]和 episode[i][2] 分别返回 $S_i$, $A_i$和 $R_{i+1}$。
End of explanation
"""


for i in range(3):
    print(generate_episode_from_limit(env))

"""
Explanation: 执行以下代码单元格以按照该策略玩二十一点。 
（代码当前会玩三次二十一点——你可以随意修改该数字，或者多次运行该单元格。该单元格旨在让你熟悉  generate_episode_from_limit 函数的输出结果。）
End of explanation
"""


{(4, 7, False): -0.38775510204081631, (18, 6, False): -0.58434296365330851, (13, 2, False): -0.43409090909090908, (6, 7, False): -0.3783783783783784, ...

"""
Explanation: 现在你已经准备好自己编写 MC 预测的实现了。你可以选择实现首次经历或所有经历 MC 预测；对于 Blackjack 环境，这两种技巧是对等的。
你的算法将有四个参数：
- env：这是 OpenAI Gym 环境的实例。
- num_episodes：这是通过智能体-环境互动生成的阶段次数。
- generate_episode：这是返回互动阶段的函数。
- gamma：这是折扣率。它必须是在 0 到 1（含）之间的值，默认值为：1。
该算法会返回以下输出结果：
- V：这是一个字典，其中 V[s] 是状态 s 的估算值。例如，如果代码返回以下输出结果：
End of explanation
"""


from collections import defaultdict
import numpy as np
import sys

def mc_prediction_v(env, num_episodes, generate_episode, gamma=1.0):
    # initialize empty dictionary of lists
    returns = defaultdict(list)
    # loop over episodes
    for i_episode in range(1, num_episodes+1):
        # monitor progress
        if i_episode % 1000 == 0:
            print("\rEpisode {}/{}.".format(i_episode, num_episodes), end="")
            sys.stdout.flush()
        
        ## TODO: complete the function
        
    return V

"""
Explanation: 则状态 (4, 7, False) 的值估算为  -0.38775510204081631。
如果你不知道如何在 Python 中使用 defaultdict，建议查看此源代码。
End of explanation
"""


from plot_utils import plot_blackjack_values

# obtain the value function
V = mc_prediction_v(env, 500000, generate_episode_from_limit)

# plot the value function
plot_blackjack_values(V)

"""
Explanation: 使用以下单元格计算并绘制状态值函数估算值。 (用于绘制值函数的代码来自此源代码，并且稍作了修改。）
要检查你的实现是否正确，应将以下图与解决方案 notebook Monte_Carlo_Solution.ipynb 中的对应图进行比较。
End of explanation
"""


def generate_episode_from_limit_stochastic(bj_env):
    episode = []
    state = bj_env.reset()
    while True:
        probs = [0.8, 0.2] if state[0] > 18 else [0.2, 0.8]
        action = np.random.choice(np.arange(2), p=probs)
        next_state, reward, done, info = bj_env.step(action)
        episode.append((state, action, reward))
        state = next_state
        if done:
            break
    return episode

"""
Explanation: 第 2 部分：MC 预测 - 动作值
在此部分，你将自己编写 MC 预测的实现（用于估算动作值函数）。  
我们首先将研究以下策略：如果点数之和超过 18，玩家将_几乎_始终停止出牌。具体而言，如果点数之和大于 18，她选择动作 STICK 的概率是 80%；如果点数之和不大于 18，她选择动作  HIT 的概率是 80%。函数 generate_episode_from_limit_stochastic 会根据该策略抽取一个阶段。 
该函数会接收以下输入：
- bj_env：这是 OpenAI Gym 的 Blackjack 环境的实例。
该算法会返回以下输出结果：
- episode: 这是一个（状态、动作、奖励）元组列表，对应的是 $(S_0, A_0, R_1, \ldots, S_{T-1}, A_{T-1}, R_{T})$， 其中 $T$ 是最终时间步。具体而言，episode[i] 返回 $(S_i, A_i, R_{i+1})$， episode[i][0]、episode[i][1]和 episode[i][2] 分别返回 $S_i$, $A_i$和 $R_{i+1}$。
End of explanation
"""


def mc_prediction_q(env, num_episodes, generate_episode, gamma=1.0):
    # initialize empty dictionaries of arrays
    returns_sum = defaultdict(lambda: np.zeros(env.action_space.n))
    N = defaultdict(lambda: np.zeros(env.action_space.n))
    Q = defaultdict(lambda: np.zeros(env.action_space.n))
    # loop over episodes
    for i_episode in range(1, num_episodes+1):
        # monitor progress
        if i_episode % 1000 == 0:
            print("\rEpisode {}/{}.".format(i_episode, num_episodes), end="")
            sys.stdout.flush()
        
        ## TODO: complete the function
        
    return Q

"""
Explanation: 现在你已经准备好自己编写 MC 预测的实现了。你可以选择实现首次经历或所有经历 MC 预测；对于 Blackjack 环境，这两种技巧是对等的。
你的算法将有四个参数：
- env: 这是 OpenAI Gym 环境的实例。
- num_episodes：这是通过智能体-环境互动生成的阶段次数。
- generate_episode：这是返回互动阶段的函数。
- gamma：这是折扣率。它必须是在 0 到 1（含）之间的值，默认值为：1。
该算法会返回以下输出结果：

Q：这是一个字典（一维数组），其中 Q[s][a] 是状态 s 和动作 a 对应的估算动作值。
End of explanation
"""


# obtain the action-value function
Q = mc_prediction_q(env, 500000, generate_episode_from_limit_stochastic)

# obtain the state-value function
V_to_plot = dict((k,(k[0]>18)*(np.dot([0.8, 0.2],v)) + (k[0]<=18)*(np.dot([0.2, 0.8],v))) \
         for k, v in Q.items())

# plot the state-value function
plot_blackjack_values(V_to_plot)

"""
Explanation: 请使用以下单元格获取动作值函数估值 $Q$。我们还绘制了相应的状态值函数。
要检查你的实现是否正确，应将以下图与解决方案 notebook Monte_Carlo_Solution.ipynb 中的对应图进行比较。
End of explanation
"""


def mc_control_GLIE(env, num_episodes, gamma=1.0):
    nA = env.action_space.n
    # initialize empty dictionaries of arrays
    Q = defaultdict(lambda: np.zeros(nA))
    N = defaultdict(lambda: np.zeros(nA))
    # loop over episodes
    for i_episode in range(1, num_episodes+1):
        # monitor progress
        if i_episode % 1000 == 0:
            print("\rEpisode {}/{}.".format(i_episode, num_episodes), end="")
            sys.stdout.flush()
        
        ## TODO: complete the function
        
    return policy, Q

"""
Explanation: 第 3 部分：MC 控制 - GLIE
在此部分，你将自己编写常量-$\alpha$ MC 控制的实现。
你的算法将有四个参数：

env: 这是 OpenAI Gym 环境的实例。
num_episodes：这是通过智能体-环境互动生成的阶段次数。
generate_episode：这是返回互动阶段的函数。
gamma：这是折扣率。它必须是在 0 到 1（含）之间的值，默认值为：1。

该算法会返回以下输出结果：


Q：这是一个字典（一维数组），其中 Q[s][a] 是状态 s 和动作 a 对应的估算动作值。


policy：这是一个字典，其中 policy[s] 会返回智能体在观察状态 s 之后选择的动作。


（你可以随意定义其他函数，以帮助你整理代码。）
End of explanation
"""


# obtain the estimated optimal policy and action-value function
policy_glie, Q_glie = mc_control_GLIE(env, 500000)

"""
Explanation: 通过以下单元格获取估算的最优策略和动作值函数。
End of explanation
"""


# obtain the state-value function
V_glie = dict((k,np.max(v)) for k, v in Q_glie.items())

# plot the state-value function
plot_blackjack_values(V_glie)

"""
Explanation: 接着，我们将绘制相应的状态值函数。
End of explanation
"""


from plot_utils import plot_policy

# plot the policy
plot_policy(policy_glie)

"""
Explanation: 最后，我们将可视化估算为最优策略的策略。
End of explanation
"""


def mc_control_alpha(env, num_episodes, alpha, gamma=1.0):
    nA = env.action_space.n
    # initialize empty dictionary of arrays
    Q = defaultdict(lambda: np.zeros(nA))
    # loop over episodes
    for i_episode in range(1, num_episodes+1):
        # monitor progress
        if i_episode % 1000 == 0:
            print("\rEpisode {}/{}.".format(i_episode, num_episodes), end="")
            sys.stdout.flush()
        
        ## TODO: complete the function
        
    return policy, Q

"""
Explanation: 真最优策略 $\pi_*$ 可以在该教科书的第 82 页找到（下文也提供了）。请将你的最终估算值与最优策略进行比较——它们能够有多接近？如果你对算法的效果不满意，请花时间调整 $\epsilon$ 的衰减率和/或使该算法运行更多个阶段，以获得更好的结果。

第 4 部分：MC 控制 - 常量-$\alpha$
在此部分，你将自己编写常量-$\alpha$ MC 控制的实现。  
你的算法将有三个参数：

env: 这是 OpenAI Gym 环境的实例。
num_episodes：这是通过智能体-环境互动生成的阶段次数。
generate_episode：这是返回互动阶段的函数。
alpha：这是更新步骤的步长参数。
gamma：这是折扣率。它必须是在 0 到 1（含）之间的值，默认值为：1。

该算法会返回以下输出结果：


Q：这是一个字典（一维数组），其中 Q[s][a] 是状态 s 和动作 a 对应的估算动作值。


policy：这是一个字典，其中 policy[s] 会返回智能体在观察状态 s 之后选择的动作。


（你可以随意定义其他函数，以帮助你整理代码。）
End of explanation
"""


# obtain the estimated optimal policy and action-value function
policy_alpha, Q_alpha = mc_control_alpha(env, 500000, 0.008)

"""
Explanation: 通过以下单元格获得估算的最优策略和动作值函数。
End of explanation
"""


# obtain the state-value function
V_alpha = dict((k,np.max(v)) for k, v in Q_alpha.items())

# plot the state-value function
plot_blackjack_values(V_alpha)

"""
Explanation: 接着，我们将绘制相应的状态值函数。
End of explanation
"""


# plot the policy
plot_policy(policy_alpha)

"""
Explanation: 最后，我们将可视化估算为最优策略的策略。
End of explanation
"""


