import pygame
import random
import math

WIDTH = 500
HEIGHT = 500

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
running = True
r = random.randint(0,255)
g = random.randint(0,255)
b = random.randint(0,255)
x = 250
y = 0
speedX= 1
speedY= 1
state = 0
while running:
    
    for event in pygame.event.get():
    
        
        if event.type == pygame.QUIT:
            running = False  
        if event.type == pygame.KEYDOWN:
         if event.key == pygame.K_RETURN:
           pygame.image.save(screen, "image.jpg")
   
    x = x + speedX
    y = y + speedY
    #screen.fill((0,0,0))
    if x > 450:
     speedX = random.randint(1,5)
     speedX = - speedX
    if x < 50:
     speedX = random.randint(1,5)
     speedX = - speedX
    if y > 0:
     speedY = - speedY
    if y < 450:
     speedY = - speedY

    pygame.draw.rect(screen,(r,g,b),(x, y, 50,50))
    pygame.display.flip()
pygame.quit()