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

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

Q = np.zeros((env.observation_space.n, env.action_space.n))

episodes = 10000
alpha = 0.1
gamma = 0.99
epsilon = 1
rewards = []

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

    for i in range(100):
        if np.random.random() < epsilon:
            action = env.action_space.sample()
        else:
            action = np.argmax(Q[state])

        next_state, reward, done, truncated, _ = env.step(action)

        target = reward if done or truncated else reward + gamma*np.max(Q[next_state])
        Q[state, action] += alpha*(target-Q[state, action])

        state = next_state
        total += reward

        if done or truncated:
            break

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

# Results
print("Q Table:")
print(Q)

print("\nMean Rewards:")
for i in range(10):
    print(i*1000, "-", (i+1)*1000, ":",
          np.mean(rewards[i*1000:(i+1)*1000]))

print("Overall Mean 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(Q[state])
    state, reward, done, truncated, _ = env.step(action)

    if done or truncated:
        break

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