# pip install numpy gymnasium pygame tensorflow
import numpy as np
import gymnasium as gym
import tensorflow as tf
import random
from collections import deque
from tensorflow.keras import layers

env = gym.make("FrozenLake-v1")

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

episodes = 5000
gamma = 0.99
epsilon = 1
lr = 0.001
memory = deque(maxlen=2000)
batch_size = 32

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

target = tf.keras.models.clone_model(model)
target.set_weights(model.get_weights())

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

# One-hot encoding
def encode(s):
    x = np.zeros(states)
    x[s] = 1
    return x.reshape(1, -1)

# Choose action
def choose_action(s):
    if np.random.random() < epsilon:
        return env.action_space.sample()
    return np.argmax(model.predict(encode(s), verbose=0)[0])

# Experience Replay
def replay():
    if len(memory) < batch_size:
        return

    batch = random.sample(memory, batch_size)

    for s, a, r, ns, done in batch:
        q = model.predict(encode(s), verbose=0)
        if done:
            q[0][a] = r
        else:
            nq = target.predict(encode(ns), verbose=0)
            q[0][a] = r + gamma*np.max(nq[0])

        model.fit(encode(s), q, epochs=1, verbose=0)

# Training
rewards = []

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

    for i in range(100):
        action = choose_action(state)
        ns, r, done, truncated, _ = env.step(action)

        memory.append((state, action, r, ns, done or truncated))
        replay()

        state = ns
        total += r

        if done or truncated:
            break

    epsilon = max(0.01, epsilon*0.995)
    rewards.append(total)

    if e % 10 == 0:
        target.set_weights(model.get_weights())

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

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

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

for i in range(100):
    action = np.argmax(model.predict(encode(state), verbose=0)[0])
    state, r, done, truncated, _ = env.step(action)

    if done or truncated:
        break

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