# pip install gymnasium tensorflow numpy pygame
import numpy as np
import gymnasium as gym
import tensorflow as tf

env = gym.make("FrozenLake-v1", is_slippery=False)

states = int(env.observation_space.n)
actions = int(env.action_space.n)

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(states,)),
    tf.keras.layers.Dense(actions, activation="linear")
])

model.compile(optimizer=tf.keras.optimizers.Adam(0.01), loss="mse")

episodes = 1000
epsilon = 1.0
gamma = 0.99
rewards = []

for e in range(episodes):
    state, _ = env.reset()
    total = 0

    for _ in range(100):
        s = tf.one_hot([state], states)

        if np.random.random() < epsilon:
            action = env.action_space.sample()
        else:
            action = int(np.argmax(model(s, training=False).numpy()[0]))

        ns, r, done, truncated, _ = env.step(action)

        q = model(s, training=False).numpy()

        if done or truncated:
            target = r
        else:
            ns_tensor = tf.one_hot([ns], states)
            target = r + gamma*np.max(
                model(ns_tensor, training=False).numpy()[0])

        q[0][action] = target
        model.train_on_batch(s, q)

        state = ns
        total += r

        if done or truncated:
            break

    epsilon = max(0.01, np.exp(-0.005*e))
    rewards.append(total)

    if e % 100 == 0:
        print("Episode:", e, "Mean Reward:",
              round(np.mean(rewards[-100:]), 3))

print("Overall Reward:", np.mean(rewards))
env.close()

# Display trained agent
env = gym.make("FrozenLake-v1", is_slippery=False, render_mode="human")
state, _ = env.reset()

for _ in range(100):
    s = tf.one_hot([state], states)
    action = int(np.argmax(model(s, training=False).numpy()[0]))

    state, r, done, truncated, _ = env.step(action)

    if done or truncated:
        break

print("Goal Reached!" if r == 1 else "Agent Failed!")
env.close()