# Modern Keras: The Comprehensive Guide to Deep Learning with the Keras API and Python

# This code is made available as supporting resource for the "Modern Keras" book by recluze. 
# See more details here: https://recluze.net/keras-book 

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

def run(episodes, is_training=True, render=False):

    env = gym.make('FrozenLake-v1', map_name="8x8", \
                    is_slippery=True, \
                    render_mode='human' if render else None)

    if(is_training):
        q = np.zeros((env.observation_space.n, \
                      env.action_space.n))
    else:
        f = open('frozen_lake8x8.pkl', 'rb')
        q = pickle.load(f)
        f.close()

    learning_rate_a = 0.9 
    discount_factor_g = 0.9 

    epsilon = 1         
    epsilon_decay_rate = 0.0001    

    rng = np.random.default_rng()   
    rewards_per_episode = np.zeros(episodes)

    for i in range(episodes):
        state = env.reset()[0] # 0 = top left, 63 = bottom right
        terminated = False      
        truncated = False       

        while(not terminated and not truncated):
            if is_training and rng.random() < epsilon:
                action = env.action_space.sample() 
                # 0 = left, 1 = down , 2 = right, 3 = up
            else:
                action = np.argmax(q[state,:])

            new_state,reward,terminated,truncated,_ = \
                        env.step(action)

            if is_training:
                q[state,action] = q[state,action] +  \
                learning_rate_a * (
                    reward + discount_factor_g * \
                    np.max(q[new_state,:]) - q[state,action]
                )

            state = new_state

        epsilon = max(epsilon - epsilon_decay_rate, 0)

        if(epsilon==0):
            learning_rate_a = 0.0001

        if reward == 1:
            rewards_per_episode[i] = 1

    env.close()

    sum_rewards = np.zeros(episodes)
    for t in range(episodes):
        sum_rewards[t] = np.sum(rewards_per_episode[ \
                                max(0, t-100):(t+1)])
    plt.plot(sum_rewards)
    plt.savefig('frozen_lake8x8.png')

    if is_training:
        f = open("frozen_lake8x8.pkl","wb")
        pickle.dump(q, f)
        f.close()

if __name__ == '__main__':
    # run(1000, is_training=True, render=True)
    # run(2000, is_training=True, render=False)
    run(1, is_training=False, render=True)
