grid_x_count = 4
grid_y_count = 4

grid = []

for y in range(grid_y_count):
    grid.append([])
    for x in range(grid_x_count):
        grid[y].append(y * grid_x_count + x + 1)

def on_key_down(key):
    for y in range(grid_y_count):
        for x in range(grid_x_count):
            if grid[y][x] == grid_x_count * grid_y_count:
                empty_x = x
                empty_y = y

    new_empty_y = empty_y

    if key == keys.DOWN:
        new_empty_y -= 1
    elif key == keys.UP:
        new_empty_y += 1

    if 0 <= new_empty_y < grid_y_count:
        changed = (grid[empty_y][empty_x], grid[new_empty_y][empty_x])
        grid[new_empty_y][empty_x], grid[empty_y][empty_x] = changed

def draw():
    screen.fill((0, 0, 0))

    for y in range(grid_y_count):
        for x in range(grid_x_count):
            if grid[y][x] == grid_x_count * grid_y_count:
                continue

            piece_size = 100
            piece_draw_size = piece_size - 1

            screen.draw.filled_rect(
                Rect(
                    x * piece_size, y * piece_size,
                    piece_draw_size, piece_draw_size
                ),
                color=(100, 20, 150)
            )

            screen.draw.text(
                str(grid[y][x]),
                (x * piece_size, y * piece_size),
                fontsize=60
            )