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

    print('empty x: ' + str(empty_x) + ', empty y: ' + str(empty_y))

    if empty_y > 0:
        changed = (grid[empty_y][empty_x], grid[empty_y - 1][empty_x])
        grid[empty_y - 1][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
            )