import gym
import random
import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib import pyplot as plt
from collections import OrderedDict

"""
Explanation: Cliff Walking Problem solved with TD(0) Algorithms: Implementation & Comparisons
1. Load Libraries & Define Environment
End of explanation
"""


print('OpenAI Gym environments for Cliff Walking Problem:')
[k for k in gym.envs.registry.env_specs.keys() if k.find('Cliff' , 0) >=0]

"""
Explanation: The OpenAI Gym toolkit includes the below environment for the "Cliff-Walking" problem:
End of explanation
"""


env = gym.make('CliffWalking-v0')

"""
Explanation: Load the Cliff-Walking environment:
End of explanation
"""


env.render()

"""
Explanation: This environment has to do about gridworld shown below, where the traveller initial position (x) and the target to achieve (reach T) has been flagged appropriately. In addition in a one of the edge of this gridwordld example there is a "Cliff" denoted with C. Reward is $-1$ on all transitions except those into the cliff region. Steppping into this region incurs a reward of $-100$ and sends the agent instantly back to the start. 
Once the environment is initialized you get the situation below. This is an episodic (undiscounted) task with start at traveller's starting point, and it is completed either when the goal is achieved, that is the traveller manage to reach the target location, T, or she may happen to step into the cliff. In this case the environment is reseted in each initial state.
End of explanation
"""


help(env)

"""
Explanation: Possible traveller's actions are of course her movements in this grid:
- "UP":  denoted by 0
- "RIGHT": denoted by 1
- "DOWN": denoted by 2
- "LEFT": denoted by 3
To get the new state at every next step of an episode, you may pass the current action into the .step() method of the environment. The environment then will return a tuple (observation, reward, done, info) each of which are explained as below:
- observation (object): agent's observation of the current environment
- reward (float): amount of reward returned after previous action
- done (bool): whether the episode has ended, in which case further step() calls will return undefined results
- info (dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning)
Note: At termination of each episode, the programmer is responsible to reset the environment.
For further details concerning the CliffWalking-v0 environment of OpenAI Gym toolkit consult the docstring below.
End of explanation
"""


%run ../PlotUtils.py

plotutls = PlotUtils()

"""
Explanation: 2. RL-Algorithms based on Temporal Difference TD(0)
2a. Load the "Temporal Difference" Python class
Load the Python class PlotUtils() which provides various plotting utilities and start a new instance.
End of explanation
"""


%run ../TD0_Utils.py 

"""
Explanation: Load the Temporal Difference Python class, TemporalDifferenceUtils():
End of explanation
"""


TD0 = TemporalDifferenceUtils(env)

"""
Explanation: Instantiate the class for the environment of interest:
End of explanation
"""


# Define Number of Episodes
n_episodes = 3e+3

# e-greedy parameters to investigate
print('Determine the epsilon parameters for the epsilon-greedy policy...\n')
epsilons = np.array([0.1, 0.13, 0.16])
print('epsilons: '.format(epsilons), '\n')

# various step-sizes (alpha) to try
print('Determine the step-sizes parameters (alphas) for the TD(0)...\n')
step_sizes = np.array([0.4])
print('step_sizes: {}'.format(step_sizes), '\n')

# Fixed discount
discount_fixed = 1

# Create a mesh-grid of trials
print('Create a dictionary of the RL-models of interest...\n')
epsilons, step_sizes = np.meshgrid(epsilons, step_sizes)
epsilons = epsilons.flatten()
step_sizes = step_sizes.flatten()

# Create a dictionary of the RL-trials of interest
RL_trials = {"baseline":
             {'epsilon': 0.1,
              'step_size': 0.5, 'discount': 1}}

for n, trial in enumerate(list(zip(epsilons, step_sizes))):
    key = 'trial_' + str(n+1)
    RL_trials[key] = {'epsilon': trial[0], 
                      'step_size': trial[1], 'discount': discount_fixed}
print('Number of RL-models to try: {}\n'.format(len(RL_trials)))

print('Let all RL-models to be trained for {0:,} episodes...\n'.format(int(n_episodes)))

rewards_per_trial_SARSA = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())
q_values_per_trial_SARSA = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())

for trial, params_dict in RL_trials.items():
    
    # Read out parameters from "params_dict"
    epsilon = params_dict['epsilon']
    step_size = params_dict['step_size']
    discount = params_dict['discount']
    
    # Apply SARSA [on-policy TD(0) Control]
    q_values, tot_rewards = TD0.sarsa_on_policy_control(env, 
                                                        n_episodes=n_episodes, 
                                                        step_size=step_size, discount=discount, epsilon=epsilon)
    
    # Update "rewards_per_trial" and "q_values_per_trial" OrderedDicts
    rewards_per_trial_SARSA[trial] = tot_rewards
    q_values_per_trial_SARSA[trial] = q_values

