import pygame
import random

# Initialize Pygame
pygame.init()

# Constants
SCREEN_WIDTH = 300
SCREEN_HEIGHT = 600
BLOCK_SIZE = 30
GRID_WIDTH = SCREEN_WIDTH // BLOCK_SIZE
GRID_HEIGHT = SCREEN_HEIGHT // BLOCK_SIZE
WHITE = (255, 255, 255)
GRAY = (128, 128, 128)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
CYAN = (0, 255, 255)
MAGENTA = (255, 0, 255)
YELLOW = (255, 255, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
ORANGE = (255, 165, 0)
FONT_SIZE = 24

# Tetrominos
tetrominos = [
    [[1, 1, 1],
     [0, 1, 0]],

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

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

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

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

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

    [[7, 7],
     [7, 7]]
]

# Initialize the screen
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Tetris")

# Load font
font = pygame.font.Font(None, FONT_SIZE)

# Function to draw a block
def draw_block(x, y, color):
    pygame.draw.rect(screen, color, pygame.Rect(x * BLOCK_SIZE, y * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE), 0)
    pygame.draw.rect(screen, BLACK, pygame.Rect(x * BLOCK_SIZE, y * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE), 1)

# Function to draw the grid
def draw_grid():
    for x in range(GRID_WIDTH):
        for y in range(GRID_HEIGHT):
            draw_block(x, y, GRAY)

# Function to create a new piece
def new_piece():
    shape = random.choice(tetrominos)
    piece = {
        'shape': shape,
        'x': GRID_WIDTH // 2 - len(shape[0]) // 2,
        'y': 0,
        'color': random.choice([RED, CYAN, MAGENTA, YELLOW, GREEN, BLUE, ORANGE])
    }
    return piece

# Function to draw a piece
def draw_piece(piece):
    shape = piece['shape']
    color = piece['color']
    for y, row in enumerate(shape):
        for x, block in enumerate(row):
            if block:
                draw_block(piece['x'] + x, piece['y'] + y, color)

# Function to check if a piece is valid
def valid(piece, offset):
    shape = piece['shape']
    for y, row in enumerate(shape):
        for x, block in enumerate(row):
            if block:
                if piece['y'] + y + offset[1] >= GRID_HEIGHT or \
                   piece['x'] + x + offset[0] < 0 or \
                   piece['x'] + x + offset[0] >= GRID_WIDTH or \
                   grid[piece['y'] + y + offset[1]][piece['x'] + x + offset[0]]:
                    return False
    return True

# Function to rotate a piece
def rotate(piece):
    piece['shape'] = [list(row) for row in zip(*reversed(piece['shape']))]

# Function to draw the score
def draw_score(score):
    score_text = font.render("Score: " + str(score), True, BLACK)
    screen.blit(score_text, (10, 10))

# Main game loop
def main():
    global grid
    clock = pygame.time.Clock()
    fall_time = 0
    fall_speed = 0.5
    piece = new_piece()
    grid = [[0] * GRID_WIDTH for _ in range(GRID_HEIGHT)]
    score = 0

    while True:
        screen.fill(WHITE)
        draw_grid()
        draw_piece(piece)
        draw_score(score)
        pygame.display.update()

        # 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 and valid(piece, (-1, 0)):
                    piece['x'] -= 1
                if event.key == pygame.K_RIGHT and valid(piece, (1, 0)):
                    piece['x'] += 1
                if event.key == pygame.K_DOWN and valid(piece, (0, 1)):
                    piece['y'] += 1
                if event.key == pygame.K_UP:
                    rotate(piece)

        # Gravity
        if pygame.time.get_ticks() - fall_time > fall_speed * 1000:
            if not valid(piece, (0, 1)):
                for y, row in enumerate(piece['shape']):
                    for x, block in enumerate(row):
                        if block:
                            grid[piece['y'] + y][piece['x'] + x] = piece['color']
                piece = new_piece()

                # Check for completed rows
                for y in range(GRID_HEIGHT):
                    if all(grid[y]):
                        grid.pop(y)
                        grid.insert(0, [0] * GRID_WIDTH)
                        score += 10

            piece['y'] += 1
            fall_time = pygame.time.get_ticks()

        clock.tick(30)

if __name__ == "__main__":
    main()