Small refactor and optimization

[psergicv]

When running the code, there was an error message "TypeError: Iterating over key states is not supported" comes from this section of your code:

keys = pygame.key.get_pressed()

for key in keys:
    # ...

Instead, you should check the state of specific keys directly, without a loop.

Also added Error Handlings in the main function. It should make the code more robust and easier to debug.