title = 'Efficiency of the RL Method\n[SARSA on-policy TD(0) Control]'
plotutls.plot_learning_curve(rewards_per_trial_SARSA, title=title, lower_reward_ratio=-100)

RL_trials

"""
Explanation: 2b. SARSA: On-Policy TD(0) Control
End of explanation
"""


# Define Number of Episodes
# n_episodes = 2e+3

# e-greedy parameters to investigate
print('Determine the epsilon parameters for the epsilon-greedy policy...\n')
epsilons = np.array([0.1, 0.13, 0.16])
print('epsilons: '.format(epsilons), '\n')

# various step-sizes (alpha) to try
print('Determine the step-sizes parameters (alphas) for the TD(0)...\n')
step_sizes = np.array([0.4])
print('step_sizes: {}'.format(step_sizes), '\n')

# Fixed discount
discount_fixed = 1

# Create a mesh-grid of trials
print('Create a dictionary of the RL-models of interest...\n')
epsilons, step_sizes = np.meshgrid(epsilons, step_sizes)
epsilons = epsilons.flatten()
step_sizes = step_sizes.flatten()

# Create a dictionary of the RL-trials of interest
RL_trials = {"baseline":
             {'epsilon': 0.1,
              'step_size': 0.5, 'discount': 1}}

for n, trial in enumerate(list(zip(epsilons, step_sizes))):
    key = 'trial_' + str(n+1)
    RL_trials[key] = {'epsilon': trial[0], 
                      'step_size': trial[1], 'discount': discount_fixed}
print('Number of RL-models to try: {}\n'.format(len(RL_trials)))

print('Let all RL-models to be trained for {0:,} episodes...\n'.format(int(n_episodes)))

rewards_per_trial_QL = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())
q_values_per_trial_QL = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())

for trial, params_dict in RL_trials.items():
    
    # Read out parameters from "params_dict"
    epsilon = params_dict['epsilon']
    step_size = params_dict['step_size']
    discount = params_dict['discount']
    
    # Apply SARSA [on-policy TD(0) Control]
    q_values, tot_rewards = TD0.q_learning_off_policy(env, 
                                                      n_episodes=n_episodes, 
                                                      step_size=step_size, discount=discount, epsilon=epsilon)
    
    # Update "rewards_per_trial" and "q_values_per_trial" OrderedDicts
    rewards_per_trial_QL[trial] = tot_rewards
    q_values_per_trial_QL[trial] = q_values

title = 'Efficiency of the RL Method\n[Q-Learning off-policy TD(0) Control]'
plotutls.plot_learning_curve(rewards_per_trial_QL, title=title)

RL_trials

"""
Explanation: 2c. Q-Learning: Off-Policy TD(0) Control
End of explanation
"""


# Define Number of Episodes
# n_episodes = 2e+3

# e-greedy parameters to investigate
print('Determine the epsilon parameters for the epsilon-greedy policy...\n')
epsilons = np.array([0.1, 0.13, 0.16])
print('epsilons: '.format(epsilons), '\n')

# various step-sizes (alpha) to try
print('Determine the step-sizes parameters (alphas) for the TD(0)...\n')
step_sizes = np.array([0.4])
print('step_sizes: {}'.format(step_sizes), '\n')

# Fixed discount
discount_fixed = 1

# Create a mesh-grid of trials
print('Create a dictionary of the RL-models of interest...\n')
epsilons, step_sizes = np.meshgrid(epsilons, step_sizes)
epsilons = epsilons.flatten()
step_sizes = step_sizes.flatten()

# Create a dictionary of the RL-trials of interest
RL_trials = {"baseline":
             {'epsilon': 0.1,
              'step_size': 0.5, 'discount': 1}}

for n, trial in enumerate(list(zip(epsilons, step_sizes))):
    key = 'trial_' + str(n+1)
    RL_trials[key] = {'epsilon': trial[0], 
                      'step_size': trial[1], 'discount': discount_fixed}
print('Number of RL-models to try: {}\n'.format(len(RL_trials)))

print('Let all RL-models to be trained for {0:,} episodes...\n'.format(int(n_episodes)))

rewards_per_trial_ExpSARSA = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())
q_values_per_trial_ExpSARSA = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())

