from PIL import Image
import os
import random

# 定义瓦片图的尺寸
TILE_SIZE = 32

# 获取目录下的所有瓦片图
def get_tiles(path):
    tiles = []
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if os.path.isfile(file_path):
            try:
                with Image.open(file_path) as tile:
                    if tile.size == (TILE_SIZE, TILE_SIZE):
                        tiles.append(tile)
            except:
                pass

    return tiles


# 将主图切割成若干块小图
def get_blocks(image):
    width, height = image.size
    blocks = []
    for y in range(0, height, TILE_SIZE):
        for x in range(0, width, TILE_SIZE):
            box = (x, y, x + TILE_SIZE, y + TILE_SIZE)
            blocks.append(image.crop(box))

    return blocks


# 获取小图和目标颜色的差异度
def color_distance(c1, c2):
    r1, g1, b1 = c1
    r2, g2, b2 = c2
    return (r1 - r2) ** 2 + (g1 - g2) ** 2 + (b1 - b2) ** 2


# 从瓦片图列表中找出最匹配目标颜色的瓦片图，并返回其索引和差异度
def find_best_match(tiles, target_color):
    best_index = 0
    best_distance = color_distance(tiles[0].getpixel((0, 0)), target_color)
    for i, tile in enumerate(tiles[1:]):
        distance = color_distance(tile.getpixel((0, 0)), target_color)
        if distance < best_distance:
            best_index = i + 1
            best_distance = distance

    return best_index, best_distance


# 生成马赛克
def create_mosaic(image, tiles):
    # 将主图切割成若干块小图
    blocks = get_blocks(image)
    # 创建和主图大小相同的画布
    mosaic = Image.new('RGB', image.size)
    # 遍历每个小图，选取和其颜色最匹配的瓦片图来填充
    for i, block in enumerate(blocks):
        print(f'\rProcessing block {i+1} of {len(blocks)}', end='')
        target_color = block.getpixel((TILE_SIZE//2, TILE_SIZE//2))
        best_index, _ = find_best_match(tiles, target_color)
        mosaic.paste(tiles[best_index], block.getbbox())

    return mosaic


if __name__ == '__main__':
    # 加载主图
    with Image.open('taj_mahal.jpg') as image:
        # 获取瓦片图
        tiles = get_tiles('tiles')
        # 生成马赛克
        mosaic = create_mosaic(image, tiles)
        mosaic.show()