if __name__ == "__main__":
    env = GridWorld(width=6, height=4)

    Q_qlearn, returns_qlearn = q_learning(env, n_episodes=500)
    Q_sarsa, returns_sarsa = sarsa(env, n_episodes=500)

    def average_last(returns, k=50):
        return np.mean(returns[-k:])

    ql_avg = average_last(returns_qlearn)
    sarsa_avg = average_last(returns_sarsa)
    print(f"Q-learning avg return (last 50 eps): {ql_avg:.1f}")
    print(f"SARSA avg return (last 50 eps):      {sarsa_avg:.1f}")

    def extract_greedy_path(Q, env):
        state = env.reset()
        path = [state]
        for _ in range(50):
            a = int(np.argmax(Q[state[0], state[1]]))
            state, _, done = env.step(a)
            path.append(state)
            if done:
                break
        return path

    print("Q-learning greedy path:", extract_greedy_path(Q_qlearn, env))
    print("SARSA greedy path:     ", extract_greedy_path(Q_sarsa, env))
