import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import tempfile
import base64
import pprint
import random
import json
import sys
import gym
import io

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

from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import Adam

"""
Explanation: Function Approximation
Previous we tried using manually created buckets to discretize continuous states, effectively mapping continuous observations to discrete states. Great, but, isn't that what Supervised Learning is for? How about we use the power of function approximators to do that mapping for us???
Let's do just that!
End of explanation
"""


def action_selection(state, model, episode, n_episodes):
    epsilon = 0.99 if episode < n_episodes//4 else 0.33 if episode < n_episodes//2 else 0.
    values = model.predict(state.reshape(1, state.shape[0]))[0]
    if np.random.random() < epsilon:
        action = np.random.randint(len(values))
    else:
        action = np.argmax(values)
    return action, epsilon

def neuro_q_learning(env, gamma = 0.99):
    nS = env.observation_space.shape[0]
    nA = env.env.action_space.n
    
    # memory bank
    memory_bank = deque()
    memory_bank_size = 100000
    
    # function approximator
    model = Sequential()
    model.add(Dense(64, input_dim=nS, activation='relu'))
    model.add(Dense(nA, activation='linear'))
    model.compile(loss='mse', optimizer='adam')

    # constant values
    n_episodes = 50000
    batch_size = 256
    training_frequency = 20
    
    # for statistics
    epsilons = []
    states = []
    actions = []
    
    # interactions
    for episode in range(n_episodes):
        state = env.reset()
        done = False
        
        # each episode
        while not done:
            states.append(state)
            
            # select action
            action, epsilon = action_selection(state, model, episode, n_episodes)
            epsilons.append(epsilon)
            actions.append(action)
            
            # save history in memory bank
            nstate, reward, done, info = env.step(action)
            memory_bank.append((state, action, reward, nstate, done))
            if len(memory_bank) > memory_bank_size:
                memory_bank.popleft()
            
            # iterate to next state
            state = nstate

        # only every few episodes enter training and update neural network weights
        if episode % training_frequency == 0 and len(memory_bank) == memory_bank_size:
            
            # randomly select batches of samples from the history
            # for training to prevent values spiking due to high 
            # correlation of sequential values
            minibatch = np.array(random.sample(memory_bank, batch_size))

            # extract values by type from the minibatch
            state_batch = np.array(minibatch[:,0].tolist())
            action_batch = np.array(minibatch[:,1].tolist())
            rewards_batch = np.array(minibatch[:,2].tolist())
            state_prime_batch = np.array(minibatch[:,3].tolist())
            is_terminal_batch = np.array(minibatch[:,4].tolist())

            # use the current neural network to predict 
            # current state values and next state values
            state_value_batch = model.predict(state_batch)
            next_state_value_batch = model.predict(state_prime_batch)

            # update the state values given the batch
            for i in range(len(minibatch)):
                if is_terminal_batch[i]:
                    state_value_batch[i, action_batch[i]] = rewards_batch[i]
                else:
                    state_value_batch[i, action_batch[i]] = rewards_batch[i] + gamma * np.max(next_state_value_batch[i])
            
            # update the neural network weights
            model.train_on_batch(state_batch, state_value_batch)

    return model, (epsilons, states, actions)

"""
Explanation: As with previous notebooks we will use an action selection function to give us an exploration strategy. Feel free to change this function and play with it to see what happens. For now, let me just provide a couple function.
End of explanation
"""


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

model, stats = neuro_q_learning(env)

"""
Explanation: Did you see that neuro_q_learning function??? I had to add comments to this one to make it easier to understand. There are a couple of things that you might want to read about. Check out the "Further reading" sections on the README.md files for that. 
For example, did you see we are now saving experiences into a memory_bank? That is called experience replay. Now, we don't just go over the saved tuples sequentially, we randomly select a batch of sample and update the neural network weights with that. This is important often due to the high correlation of the samples.
Also, did you see how we update the state_value_batch with each sample? That is just very cool, we are not using alpha here anymore, the learning rate is now used for gradient descent when we are learning the network weights.
Ok, let's run this algorithm!
Beware, this will take some time... 15-25 min, so go get your favorite beverage.
End of explanation
"""


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

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

strm = ''
for video_path, meta_path in videos:
    video = io.open(video_path, 'r+b').read()
    encoded = base64.b64encode(video)
    
    with open(meta_path) as data_file:    
        meta = json.load(data_file)

    html_tag = """
    <h2>{0}<h2/>
    <video width="960" height="540" controls>
        <source src="data:video/mp4;base64,{1}" type="video/mp4" />
    </video>"""
    strm += html_tag.format('Episode ' + str(meta['episode_id']), encoded.decode('ascii'))
HTML(data=strm)

"""
Explanation: Hope you enjoyed your vacations. :)
Back to it. Let's look at a couple of episodes to see how the agent did.
End of explanation
"""


env.close()

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

"""
Explanation: Let's close up the environment and upload (in case you want to).
End of explanation
"""


epsilons, states, actions = stats

plt.plot(np.arange(len(epsilons)), epsilons, '.')

hist, bins = np.histogram(actions, bins=3)
width = 0.7 * (bins[1] - bins[0])
center = (bins[:-1] + bins[1:]) / 2
plt.bar(center, hist, align='center', width=width)
plt.show()

"""
Explanation: Now, let's look at some of the statistics
End of explanation
"""


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

for episode in range(100):
    state = env.reset()
    done = False
    while not done:
        action = np.argmax(model.predict(state.reshape(1, 4))[0])
        nstate, reward, done, info = env.step(action)
        state = nstate

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

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

