import matplotlib.pyplot as plt

import numpy as np
import pandas as pd
import tempfile
import pprint
import math
import json
import sys
import gym

from gym import wrappers
from subprocess import check_output
from IPython.display import HTML

"""
Explanation: Model-Free Reinforcement Learning
Remember how in last Notebook we felt like cheating by using directions calculated from the map of the environment?? Well, model-free reinforcement learning deals with that. Model-free refers to the fact that algorithms unders this category do not need a model of the environment, also known as MDP, to calculate optimal policies.
In this notebook, we will look at what is perhaps the most popular model-free reinforcement learning algorithm, q-learning. Q-learning run without needing a map of the environment, it works by balancing the need for exploration with the need for exploiting previously explored knowledge. Let's take a look.
End of explanation
"""


def action_selection(state, Q, episode, n_episodes):
    epsilon = max(0, episode/n_episodes*2)
    if np.random.random() < epsilon:
        action = np.random.randint(len(Q[0]))
    else:
        action = np.argmax(Q[state])
    return action, epsilon

Q = [[0]]
n_episodes = 10000
epsilons = []
for episode in range(n_episodes//2, -n_episodes//2, -1):
    _, epsilon = action_selection(0, Q, episode, n_episodes)
    epsilons.append(epsilon)
plt.plot(np.arange(len(epsilons)), epsilons, '.')
plt.ylabel('Probability')
plt.xlabel('Episode')

"""
Explanation: Q-Learning
The function below, action_selection is an important aspect of reinforcement learning algorithms. The fact is, when you have possibly conflicting needs, explore vs exploit, you enter into a difficult situation, dilemma. The Exploration vs Exploitation Dilemma is at the core of reinforcement learning and it is good for you to think about it for a little while. How much do you need to explore an environment before you exploit it?
In the function below we use one of the many alternatives which is we explore a lot at the begining and decay the amount of exploration as we increase the number of episodes. Let's take a look at what the function looks like:
End of explanation
"""


def q_learning(env, alpha = 0.9, gamma = 0.9):
    nS = env.env.observation_space.n
    nA = env.env.action_space.n
    
    Q = np.random.random((nS, nA)) * 2.0
    n_episodes = 10000
    
    for episode in range(n_episodes//2, -n_episodes//2, -1):
        state = env.reset()
        done = False
        while not done:
            action, _ = action_selection(state, Q, episode, n_episodes)
            nstate, reward, done, info = env.step(action)
            Q[state][action] += alpha * (reward + gamma * Q[nstate].max() * (not done) - Q[state][action])
            state = nstate
    return Q

"""
Explanation: See that? So, at episode 0 we have 100% change of acting randomly, all the way down to 0 when we stop exploring and instead always select the action that we think would maximizing the discounted future rewards. 
Again, this is a way of doing this, there are many and you surely should be thinking about better ways of doing so.
Next, let me show you what Q-Learning looks like:
End of explanation
"""


mdir = tempfile.mkdtemp()
env = gym.make('FrozenLake-v0')
env = wrappers.Monitor(env, mdir, force=True)

Q = q_learning(env)

"""
Explanation: Nice, right? You just pass it an environment, nS and nA are the number of states and actions respectively. 
Q is a table of states as rows and actions as columns that will hold the expected reward the agent expects to get for taking action 'a' on state 's'. You can see how we initialize Q(s,a)'s to a random value, but also we multiply that by 2. You may ask, why is this? This is called "Optimism in the face of uncertainty" and it is a common reinforcement learning technique for encouraging agents to explore. Think about it on an intuitive level. If you think positively most of the time, if you receive a low balling job offer, you are going to pass on it and potentially get a better offer later. Worst case, you don't find any better offer and after 'adjusting' your estimates you will think an offer like the "low balling" one you got wasn't that bad after all. The same applies to reinforcement learning agent, cool right?
Then, I go on a loop for n_episodes using the action_selection function as described above. Don't pay too much attention to the range start and end, that is just the way I get the exploration strategy the way I showed. You should not like it, I don't like it. You will have a chance to make it better.
For now, let's unleash this agent and see how it does!!!
End of explanation
"""


videos = np.array(env.videos)
n_videos = 5

idxs = np.linspace(0, len(videos) - 1, n_videos).astype(int)
videos = videos[idxs,:]

urls = []
for i in range(n_videos):
    out = check_output(["asciinema", "upload", videos[i][0]])
    out = out.decode("utf-8").replace('\n', '').replace('\r', '')
    urls.append([out])
videos = np.concatenate((videos, urls), axis=1)

strm = ''
for video_path, meta_path, url in videos:

    with open(meta_path) as data_file:    
        meta = json.load(data_file)
    castid = url.split('/')[-1]
    html_tag = """
    <h2>{0}
    <script type="text/javascript" 
        src="https://asciinema.org/a/{1}.js" 
        id="asciicast-{1}" 
        async data-autoplay="true" data-size="big">
    </script>
    """
    strm += html_tag.format('Episode ' + str(meta['episode_id']),
                               castid)
HTML(data=strm)

"""
Explanation: Let's look at a couple of the episodes in more detail.
End of explanation
"""


V = np.max(Q, axis=1)
V

pi = np.argmax(Q, axis=1)
pi

"""
Explanation: Nice!!!
You can see the progress of this agent. From total caos completely sinking into holes, to sliding into the goal fairly consistently.
Let's inspect the Values and Policies.
End of explanation
"""


env.close()

gym.upload(mdir, api_key='<YOUR OPENAI API KEY>')

"""
Explanation: Fair enough, let's close this environment and you will have a chance to submit to your OpenAI account. After that, you will have a chance to modify the action_selection to try something different.
End of explanation
"""


def action_selection(state, Q, episode, n_episodes, decay=0.0006, initial=1.00):
    epsilon = initial * math.exp(-decay*episode)
    if np.random.random() < epsilon:
        action = np.random.randint(len(Q[0]))
    else:
        action = np.argmax(Q[state])
    return action, epsilon

"""
Explanation: Your turn
Maybe you want to try an exponential decay?? (http://www.miniwebtool.com/exponential-decay-calculator/)
P(t) = P0e-rt
where: 
* P(t) = the amount of some quantity at time t 
* P0 = initial amount at time t = 0 
* r = the decay rate 
* t = time (number of periods)
End of explanation
"""


Q = [[0]]
n_episodes = 10000
epsilons = []
for episode in range(n_episodes):
    _, epsilon = action_selection(0, Q, episode, n_episodes)
    epsilons.append(epsilon)
plt.plot(np.arange(len(epsilons)), epsilons, '.')
plt.ylabel('Probability')
plt.xlabel('Episode')

"""
Explanation: Use the following code to test your new exploration strategy:
End of explanation
"""


def q_learning(env, alpha = 0.9, gamma = 0.9):
    nS = env.env.observation_space.n
    nA = env.env.action_space.n
    
    Q = np.random.random((nS, nA)) * 2.0
    n_episodes = 10000
    
    for episode in range(n_episodes):
        state = env.reset()
        done = False
        while not done:
            action, _ = action_selection(state, Q, episode, n_episodes)
            nstate, reward, done, info = env.step(action)
            Q[state][action] += alpha * (reward + gamma * Q[nstate].max() * (not done) - Q[state][action])
            state = nstate
    return Q

mdir = tempfile.mkdtemp()
env = gym.make('FrozenLake-v0')
env = wrappers.Monitor(env, mdir, force=True)

Q = q_learning(env)

"""
Explanation: Let's redefine the q_learning function we had above and run it against the environment again.
End of explanation
"""


videos = np.array(env.videos)
n_videos = 5

idxs = np.linspace(0, len(videos) - 1, n_videos).astype(int)
videos = videos[idxs,:]

urls = []
for i in range(n_videos):
    out = check_output(["asciinema", "upload", videos[i][0]])
    out = out.decode("utf-8").replace('\n', '').replace('\r', '')
    urls.append([out])
videos = np.concatenate((videos, urls), axis=1)

strm = ''
for video_path, meta_path, url in videos:

    with open(meta_path) as data_file:    
        meta = json.load(data_file)
    castid = url.split('/')[-1]
    html_tag = """
    <h2>{0}
    <script type="text/javascript" 
        src="https://asciinema.org/a/{1}.js" 
        id="asciicast-{1}" 
        async data-autoplay="true" data-size="big">
    </script>
    """
    strm += html_tag.format('Episode ' + str(meta['episode_id']),
                               castid)
HTML(data=strm)

"""
Explanation: Curious to see how the new agent did?? Let's check it out!
End of explanation
"""


V = np.max(Q, axis=1)
V

pi = np.argmax(Q, axis=1)
pi

"""
Explanation: Did it do good??? This isn't an easy thing, take your time. Be sure to look into the Notebook solution if you want an idea.
For now, let's take a look at the value function and policy the agent came up with.
End of explanation
"""


env.close()

gym.upload(mdir, api_key='<YOUR OPENAI API KEY>')

"""
Explanation: Good??? Nice!
Let's wrap-up!
End of explanation
"""


