# 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()

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

    total = 0

    for i in range(100):
        next_state, reward, done, truncated, _ = env.step(action)

        if np.random.random() < epsilon:
            next_action = env.action_space.sample()
        else:
            next_action = np.argmax(Q[next_state])

        target = reward if done or truncated else reward + gamma*Q[next_state, next_action]

        Q[state, action] += alpha*(target-Q[state, action])

        state = next_state
        action = next_action
        total += reward

        if done or truncated:
            break

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

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

print("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()