strm = ''
for video_path, meta_path in videos:
    video = io.open(video_path, 'r+b').read()
    encoded = base64.b64encode(video)
    
    with open(meta_path) as data_file:    
        meta = json.load(data_file)

    html_tag = """
    <h2>{0}<h2/>
    <video width="960" height="540" controls>
        <source src="data:video/mp4;base64,{1}" type="video/mp4" />
    </video>"""
    strm += html_tag.format('Episode ' + str(meta['episode_id']), encoded.decode('ascii'))
HTML(data=strm)

env.close()

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

"""
Explanation: Cool, not bad at all!
Let's test this fully trained agent and see how it does.
End of explanation
"""


def action_selection(state, model, episode, n_episodes):
    epsilon = 0.99 if episode < n_episodes//4 else 0.33 if episode < n_episodes//2 else 0.
    values = model.predict(state.reshape(1, 4))[0]
    if np.random.random() < epsilon:
        action = np.random.randint(len(values))
    else:
        action = np.argmax(values)
    return action, epsilon

def neuro_q_learning(env, gamma = 0.99):
    nS = env.observation_space.shape[0]
    nA = env.env.action_space.n
    
    # memory bank
    memory_bank = deque()
    memory_bank_size = 100000
    
    # function approximator
    #### THIS IS A WEAK NEURAL NETWORK
    #### CAN YOU HELP THE AGENT LEARN MORE
    #### COMPLEX OBSERVATIONS?????
    model = Sequential()
    model.add(Dense(64, input_dim=nS, activation='relu'))
    model.add(Dense(nA, activation='linear'))
    model.compile(loss='mse', optimizer='adam')

    # constant values
    n_episodes = 50000
    batch_size = 256
    training_frequency = 20
    
    # for statistics
    epsilons = []
    states = []
    actions = []
    
    # interactions
    for episode in range(n_episodes):
        state = env.reset()
        done = False
        
        # each episode
        while not done:
            states.append(state)
            
            # select action
            action, epsilon = action_selection(state, model, episode, n_episodes)
            epsilons.append(epsilon)
            actions.append(action)
            
            # save history in memory bank
            nstate, reward, done, info = env.step(action)
            memory_bank.append((state, action, reward, nstate, done))
            if len(memory_bank) > memory_bank_size:
                memory_bank.popleft()
            
            # iterate to next state
            state = nstate

        # only every few episodes enter training and update neural network weights
        if episode % training_frequency == 0 and len(memory_bank) == memory_bank_size:
            
            # randomly select batches of samples from the history
            # for training to prevent values spiking due to high 
            # correlation of sequential values
            minibatch = np.array(random.sample(memory_bank, batch_size))

            # extract values by type from the minibatch
            state_batch = np.array(minibatch[:,0].tolist())
            action_batch = np.array(minibatch[:,1].tolist())
            rewards_batch = np.array(minibatch[:,2].tolist())
            state_prime_batch = np.array(minibatch[:,3].tolist())
            is_terminal_batch = np.array(minibatch[:,4].tolist())

            # use the current neural network to predict 
            # current state values and next state values
            state_value_batch = model.predict(state_batch)
            next_state_value_batch = model.predict(state_prime_batch)

            # update the state values given the batch
            for i in range(len(minibatch)):
                if is_terminal_batch[i]:
                    state_value_batch[i, action_batch[i]] = rewards_batch[i]
                else:
                    state_value_batch[i, action_batch[i]] = rewards_batch[i] + gamma * np.max(next_state_value_batch[i])
            
            # update the neural network weights
            model.train_on_batch(state_batch, state_value_batch)

    return model, (epsilons, states, actions)

"""
Explanation: Your turn
Nice, the agent is good, but it could use some help. It seems the agent is learning, but perhaps not enough. Try modifying the neural network... Maybe more layers? Larger hidden layers?? Try things out, it's now your turn.
End of explanation
"""


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

model, stats = neuro_q_learning(env)

"""
Explanation: Let's test thing puppy.
End of explanation
"""


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

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

strm = ''
for video_path, meta_path in videos:
    video = io.open(video_path, 'r+b').read()
    encoded = base64.b64encode(video)
    
    with open(meta_path) as data_file:    
        meta = json.load(data_file)

    html_tag = """
    <h2>{0}<h2/>
    <video width="960" height="540" controls>
        <source src="data:video/mp4;base64,{1}" type="video/mp4" />
    </video>"""
    strm += html_tag.format('Episode ' + str(meta['episode_id']), encoded.decode('ascii'))
HTML(data=strm)

env.close()

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

"""
Explanation: Wake upp!!!
End of explanation
"""


epsilons, states, actions = stats

plt.plot(np.arange(len(epsilons)), epsilons, '.')

hist, bins = np.histogram(actions, bins=3)
width = 0.7 * (bins[1] - bins[0])
center = (bins[:-1] + bins[1:]) / 2
plt.bar(center, hist, align='center', width=width)
plt.show()

"""
Explanation: You happy with it????? Are you sure????
If you have any issues, check out my solution. But don't do befere you spend a good deal of time on this one.
Let's look into the stats.
End of explanation
"""


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

for episode in range(100):
    state = env.reset()
    done = False
    while not done:
        action = np.argmax(model.predict(state.reshape(1, 4))[0])
        nstate, reward, done, info = env.step(action)
        state = nstate

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

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

strm = ''
for video_path, meta_path in videos:
    video = io.open(video_path, 'r+b').read()
    encoded = base64.b64encode(video)
    
    with open(meta_path) as data_file:    
        meta = json.load(data_file)

    html_tag = """
    <h2>{0}<h2/>
    <video width="960" height="540" controls>
        <source src="data:video/mp4;base64,{1}" type="video/mp4" />
    </video>"""
    strm += html_tag.format('Episode ' + str(meta['episode_id']), encoded.decode('ascii'))
HTML(data=strm)

env.close()

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

"""
Explanation: How about running your fully-trained greedy agent?
End of explanation
"""


