import pygame
import random

# Initialize Pygame
pygame.init()

# Set up the screen
screen_width = 400
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("Tetris")

# Define colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

# Define the shapes of the Tetrominos
tetrominos = [
    [[1, 1, 1],
     [0, 1, 0]],

    [[0, 2, 2],
     [2, 2, 0]],

    [[3, 3],
     [3, 3]],

    [[0, 4, 0],
     [4, 4, 4]],

    [[5, 0],
     [5, 0],
     [5, 5]],

    [[6, 6, 6, 6]]
]

# Define the colors of the Tetrominos
tetromino_colors = [RED, GREEN, BLUE, (255, 165, 0), (128, 0, 128), (255, 0, 255), (255, 255, 0)]

# Define a class for the Tetrominos
class Tetromino:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.shape = random.choice(tetrominos)
        self.color = random.choice(tetromino_colors)
        self.rotation = 0

    def draw(self):
        for y, row in enumerate(self.shape):
            for x, cell in enumerate(row):
                if cell:
                    pygame.draw.rect(screen, self.color, pygame.Rect(int((self.x + x) * 30), int((self.y + y) * 30), 30, 30), 0)

# Define the main function
def main():
    # Game variables
    clock = pygame.time.Clock()
    falling_piece = Tetromino(3, 0)

    # Game loop
    while True:
        screen.fill(BLACK)

        # Event handling
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_LEFT:
                    falling_piece.x -= 1
                elif event.key == pygame.K_RIGHT:
                    falling_piece.x += 1
                elif event.key == pygame.K_DOWN:
                    falling_piece.y += 1

        # Draw the falling piece
        falling_piece.draw()

        # Update the display
        pygame.display.flip()

        # Cap the frame rate
        clock.tick(30)

# Run the main function
if __name__ == "__main__":
    main()