import gym
import numpy as np


def nstep_on_policy_return(v, done, states, rewards):
    """ Calculate un-discounted n-step on-policy return per (7.1) """
    assert len(states) - 1 == len(rewards)

    if not rewards:
        # Append value of the n-th state unless in the termination phase
        return 0 if done else v[states[0]]

    sub_return = nstep_on_policy_return(v, done, states[1:], rewards[1:])
    return rewards[0] + sub_return


def td_on_policy_prediction(env, policy, n, num_episodes, alpha=0.5, tderr=False):
    """ n-step TD algorithm for on-policy value prediction per Chapter 7.1. Value function updates are
     calculated by summing TD errors per Exercise 7.2 (tderr=True) or with (7.2) (tderr=False). """
    assert type(env.action_space) == gym.spaces.Discrete
    assert type(env.observation_space) == gym.spaces.Discrete

    # Number of available actions and states
    n_state, n_action = env.observation_space.n, env.action_space.n,
    assert policy.shape == (n_state, n_action)

    # Initialization of value function
    v = np.ones([n_state], dtype=np.float) * 0.5

    history = []
    for episode in range(num_episodes):
        # Reset the environment and initialize n-step rewards and states
        state = env.reset()
        nstep_states = [state]
        nstep_rewards = []

        dv = np.zeros_like(v)

        done = False
        while nstep_rewards or not done:
            if not done:
                # Step the environment forward and check for termination
                action = np.random.choice(n_action, p=policy[state])
                state, reward, done, info = env.step(action)

                # Accumulate n-step rewards and states
                nstep_rewards.append(reward)
                nstep_states.append(state)

                # Keep accumulating until there's enough for the first n-step update
                if len(nstep_rewards) < n:
                    continue
                assert len(nstep_states) - 1 == len(nstep_rewards) == n

            # Calculate un-discounted n-step return per (7.1)
            v_target = nstep_on_policy_return(v, done, nstep_states, nstep_rewards)

            if tderr is True:
                # Accumulate TD errors over the episode while v is kept constant per Exercise 7.2
                dv[nstep_states[0]] += alpha * (v_target - v[nstep_states[0]])
            else:
                # Update value function toward the target per (7.2)
                v[nstep_states[0]] += alpha * (v_target - v[nstep_states[0]])

            # Remove the used n-step reward and states
            del nstep_rewards[0]
            del nstep_states[0]

            # Update value function with the sum of TD errors accumulated during the episode
            v += dv

        history += [np.copy(v)]
    return history

import gym
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

import randomwalk
from td import td_on_policy_prediction

env = gym.make('RandomWalk-v0')
policy = np.ones([env.observation_space.n, env.action_space.n], dtype=np.float)
size = env.observation_space.n

n = 4
history0 = td_on_policy_prediction(env, policy, n=n, num_episodes=10,
                                   alpha=0.1, tderr=False)
history1 = td_on_policy_prediction(env, policy, n=n, num_episodes=10,
                                   alpha=0.1, tderr=True)

v_star = np.arange(1, size + 1) / (size + 1)



plt.figure()
plt.title(f"Learning Curves (n={n})")
plt.xlabel("Episodes")
plt.ylabel("RMS Error")
rms0 = np.sqrt((np.mean(history0 - v_star, axis=1) ** 2))
rms1 = np.sqrt((np.mean(history1 - v_star, axis=1) ** 2))
plt.plot(rms0, label='Online')
plt.plot(rms1, label='TD Error Sum')
plt.legend()
plt.show()


plt.figure()
plt.title(f"Value Function (n={n})")
plt.xlabel("State")
plt.ylabel("Value")
plt.xticks(range(size + 1), [chr(ord('A') + i) for i in range(size + 1)])
plt.plot(history0[-1], 'o-', label='Online')
plt.plot(history1[-1], 'x-', label='TD Error Sum')
plt.plot(v_star, 'o--', color='gray', label='Theoretical')
plt.legend()
plt.show()