for trial, params_dict in RL_trials.items():
    
    # Read out parameters from "params_dict"
    epsilon = params_dict['epsilon']
    step_size = params_dict['step_size']
    discount = params_dict['discount']
    
    # Apply SARSA [on-policy TD(0) Control]
    q_values, tot_rewards = TD0.expected_sarsa_on_policy(env, 
                                                         n_episodes=n_episodes, 
                                                         step_size=step_size, discount=discount, epsilon=epsilon)
    
    # Update "rewards_per_trial" and "q_values_per_trial" OrderedDicts
    rewards_per_trial_ExpSARSA[trial] = tot_rewards
    q_values_per_trial_ExpSARSA[trial] = q_values

title = 'Efficiency of the RL Method\n[Expected SARSA on-policy TD(0) Control]'
plotutls.plot_learning_curve(rewards_per_trial_ExpSARSA, title=title)

RL_trials

"""
Explanation: 2d. On-Policy Expected SARSA
End of explanation
"""


# Define Number of Episodes
# n_episodes = 2e+3

# e-greedy parameters to investigate
print('Determine the epsilon parameters for the epsilon-greedy policy...\n')
epsilons = np.array([0.1, 0.13, 0.16])
print('epsilons: '.format(epsilons), '\n')

# various step-sizes (alpha) to try
print('Determine the step-sizes parameters (alphas) for the TD(0)...\n')
step_sizes = np.array([0.4])
print('step_sizes: {}'.format(step_sizes), '\n')

# Fixed discount
discount_fixed = 1

# Create a mesh-grid of trials
print('Create a dictionary of the RL-models of interest...\n')
epsilons, step_sizes = np.meshgrid(epsilons, step_sizes)
epsilons = epsilons.flatten()
step_sizes = step_sizes.flatten()

# Create a dictionary of the RL-trials of interest
RL_trials = {"baseline":
             {'epsilon': 0.1,
              'step_size': 0.5, 'discount': 1}}

for n, trial in enumerate(list(zip(epsilons, step_sizes))):
    key = 'trial_' + str(n+1)
    RL_trials[key] = {'epsilon': trial[0], 
                      'step_size': trial[1], 'discount': discount_fixed}
print('Number of RL-models to try: {}\n'.format(len(RL_trials)))

print('Let all RL-models to be trained for {0:,} episodes...\n'.format(int(n_episodes)))

rewards_per_trial_DSARSA = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())
q_values_per_trial_DSARSA = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())

for trial, params_dict in RL_trials.items():
    
    # Read out parameters from "params_dict"
    epsilon = params_dict['epsilon']
    step_size = params_dict['step_size']
    discount = params_dict['discount']
    
    # Apply SARSA [on-policy TD(0) Control]
    q_values_1, tot_rewards = TD0.sarsa_on_policy_control(env, 
                                                        n_episodes=n_episodes, 
                                                        step_size=step_size, discount=discount, epsilon=epsilon,
                                                        double_learning=True)
    
    # Update "rewards_per_trial" and "q_values_per_trial" OrderedDicts
    rewards_per_trial_DSARSA[trial] = tot_rewards
    q_values_per_trial_DSARSA[trial] = q_values_1

title = 'Efficiency of the RL Method\n[Double SARSA: on-policy TD(0) Control]'
plotutls.plot_learning_curve(rewards_per_trial_DSARSA, title=title)

RL_trials

"""
Explanation: 3. Double Learning: a method to mitigate maximization bias
3a. Double SARSA: On-Policy TD(0) Control
End of explanation
"""


# Define Number of Episodes
# n_episodes = 2e+3

# e-greedy parameters to investigate
print('Determine the epsilon parameters for the epsilon-greedy policy...\n')
epsilons = np.array([0.1, 0.13, 0.16])
print('epsilons: '.format(epsilons), '\n')

# various step-sizes (alpha) to try
print('Determine the step-sizes parameters (alphas) for the TD(0)...\n')
step_sizes = np.array([0.4])
print('step_sizes: {}'.format(step_sizes), '\n')

# Fixed discount
discount_fixed = 1

# Create a mesh-grid of trials
print('Create a dictionary of the RL-models of interest...\n')
epsilons, step_sizes = np.meshgrid(epsilons, step_sizes)
epsilons = epsilons.flatten()
step_sizes = step_sizes.flatten()

# Create a dictionary of the RL-trials of interest
RL_trials = {"baseline":
             {'epsilon': 0.1,
              'step_size': 0.5, 'discount': 1}}

