#!!!credit:
# >> gridworld taken from https://github.com/dennybritz <<
!wget https://jbdel.github.io/teaching/nn/ressources/rl/lib/envs/gridworld.py

import numpy as np
import pprint
import sys
from gridworld import GridworldEnv


pp = pprint.PrettyPrinter(indent=2)
env = GridworldEnv()
env._render()

# P[s][a] = (prob, next_state, reward, is_done)
# a = 0 -> up
# a = 1 -> right
# a = 2 -> down
# a = 3 -> left

print(env.P[0][2])
print(env.P[13][0])

def value_iteration(env, theta=0.0001, discount_factor=0.99):
    """
      env.P[s][a] la liste des transition (prob, next_state, reward, done)
      env.nS nombre d'états dans l'environnement
      env.nA nombre d'actions dans l'environnement
      theta: facteur d'arrêt
      discount_factor: le facteur de pondération de la récompense
        
    Retourne:
        Un tuple (policy, V) la politique optimale pi* et la fonction de 
        valeur optimale v*
    """

    def chercher_s_prime(state, V):
        A = np.zeros(env.nA)
        #appliquer ici r(s,a) + sum (discount_factor * prob * V(s'))
        for a in range(env.nA):
            for prob, next_state, reward, done in env.P[state][a]:
                A[a] += ######## A COMPLETER ###########
        return A

    V = np.zeros(env.nS)
    while True:
        # Condition d'arrêt
        delta = 0
        # MAJ de la fonction de valeur pour chaque etat s
        for s in range(env.nS):
            # on va chercher la valeur de s_prime pour chaque action 
            A = chercher_s_prime(s, V)
            best_action_value = np.max(A)
            # est-ce que une MAJ d'état est au dessus du delta
            delta = max(delta, np.abs(best_action_value - V[s]))
            # MAJ de la fonction de valeur de l'état courant 
            V[s] = ######## A COMPLETER ###########        
        # condition d'arrêt 
        if delta < theta:
            break

    # Extraction de la police optimale
    policy = np.zeros([env.nS, env.nA])
    for s in range(env.nS):
        # on regarde les s_prime pour trouver quelle action est la meilleure        
        A = chercher_s_prime(s, V)
        best_action = np.argmax(A)
        policy[s, best_action] = 1.0

    return policy, V

policy, v = value_iteration(env)

print("Reshaped Grid Policy (0=up, 1=right, 2=down, 3=left):")
print(np.reshape(np.argmax(policy, axis=1), env.shape))
print("")

print("Value Function:")
print(v)
print("")

print("Reshaped Grid Value Function:")
print(v.reshape(env.shape))
print("")


# Test the value function
expected_v = np.array([ 0, -1, -2, -3, -1, -2, -3, -2, -2, -3, -2, -1, -3, -2, -1,  0])
np.testing.assert_array_almost_equal(v, expected_v, decimal=1)

!wget https://jbdel.github.io/teaching/nn/ressources/rl/lib/envs/cliff_walking.py
!wget https://jbdel.github.io/teaching/nn/ressources/rl/lib/plotting.py

%matplotlib inline

import gym
import itertools
import matplotlib
import numpy as np
import sys

from collections import defaultdict
from cliff_walking import CliffWalkingEnv
import plotting

matplotlib.style.use('ggplot')

env = CliffWalkingEnv()
env.render()

#0=up, 1=right, 2=down, 3=left
print(env.step(0))
env.render()

print(env.step(1))
env.render()

print(env.step(1))
env.render()

print(env.step(2))
env.render()

def make_epsilon_greedy_policy(Q, epsilon, nA):
    """
    Creates an epsilon-greedy policy based on a given Q-function and epsilon.
    
    Args:
        Q: A dictionary that maps from state -> action-values.
            Each value is a numpy array of length nA (see below)
        epsilon: The probability to select a random action. Float between 0 and 1.
        nA: Number of actions in the environment.
    
    Returns:
        A function that takes the observation as an argument and returns
        the probabilities for each action in the form of a numpy array of length nA.
    
    """
    def policy_fn(observation):
        A = np.ones(nA, dtype=float) * epsilon / nA
        best_action = np.argmax(Q[observation])
        A[best_action] += (1.0 - epsilon)
        return A
    return policy_fn

def q_learning(env, num_episodes, discount_factor=1.0, alpha=0.5, epsilon=0.1):
    """
    Q-Learning algorithm: Off-policy TD control. Finds the optimal greedy policy
    while following an epsilon-greedy policy
    
    Args:
        env: OpenAI environment.
        num_episodes: Number of episodes to run for.
        discount_factor: Gamma discount factor.
        alpha: TD learning rate.
        epsilon: Chance to sample a random action. Float between 0 and 1.
    
    Returns:
        A tuple (Q, episode_lengths).
        Q is the optimal action-value function, a dictionary mapping state -> action values.
        stats is an EpisodeStats object with two numpy arrays for episode_lengths and episode_rewards.
    """

    # The final action-value function.
    # A nested dictionary that maps state -> (action -> action-value).
    Q = defaultdict(lambda: np.zeros(env.action_space.n))

    # Keeps track of useful statistics
    stats = plotting.EpisodeStats(
        episode_lengths=np.zeros(num_episodes),
        episode_rewards=np.zeros(num_episodes))

    # The policy we're following
    policy = make_epsilon_greedy_policy(Q, epsilon, env.action_space.n)

    for i_episode in range(num_episodes):
        # Print out which episode we're on, useful for debugging.
        if (i_episode + 1) % 100 == 0:
            print("\rEpisode {}/{}.".format(i_episode + 1, num_episodes), end="")
            sys.stdout.flush()

        # Reset the environment and pick the first action
        state = env.reset()

        # One step in the environment
        # total_reward = 0.0
        for t in itertools.count():

            # Take a step
            action_probs = policy(state)
            action = np.random.choice(np.arange(len(action_probs)), p=action_probs)
            next_state, reward, done, _ = env.step(action)

            # Update statistics
            stats.episode_rewards[i_episode] += reward
            stats.episode_lengths[i_episode] = t

            # TD Update
            best_next_action = # à compléter
            td_target = # à compléter
            td_delta = td_target - Q[state][action]
            Q[state][action] += # à compléter

            if done:
                break

            state = next_state

    return Q, stats

Q, stats = q_learning(env, 500)

policy = [np.argmax(Q[s]) for s in range(env.nS)]
print("Reshaped Grid Policy (0=up, 1=right, 2=down, 3=left):")
print(np.reshape(policy, env.shape))
print("")

plotting.plot_episode_stats(stats)

env = CliffWalkingEnv(wind=True)
env.reset()

env.step(0)
env.step(1)
env.render()
#aller en haut
print("UP",env.P[env.s][0])
#aller a droite
print("RIGHT",env.P[env.s][1])

Q, stats = q_learning(env, 2000)

policy = [np.argmax(Q[s]) for s in range(env.nS)]
print("Reshaped Grid Policy (0=up, 1=right, 2=down, 3=left):")
print(np.reshape(policy, env.shape))
print("")