ns = [2, 4, 8]
alphas = np.linspace(0, 0.5, num=20)
rms10int0 = np.zeros([len(ns), len(alphas)])
rms10int1 = np.zeros([len(ns), len(alphas)])

for i, n in enumerate(ns):
    for j, alpha in enumerate(alphas):
        env.seed(7)
        history0 = td_on_policy_prediction(env, policy, n=n, num_episodes=10,
                                           alpha=alpha, tderr=False)
        rmsint0 = np.sqrt((history0 - v_star) ** 2).mean()
        rms10int0[i, j] = rmsint0

        env.seed(7)
        history1 = td_on_policy_prediction(env, policy, n=n, num_episodes=10,
                                           alpha=alpha, tderr=True)
        rmsint1 = np.sqrt((history1 - v_star) ** 2).mean()
        rms10int1[i, j] = rmsint1



matplotlib.rcParams['figure.figsize'] = [10, 10]

plt.figure()
plt.title(f"RMS Mean")
plt.xlabel(r"$\alpha$"); plt.ylim([0, 0.55])
plt.ylabel("RMS Mean Error Over First 10 Episodes")
for i, n in enumerate(ns):
    plt.plot(alphas, rms10int0[i, :], 'o-', label=f'Online (n={n})')
for i, n in enumerate(ns):
    plt.plot(alphas, rms10int1[i, :], 'x-', label=f'TD Error Sum (n={n})')
plt.legend(loc=0)
plt.show()

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from tqdm import tqdm

# all states
N_STATES = 19

# discount
GAMMA = 1

# all states but terminal states
STATES = np.arange(1, N_STATES + 1)

# start from the middle state
START_STATE = 10

# two terminal states
# an action leading to the left terminal state has reward -1
# an action leading to the right terminal state has reward 1
END_STATES = [0, N_STATES + 1]

# true state value from bellman equation
TRUE_VALUE = np.arange(-20, 22, 2) / 20.0
TRUE_VALUE[0] = TRUE_VALUE[-1] = 0

# n-steps TD method
# @value: values for each state, will be updated
# @n: # of steps
# @alpha: # step size
def temporal_difference(value, n, alpha):
    # initial starting state
    state = START_STATE

    # arrays to store states and rewards for an episode
    # space isn't a major consideration, so I didn't use the mod trick
    states = [state]
    rewards = [0]

    # track the time
    time = 0

    # the length of this episode
    T = float('inf')
    while True:
        # go to next time step
        time += 1

        if time < T:
            # choose an action randomly
            if np.random.binomial(1, 0.5) == 1:
                next_state = state + 1
            else:
                next_state = state - 1

            if next_state == 0:
                reward = -1
            elif next_state == 20:
                reward = 1
            else:
                reward = 0

            # store new state and new reward
            states.append(next_state)
            rewards.append(reward)

            if next_state in END_STATES:
                T = time

        # get the time of the state to update
        update_time = time - n
        if update_time >= 0:
            returns = 0.0
            # calculate corresponding rewards
            for t in range(update_time + 1, min(T, update_time + n) + 1):
                returns += pow(GAMMA, t - update_time - 1) * rewards[t]
            # add state value to the return
            if update_time + n <= T:
                returns += pow(GAMMA, n) * value[states[(update_time + n)]]
            state_to_update = states[update_time]
            # update the state value
            if not state_to_update in END_STATES:
                value[state_to_update] += alpha * (returns - value[state_to_update])
        if update_time == T - 1:
            break
        state = next_state