for n, trial in enumerate(list(zip(epsilons, step_sizes))):
    key = 'trial_' + str(n+1)
    RL_trials[key] = {'epsilon': trial[0], 
                      'step_size': trial[1], 'discount': discount_fixed}
print('Number of RL-models to try: {}\n'.format(len(RL_trials)))

print('Let all RL-models to be trained for {0:,} episodes...\n'.format(int(n_episodes)))

rewards_per_trial_DQL = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())
q_values_per_trial_DQL = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())

for trial, params_dict in RL_trials.items():
    
    # Read out parameters from "params_dict"
    epsilon = params_dict['epsilon']
    step_size = params_dict['step_size']
    discount = params_dict['discount']
    
    # Apply SARSA [on-policy TD(0) Control]
    q_values, tot_rewards = TD0.q_learning_off_policy(env, 
                                                      n_episodes=n_episodes, 
                                                      step_size=step_size, discount=discount, epsilon=epsilon,
                                                      double_learning=True)
    
    # Update "rewards_per_trial" and "q_values_per_trial" OrderedDicts
    rewards_per_trial_DQL[trial] = tot_rewards
    q_values_per_trial_DQL[trial] = q_values

title = 'Efficiency of the RL Method\n[Double Q-Learning: off-policy TD(0) Control]'
plotutls.plot_learning_curve(rewards_per_trial_DQL, title=title, lower_reward_ratio=-100)

RL_trials

"""
Explanation: 3b. Double Q-Learning: Off-Policy TD(0) Control
End of explanation
"""


# Define Number of Episodes
# n_episodes = 2e+3

# e-greedy parameters to investigate
print('Determine the epsilon parameters for the epsilon-greedy policy...\n')
epsilons = np.array([0.1, 0.13, 0.16])
print('epsilons: '.format(epsilons), '\n')

# various step-sizes (alpha) to try
print('Determine the step-sizes parameters (alphas) for the TD(0)...\n')
step_sizes = np.array([0.4])
print('step_sizes: {}'.format(step_sizes), '\n')

# Fixed discount
discount_fixed = 1

# Create a mesh-grid of trials
print('Create a dictionary of the RL-models of interest...\n')
epsilons, step_sizes = np.meshgrid(epsilons, step_sizes)
epsilons = epsilons.flatten()
step_sizes = step_sizes.flatten()

# Create a dictionary of the RL-trials of interest
RL_trials = {"baseline":
             {'epsilon': 0.1,
              'step_size': 0.5, 'discount': 1}}

for n, trial in enumerate(list(zip(epsilons, step_sizes))):
    key = 'trial_' + str(n+1)
    RL_trials[key] = {'epsilon': trial[0], 
                      'step_size': trial[1], 'discount': discount_fixed}
print('Number of RL-models to try: {}\n'.format(len(RL_trials)))

print('Let all RL-models to be trained for {0:,} episodes...\n'.format(int(n_episodes)))

rewards_per_trial_DExpSARSA = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())
q_values_per_trial_DExpSARSA = OrderedDict((label, np.array([])) for label, _ in RL_trials.items())

for trial, params_dict in RL_trials.items():
    
    # Read out parameters from "params_dict"
    epsilon = params_dict['epsilon']
    step_size = params_dict['step_size']
    discount = params_dict['discount']
    
    # Apply SARSA [on-policy TD(0) Control]
    q_values, tot_rewards = TD0.expected_sarsa_on_policy(env, 
                                                         n_episodes=n_episodes, 
                                                         step_size=step_size, discount=discount, epsilon=epsilon, 
                                                         double_learning=True)
    
    # Update "rewards_per_trial" and "q_values_per_trial" OrderedDicts
    rewards_per_trial_DExpSARSA[trial] = tot_rewards
    q_values_per_trial_DExpSARSA[trial] = q_values

title = 'Efficiency of the RL Method\n[Double Expected SARSA: on-policy TD(0) Control]'
plotutls.plot_learning_curve(rewards_per_trial_DQL, title=title, lower_reward_ratio=-50)

RL_trials

"""
Explanation: 3c. Double Expected SARSA: On-Policy TD(0) Control
End of explanation
"""


winning_trial = 'baseline'
rewards_per_trial_best_models = OrderedDict([('Model_SARSA', np.array([])),
                                             ('Model_DSARSA', np.array([])),
                                             ('Model_QL', np.array([])),
                                             ('Model_DQL', np.array([])),
                                             ('Model_ExpSARSA', np.array([])),
                                             ('Model_DExpSARSA', np.array([]))])

