import pygame
import random

WIDTH = 500
HEIGHT = 500
FPS = 60

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()
X = 250
speedY = 0
statefigure = 1
running = True
while running:
    
    clock.tick(FPS)
    
    for event in pygame.event.get():
        
        if event.type == pygame.QUIT:
            running = False
        
        if event.type == pygame.KEYDOWN:
         if event.key == pygame.K_RETURN:
             if (statefigure == 4):
                 statefigure = 1
             if (statefigure == 3):
                statefigure = 4
             if (statefigure == 2):
                statefigure = 3
             if (statefigure == 1):
                statefigure = 2
   
    
    speedY = speedY + 1
    screen.fill(BLACK)
    if (statefigure == 1):
          pygame.draw.rect(screen,RED, (X,speedY,20,20))
          pygame.draw.rect(screen,GREEN,(X,speedY+20,20,20))
          pygame.draw.rect(screen,BLUE,(X,speedY+40,20,20))
          pygame.draw.rect(screen,WHITE,(X,speedY+60,20,20))
    if (statefigure == 2):
          pygame.draw.rect(screen,RED, (X,speedY,20,20))
          pygame.draw.rect(screen,GREEN,(X+20,speedY,20,20))
          pygame.draw.rect(screen,BLUE,(X+40,speedY,20,20))
          pygame.draw.rect(screen,WHITE,(X+60,speedY,20,20))
    if (statefigure == 3):
          pygame.draw.rect(screen,RED, (X,speedY,20,20))
          pygame.draw.rect(screen,GREEN,(X,speedY-20,20,20))
          pygame.draw.rect(screen,BLUE,(X,speedY-40,20,20))
          pygame.draw.rect(screen,WHITE,(X,speedY-60,20,20))
    if (statefigure == 4):
        pygame.draw.rect(screen,RED, (X,speedY,20,20))
        pygame.draw.rect(screen,GREEN,(X-20,speedY,20,20))
        pygame.draw.rect(screen,BLUE,(X-40,speedY,20,20))
        pygame.draw.rect(screen,WHITE,(X-60,speedY,20,20))     
    pygame.display.flip()
    if speedY > 475:
        speedY = 0
pygame.quit()