import pygame
import random

# Initialize Pygame
pygame.init()

# Define constants
SCREEN_WIDTH = 300
SCREEN_HEIGHT = 600
BLOCK_SIZE = 30
GRID_WIDTH = SCREEN_WIDTH // BLOCK_SIZE
GRID_HEIGHT = SCREEN_HEIGHT // BLOCK_SIZE
FPS = 30
SLOW_SPEED = FPS // 2  # Number of frames between each automatic downward movement (slower)
FAST_SPEED = FPS * 10  # Number of frames between each automatic downward movement when spacebar is pressed

# Define colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
CYAN = (0, 255, 255)
MAGENTA = (255, 0, 255)
YELLOW = (255, 255, 0)
ORANGE = (255, 165, 0)

# Define shapes
SHAPES = [
    [[1, 1, 1],
     [0, 1, 0]],

    [[1, 1, 1, 1]],

    [[1, 1],
     [1, 1]],

    [[1, 1, 0],
     [0, 1, 1]],

    [[0, 1, 1],
     [1, 1, 0]],

    [[1, 1, 1],
     [1, 0, 0]],

    [[1, 1, 1],
     [0, 0, 1]]
]

# Define Tetris class
class Tetris:
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("Tetris")
        self.clock = pygame.time.Clock()

        self.grid = [[0] * GRID_WIDTH for _ in range(GRID_HEIGHT)]
        self.current_piece = self.new_piece()
        self.score = 0
        self.fall_counter = 0  # Counter for controlling the falling speed
        self.fall_speed = SLOW_SPEED

    def new_piece(self):
        return {
            'shape': random.choice(SHAPES),
            'x': GRID_WIDTH // 2 - 1,
            'y': 0,
            'color': random.choice([RED, GREEN, BLUE, CYAN, MAGENTA, YELLOW, ORANGE])
        }

    def draw_grid(self):
        for y in range(GRID_HEIGHT):
            for x in range(GRID_WIDTH):
                pygame.draw.rect(self.screen, WHITE, (x * BLOCK_SIZE, y * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE), 1)
                if self.grid[y][x] != 0:
                    pygame.draw.rect(self.screen, self.grid[y][x], (x * BLOCK_SIZE + 1, y * BLOCK_SIZE + 1, BLOCK_SIZE - 2, BLOCK_SIZE - 2))

    def draw_piece(self, piece):
        shape = piece['shape']
        color = piece['color']
        for y in range(len(shape)):
            for x in range(len(shape[y])):
                if shape[y][x] == 1:
                    pygame.draw.rect(self.screen, color, ((piece['x'] + x) * BLOCK_SIZE + 1, (piece['y'] + y) * BLOCK_SIZE + 1, BLOCK_SIZE - 2, BLOCK_SIZE - 2))

    def move_piece(self, dx, dy):
        if self.valid_move(self.current_piece, dx, dy):
            self.current_piece['x'] += dx
            self.current_piece['y'] += dy
            return True
        return False

    def valid_move(self, piece, dx, dy):
        shape = piece['shape']
        x = piece['x']
        y = piece['y']
        for i in range(len(shape)):
            for j in range(len(shape[i])):
                if shape[i][j] == 1:
                    new_x = x + j + dx
                    new_y = y + i + dy
                    if new_x < 0 or new_x >= GRID_WIDTH or new_y >= GRID_HEIGHT or (new_y >= 0 and self.grid[new_y][new_x] != 0):
                        return False
        return True

    def rotate_piece(self):
        rotated_shape = [list(row) for row in zip(*self.current_piece['shape'][::-1])]
        if self.valid_move({'shape': rotated_shape, 'x': self.current_piece['x'], 'y': self.current_piece['y']}, 0, 0):
            self.current_piece['shape'] = rotated_shape

    def place_piece(self):
        shape = self.current_piece['shape']
        x = self.current_piece['x']
        y = self.current_piece['y']
        color = self.current_piece['color']
        for i in range(len(shape)):
            for j in range(len(shape[i])):
                if shape[i][j] == 1:
                    self.grid[y + i][x + j] = color

    def clear_lines(self):
        lines_cleared = 0
        for y in range(GRID_HEIGHT):
            if all(self.grid[y]):
                del self.grid[y]
                self.grid.insert(0, [0] * GRID_WIDTH)
                lines_cleared += 1
        self.score += lines_cleared * 100

    def game_over(self):
        return any(self.grid[0])

    def run(self):
        while True:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    return
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_LEFT:
                        self.move_piece(-1, 0)
                    elif event.key == pygame.K_RIGHT:
                        self.move_piece(1, 0)
                    elif event.key == pygame.K_DOWN:
                        self.fall_speed = FAST_SPEED
                    elif event.key == pygame.K_UP:
                        self.rotate_piece()
                    elif event.key == pygame.K_SPACE:
                        while self.move_piece(0, 1):
                            pass

                elif event.type == pygame.KEYUP:
                    if event.key == pygame.K_DOWN:
                        self.fall_speed = SLOW_SPEED

            self.fall_counter += 1
            if self.fall_counter >= self.fall_speed:
                self.fall_counter = 0
                if not self.move_piece(0, 1):
                    self.place_piece()
                    self.clear_lines()
                    if self.game_over():
                        print("Game Over!")
                        pygame.quit()
                        return
                    self.current_piece = self.new_piece()
                    self.fall_speed = SLOW_SPEED

            self.screen.fill(BLACK)
            self.draw_grid()
            self.draw_piece(self.current_piece)
            pygame.display.flip()
            self.clock.tick(FPS)

            pygame.display.set_caption("Tetris | Score: {}".format(self.score))

# Run the game
if __name__ == '__main__':
    Tetris().run()