rewards_per_trial_best_models['Model_SARSA'] = rewards_per_trial_SARSA[winning_trial]
rewards_per_trial_best_models['Model_DSARSA'] = rewards_per_trial_DSARSA[winning_trial]
rewards_per_trial_best_models['Model_QL'] = rewards_per_trial_QL[winning_trial]
rewards_per_trial_best_models['Model_DQL'] = rewards_per_trial_DQL[winning_trial]
rewards_per_trial_best_models['Model_ExpSARSA'] = rewards_per_trial_ExpSARSA[winning_trial]
rewards_per_trial_best_models['Model_DExpSARSA'] = rewards_per_trial_DExpSARSA[winning_trial] 

title = 'Efficiency of the RL Method\n[SARSA vs Q-Learning and Expected SARSA Winning Models]'
plotutls.plot_learning_curve(rewards_per_trial_best_models, title=title, lower_reward_ratio=-100)

title = 'Efficiency of the RL Method\n[SARSA vs Q-Learning and Expected SARSA Winning Models]'
plotutls.plot_learning_curve(rewards_per_trial_best_models, title=title, lower_reward_ratio=-35)

"""
Explanation: 4. Comparison of SARSA, Q-Learning and Expected SARSA best models
After an initial transient, Q-learning learns values for the optimal policy, the ones that travels right along the edge of the cliff. Unfortunately, this results the traveller fall-off the cliff ocasionally, because of the $\varepsilon$-greedy selection. SARSA, on the other hand, takes the action selection into account and learns the longer but safer path, through the upper part of the grid. Although, Q-learning actually learns the values of the optimal policy, its online performance is worse than that of SARSA, which learns the roundabout policy.
End of explanation
"""


# print optimal policy
def print_optimal_policy(q_values, grid_height=4, grid_width=12):
    # Define a helper dictionary of actions
    actions_dict = {}
    actions = ['UP', 'RIGHT', 'DOWN', 'LEFT']
    for k, v in zip(actions, range(0, len(actions))):
        actions_dict[k] = v
        
    # Define the position of target dstination
    GOAL = [3, 11]
        
    # Reshape the "q_values" table to follow grid-world dimensionality
    q_values = q_values.reshape((grid_height, grid_width, len(actions)))
        
    optimal_policy = []
    for i in range(0, grid_height):
        optimal_policy.append([])
        for j in range(0, grid_width):
            if [i, j] == GOAL:
                optimal_policy[-1].append('G')
                continue
            bestAction = np.argmax(q_values[i, j, :])
            if bestAction == actions_dict['UP']:
                optimal_policy[-1].append('\U00002191')
            elif bestAction == actions_dict['RIGHT']:
                optimal_policy[-1].append('\U00002192')
            elif bestAction == actions_dict['DOWN']:
                optimal_policy[-1].append('\U00002193')
            elif bestAction == actions_dict['LEFT']:
                optimal_policy[-1].append('\U00002190')
    for row in optimal_policy:
        print(*row)


"""
Explanation: Note: In "CliffWalking-v0" environment the traveler can choose one of the below actions as she navigates through the grid:
- "UP":  denoted by 0
- "RIGHT": denoted by 1
- "DOWN": denoted by 2
- "LEFT": denoted by 3
End of explanation
"""


winning_trial = 'baseline'

print_optimal_policy(q_values_per_trial_SARSA[winning_trial], grid_height=4, grid_width=12)

"""
Explanation: 5. Learned Policies
SARSA on-Policy TD(0) Control:
Winning trial:
End of explanation
"""


print_optimal_policy(q_values_per_trial_DSARSA[winning_trial], grid_height=4, grid_width=12)

"""
Explanation: Double SARSA on-Policy TD(0) Control:
End of explanation
"""


print_optimal_policy(q_values_per_trial_QL[winning_trial], grid_height=4, grid_width=12)

"""
Explanation: Q-Learning off-Policy TD(0) Control:
End of explanation
"""


print_optimal_policy(q_values_per_trial_DQL[winning_trial], grid_height=4, grid_width=12)

"""
Explanation: Double Q-Learning off-Policy TD(0) Control:
End of explanation
"""


print_optimal_policy(q_values_per_trial_ExpSARSA[winning_trial], grid_height=4, grid_width=12)

"""
Explanation: Expected SARSA on-Policy TD(0) Control:
End of explanation
"""


print_optimal_policy(q_values_per_trial_DExpSARSA[winning_trial], grid_height=4, grid_width=12)

"""
Explanation: Double Expected SARSA on-Policy TD(0) Control:
End of explanation
"""


