import pygame
import random
import math

WIDTH = 500
HEIGHT = 500

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
running = True

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

    if state == 0:
     for x in range(0,500,50):
      for y in range(0,500,50):
       r = random.randint(0,255)
       g = random.randint(0,255)
       b = random.randint(0,255)
       pygame.draw.rect(screen,(r,g,b),(x, y, 50,50))
     state = 1
     pygame.display.flip()
pygame.quit()