import pygame
import random

# Initialize Pygame
pygame.init()

# Define constants
SCREEN_WIDTH, SCREEN_HEIGHT = 300, 600
GRID_SIZE = 30
GRID_WIDTH, GRID_HEIGHT = SCREEN_WIDTH // GRID_SIZE, SCREEN_HEIGHT // GRID_SIZE

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

# Colors for shapes
COLORS = [
    (0, 255, 255),
    (255, 255, 0),
    (128, 0, 128),
    (255, 0, 0),
    (0, 0, 255),
    (0, 255, 0),
    (255, 165, 0)
]

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_shape = self.new_shape()
        self.next_shape = self.new_shape()
        self.shape_pos = [0, GRID_WIDTH // 2 - len(self.current_shape[0]) // 2]
        self.fall_time = 0
        self.level = 1
        self.score = 0

    def new_shape(self):
        return random.choice(SHAPES), random.choice(COLORS)

    def rotate_shape(self):
        shape, color = self.current_shape
        rotated = [[shape[y][x] for y in range(len(shape))] for x in range(len(shape[0]) - 1, -1, -1)]
        if not self.collision(rotated, self.shape_pos):
            self.current_shape = rotated, color

    def collision(self, shape, offset):
        off_y, off_x = offset
        for y, row in enumerate(shape):
            for x, cell in enumerate(row):
                if cell and (y + off_y >= GRID_HEIGHT or x + off_x < 0 or x + off_x >= GRID_WIDTH or self.grid[y + off_y][x + off_x]):
                    return True
        return False

    def clear_lines(self):
        new_grid = [row for row in self.grid if any(cell == 0 for cell in row)]
        cleared_lines = GRID_HEIGHT - len(new_grid)
        self.grid = [[0] * GRID_WIDTH for _ in range(cleared_lines)] + new_grid
        self.score += cleared_lines ** 2
        self.level = 1 + self.score // 10

    def place_shape(self):
        shape, color = self.current_shape
        for y, row in enumerate(shape):
            for x, cell in enumerate(row):
                if cell:
                    self.grid[self.shape_pos[0] + y][self.shape_pos[1] + x] = color
        self.clear_lines()
        self.current_shape = self.next_shape
        self.next_shape = self.new_shape()
        self.shape_pos = [0, GRID_WIDTH // 2 - len(self.current_shape[0][0]) // 2]
        if self.collision(self.current_shape[0], self.shape_pos):
            self.__init__()  # Restart the game if new shape collides at spawn

    def move(self, dx):
        new_pos = [self.shape_pos[0], self.shape_pos[1] + dx]
        if not self.collision(self.current_shape[0], new_pos):
            self.shape_pos = new_pos

    def drop(self):
        self.fall_time += self.clock.get_rawtime()
        self.clock.tick()
        if self.fall_time > 1000 // self.level:
            self.fall_time = 0
            self.shape_pos[0] += 1
            if self.collision(self.current_shape[0], self.shape_pos):
                self.shape_pos[0] -= 1
                self.place_shape()

    def hard_drop(self):
        while not self.collision(self.current_shape[0], [self.shape_pos[0] + 1, self.shape_pos[1]]):
            self.shape_pos[0] += 1
        self.place_shape()

    def draw_grid(self):
        for y in range(GRID_HEIGHT):
            for x in range(GRID_WIDTH):
                if self.grid[y][x]:
                    pygame.draw.rect(self.screen, self.grid[y][x], (x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE))
        for x in range(GRID_WIDTH):
            pygame.draw.line(self.screen, (128, 128, 128), (x * GRID_SIZE, 0), (x * GRID_SIZE, SCREEN_HEIGHT))
        for y in range(GRID_HEIGHT):
            pygame.draw.line(self.screen, (128, 128, 128), (0, y * GRID_SIZE), (SCREEN_WIDTH, y * GRID_SIZE))

    def draw_shape(self):
        shape, color = self.current_shape
        for y, row in enumerate(shape):
            for x, cell in enumerate(row):
                if cell:
                    pygame.draw.rect(self.screen, color, ((self.shape_pos[1] + x) * GRID_SIZE, (self.shape_pos[0] + y) * GRID_SIZE, GRID_SIZE, GRID_SIZE))

    def draw_next_shape(self):
        shape, color = self.next_shape
        for y, row in enumerate(shape):
            for x, cell in enumerate(row):
                if cell:
                    pygame.draw.rect(self.screen, color, (SCREEN_WIDTH + x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE))

    def run(self):
        running = True
        while running:
            self.screen.fill((0, 0, 0))
            self.drop()
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_LEFT:
                        self.move(-1)
                    if event.key == pygame.K_RIGHT:
                        self.move(1)
                    if event.key == pygame.K_DOWN:
                        self.drop()
                    if event.key == pygame.K_UP:
                        self.rotate_shape()
                    if event.key == pygame.K_SPACE:
                        self.hard_drop()
            self.draw_grid()
            self.draw_shape()
            pygame.display.flip()

# Run the game
if __name__ == "__main__":
    game = Tetris()
    game.run()
    pygame.quit()