# Figure 7.2, it will take quite a while
def figure7_2():
    # all possible steps
    steps = np.power(2, np.arange(0, 10))

    # all possible alphas
    alphas = np.arange(0, 1.1, 0.1)

    # each run has 10 episodes
    episodes = 10

    # perform 100 independent runs
    runs = 100

    # track the errors for each (step, alpha) combination
    errors = np.zeros((len(steps), len(alphas)))
    for run in tqdm(range(0, runs)):
        for step_ind, step in enumerate(steps):
            for alpha_ind, alpha in enumerate(alphas):
                # print('run:', run, 'step:', step, 'alpha:', alpha)
                value = np.zeros(N_STATES + 2)
                for ep in range(0, episodes):
                    temporal_difference(value, step, alpha)
                    # calculate the RMS error
                    errors[step_ind, alpha_ind] += np.sqrt(np.sum(np.power(value - TRUE_VALUE, 2)) / N_STATES)
    # take average
    errors /= episodes * runs

    for i in range(0, len(steps)):
        plt.plot(alphas, errors[i, :], label='n = %d' % (steps[i]))
    plt.xlabel('alpha')
    plt.ylabel('RMS error')
    plt.ylim([0.25, 0.55])
    plt.legend()

    plt.savefig('../rl7/images/figure_7_2.png')
    plt.close()

if __name__ == '__main__':
    figure7_2()

def tree_backup(env, num_episodes, n, alpha, gamma=1.0):
    """
    env: 游戏环境
    num_episodes:总的训练幕数
    alpha: 增量步长
    gamma: 折扣因子
    return:预测的状态价值
    """
    # 初始化存储状态-动作轨迹列表
    state_action_trajectory = np.zeros((n+1, 2), dtype=int)
    # 初始化收益值的轨迹列表
    reward_trajectory = np.zeros(n)
    # 环境总的可选动作数
    nA = env.action_space.n
    # 初始化状态-动作价值表
    Q = defaultdict(lambda: np.zeros(nA))
    # 开始训练迭代
    for each_episode in range(1, num_episodes + 1):
        # 进度显示
        print("Episode {}/{}".format(each_episode, num_episodes), end="\r")
        sys.stdout.flush()
        # 环境重置
        state = env.reset()
        # 设置逐渐递减的epsilon
        epsilon = 1 / each_episode
        # 随机选择动作
        action = np.random.randint(4)
        # 记录初始时刻的状态-动作
        state_action_trajectory[0] = [state, action]
        # 初始化采样时间
        t = 0
        T = float('inf')
        while True:
            # 若非幕终止则继续动作
            if t < T:
                # 执行动作
                next_state, reward, done, info = env.step(action)
                reward_trajectory[np.mod(t, n)] = reward
                # 如果下个时刻幕终止则记录终止时刻与状态
                if done:
                    T = t + 1
                    next_action = action
                # 否则记录下状态-动作二元组
                else:
                    next_action = np.random.randint(4)
                    action = next_action
                state_action_trajectory[np.mod(t + 1, n + 1)] = [next_state, next_action]
            # 定义更新时刻
            tau = t - n + 1
            # 最开始的n-1个时刻后开始更新，
            if tau >= 0:
                # 记录下一采样时刻的回报估计值
                if t+1 >= T:
                    returns = reward_trajectory[np.mod(T-1, n)]
                else:
                    next_state_value = sum(greedy_pro(Q, next_state, act, nA, epsilon) * Q[next_state][act] for act in range(4))
                    returns = reward_trajectory[np.mod(t, n)] + gamma * next_state_value
                for k in range(min(t, T-1), tau+1, -1):
                    state_k = state_action_trajectory[np.mod(k, n + 1), 0]
                    action_k = state_action_trajectory[np.mod(k, n + 1), 1]
                    next_state_value_part = sum(greedy_pro(Q, state_k, act, nA, epsilon) * Q[state_k][act] if act!=action_k else 0 for act in range(4))
                    returns = reward_trajectory[np.mod(k-1, n)] + gamma * next_state_value_part + gamma * greedy_pro(Q, state_k, action_k, nA, epsilon) * returns
                state_tau = state_action_trajectory[np.mod(tau, n + 1), 0]
                action_tau = state_action_trajectory[np.mod(tau, n + 1), 1]
                Q[state_tau][action_tau] += alpha * (returns - Q[state_tau][action_tau])
            if tau == T - 1:
                break
            t += 1
    return Q