import pygame
import numpy as np
import random

# Pygame setup
pygame.init()
CELL_SIZE = 30
ROWS, COLS = 20, 10
WIDTH, HEIGHT = COLS * CELL_SIZE, ROWS * CELL_SIZE
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()

# Tetromino definitions
TETROMINOS = {
    'I': np.array([[1, 1, 1, 1]]),
    'O': np.array([[1, 1], [1, 1]]),
    'T': np.array([[0, 1, 0], [1, 1, 1]]),
    'S': np.array([[0, 1, 1], [1, 1, 0]]),
    'Z': np.array([[1, 1, 0], [0, 1, 1]]),
    'J': np.array([[1, 0, 0], [1, 1, 1]]),
    'L': np.array([[0, 0, 1], [1, 1, 1]])
}

# Colors for pieces
COLORS = {
    0: (0, 0, 0),
    1: (255, 255, 255)
}

# Game state
grid = np.zeros((ROWS, COLS), dtype=int)

def new_piece():
    shape = random.choice(list(TETROMINOS.values()))
    return shape, [0, COLS // 2 - shape.shape[1] // 2]

def check_collision(grid, piece, offset):
    x, y = offset
    h, w = piece.shape
    if x + h > ROWS or y < 0 or y + w > COLS:
        return True
    area = grid[x:x+h, y:y+w]
    return np.any(area & piece)

def place_piece(grid, piece, offset):
    x, y = offset
    h, w = piece.shape
    grid[x:x+h, y:y+w] |= piece

def clear_lines(grid):
    full_rows = [i for i in range(ROWS) if all(grid[i])]
    for i in full_rows:
        grid[1:i+1] = grid[0:i]
        grid[0] = np.zeros(COLS)

def rotate_piece(piece):
    return np.rot90(piece, -1)

def draw_grid(surface, grid):
    for i in range(ROWS):
        for j in range(COLS):
            pygame.draw.rect(surface, COLORS[grid[i][j]],
                             (j * CELL_SIZE, i * CELL_SIZE, CELL_SIZE, CELL_SIZE), 0)
            pygame.draw.rect(surface, (40, 40, 40),
                             (j * CELL_SIZE, i * CELL_SIZE, CELL_SIZE, CELL_SIZE), 1)

def draw_piece(surface, piece, offset):
    x, y = offset
    for i in range(piece.shape[0]):
        for j in range(piece.shape[1]):
            if piece[i][j]:
                pygame.draw.rect(surface, COLORS[1],
                                 ((y + j) * CELL_SIZE, (x + i) * CELL_SIZE, CELL_SIZE, CELL_SIZE))

# Start game
piece, pos = new_piece()
drop_timer = 0
running = True

while running:
    screen.fill((0, 0, 0))
    draw_grid(screen, grid)
    draw_piece(screen, piece, pos)
    pygame.display.flip()
    clock.tick(60)
    drop_timer += 1

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            new_pos = pos[:]
            if event.key == pygame.K_LEFT:
                new_pos[1] -= 1
            elif event.key == pygame.K_RIGHT:
                new_pos[1] += 1
            elif event.key == pygame.K_DOWN:
                new_pos[0] += 1
            elif event.key == pygame.K_UP:
                rotated = rotate_piece(piece)
                if not check_collision(grid, rotated, pos):
                    piece = rotated
                continue

            if not check_collision(grid, piece, new_pos):
                pos = new_pos

    if drop_timer > 30:
        drop_pos = [pos[0] + 1, pos[1]]
        if not check_collision(grid, piece, drop_pos):
            pos = drop_pos
        else:
            place_piece(grid, piece, pos)
            clear_lines(grid)
            piece, pos = new_piece()
            if check_collision(grid, piece, pos):
                print("Game Over")
                running = False
        drop_timer = 0

pygame.quit()