"""
Wave Function Collapse (波函数坍缩) 图像生成器

这个程序使用 WFC 算法从输入图像中提取瓦片，并生成新的相似纹理图像。
算法通过约束传播确保相邻瓦片的边缘匹配，从而生成连贯的图像。
"""

import os
import random
import time
from datetime import datetime
from typing import Optional

import pygame

# ============================================================================
# 初始化
# ============================================================================

pygame.init()

MAX_INT: int = 2**31 - 1


# ============================================================================
# 配置类
# ============================================================================

class Config:
    """全局配置参数"""

    # 瓦片来源设置
    TILE_SOURCE_MODE: str = "folder"  # "image" 或 "folder" - 从图像切割或从文件夹读取
    TILE_FOLDER_PATH: str = "tiles"  # 当模式为"folder"时，从此文件夹读取PNG文件作为瓦片

    # 输入图像设置
    ORIGINAL_IMG_PATH: str = "Water.png"  # 源图像路径（仅在image模式下使用）
    TILE_SIZE: int = 4  # 瓦片尺寸（像素），决定提取多大的图案块（仅在image模式下使用）

    # 输出图像设置
    OUTPUT_WIDTH: int = 24  # 输出图像宽度（瓦片数量）
    OUTPUT_HEIGHT: int = 24  # 输出图像高度（瓦片数量）
    SCALE: float = 1.0  # 显示缩放比例

    # 瓦片变换设置
    ALLOW_ROTATE: bool = True  # 是否允许旋转瓦片（增加变化性）
    ALLOW_FLIP: bool = True  # 是否允许翻转瓦片（增加变化性）

    # 显示设置
    SHOW_GENERATION: bool = True  # 是否实时显示生成过程
    GENERATION_DELAY: float = 0.0001  # 每步之间的延迟（秒），0表示最快速度

    # 输出设置
    SAVE_OUTPUT: bool = True  # 是否保存生成的图像
    OUTPUT_DIR: str = "output"  # 输出目录

    # Pygame显示对象（延迟创建）
    DISPLAY: pygame.Surface = None


# 确保输出尺寸为整数
Config.OUTPUT_WIDTH = int(Config.OUTPUT_WIDTH)
Config.OUTPUT_HEIGHT = int(Config.OUTPUT_HEIGHT)


# ============================================================================
# 瓦片类
# ============================================================================

class Tile:
    """
    表示一个图像瓦片（图案块）

    瓦片是从源图像中提取的小块图案，包含完整的像素数据和边缘信息。
    边缘信息用于WFC算法中的瓦片匹配。
    """

    def __init__(self, pixels: list[list[tuple]]):
        """
        初始化瓦片

        参数:
            pixels: 2D数组 [y][x]，每个元素是RGBA颜色元组
        """
        self.pixels = pixels
        self.height = len(pixels)
        self.width = len(pixels[0]) if pixels else 0

        # 提取四条边缘（用于匹配相邻瓦片）
        self.edges = {
            'up': tuple(pixels[0]),  # 顶部边缘（第一行）
            'right': tuple(row[-1] for row in pixels),  # 右侧边缘（最后一列）
            'down': tuple(pixels[-1]),  # 底部边缘（最后一行）
            'left': tuple(row[0] for row in pixels)  # 左侧边缘（第一列）
        }

        # 创建Pygame Surface用于渲染
        self.surface = pygame.Surface((self.width, self.height), pygame.SRCALPHA)
        for y in range(self.height):
            for x in range(self.width):
                self.surface.set_at((x, y), pixels[y][x])

    def rotate_90(self) -> 'Tile':
        """
        顺时针旋转90度

        返回:
            旋转后的新瓦片对象
        """
        new_pixels = [
            [self.pixels[self.height - 1 - x][y] for x in range(self.height)]
            for y in range(self.width)
        ]
        return Tile(new_pixels)

    def flip_horizontal(self) -> 'Tile':
        """
        水平翻转（左右镜像）

        返回:
            翻转后的新瓦片对象
        """
        new_pixels = [row[::-1] for row in self.pixels]
        return Tile(new_pixels)


# ============================================================================
# 瓦片加载和处理函数
# ============================================================================

def load_tiles_from_image(image_path: str, tile_size: int) -> list[Tile]:
    """
    使用滑动窗口从图像中提取瓦片

    通过在源图像上滑动一个tile_size×tile_size的窗口，
    每次移动1个像素，提取所有可能的图案块。

    参数:
        image_path: 源图像路径
        tile_size: 瓦片尺寸（正方形边长）

    返回:
        提取的唯一瓦片列表
    """
    img = pygame.image.load(image_path)
    width, height = img.get_size()

    tiles = []
    tile_set = set()  # 用于去重（避免重复的图案）

    # 滑动窗口提取瓦片（步长为1像素）
    for y in range(height - tile_size + 1):
        for x in range(width - tile_size + 1):
            # 提取tile_size × tile_size的像素块
            pixels = []
            for dy in range(tile_size):
                row = []
                for dx in range(tile_size):
                    color = img.get_at((x + dx, y + dy))
                    row.append((color.r, color.g, color.b, color.a))
                pixels.append(row)

            # 使用边缘元组作为唯一标识进行去重
            tile = Tile(pixels)
            edge_signature = (
                tile.edges['up'],
                tile.edges['right'],
                tile.edges['down'],
                tile.edges['left']
            )

            if edge_signature not in tile_set:
                tile_set.add(edge_signature)
                tiles.append(tile)

    print(f"从图像提取了 {len(tiles)} 个唯一瓦片")
    return tiles


def load_tiles_from_folder(folder_path: str) -> list[Tile]:
    """
    从文件夹中加载所有PNG图像作为瓦片

    读取指定文件夹中的所有PNG文件，每个文件作为一个瓦片。
    所有瓦片必须具有相同的尺寸。

    参数:
        folder_path: 包含PNG瓦片文件的文件夹路径

    返回:
        加载的瓦片列表
    """
    if not os.path.exists(folder_path):
        raise FileNotFoundError(f"瓦片文件夹不存在: {folder_path}")

    # 获取所有PNG文件
    png_files = [f for f in os.listdir(folder_path) if f.lower().endswith('.png')]

    if not png_files:
        raise ValueError(f"文件夹中没有找到PNG文件: {folder_path}")

    print(f"在文件夹 '{folder_path}' 中找到 {len(png_files)} 个PNG文件")

    tiles = []
    tile_set = set()
    expected_size = None

    for i, filename in enumerate(sorted(png_files)):
        filepath = os.path.join(folder_path, filename)

        try:
            img = pygame.image.load(filepath)
            width, height = img.get_size()

            # 检查所有瓦片尺寸是否一致
            if expected_size is None:
                expected_size = (width, height)
                print(f"  瓦片尺寸: {width}x{height} 像素")
            elif (width, height) != expected_size:
                print(f"  ⚠ 跳过 '{filename}': 尺寸不匹配 ({width}x{height}, 期望 {expected_size[0]}x{expected_size[1]})")
                continue

            # 提取所有像素
            pixels = []
            for y in range(height):
                row = []
                for x in range(width):
                    color = img.get_at((x, y))
                    row.append((color.r, color.g, color.b, color.a))
                pixels.append(row)

            # 创建瓦片并去重
            tile = Tile(pixels)
            edge_signature = (
                tile.edges['up'],
                tile.edges['right'],
                tile.edges['down'],
                tile.edges['left']
            )

            if edge_signature not in tile_set:
                tile_set.add(edge_signature)
                tiles.append(tile)
                print(f"  ✓ 加载: {filename}")
            else:
                print(f"  - 跳过重复: {filename}")

        except Exception as e:
            print(f"  ✗ 加载失败 '{filename}': {e}")
            continue

    if not tiles:
        raise ValueError("没有成功加载任何瓦片")

    print(f"成功加载 {len(tiles)} 个唯一瓦片")
    return tiles


def generate_tile_variants(
    base_tiles: list[Tile],
    allow_rotate: bool,
    allow_flip: bool
) -> list[Tile]:
    """
    生成瓦片的旋转和翻转变体

    通过旋转和翻转原始瓦片，可以增加可用瓦片的数量和多样性，
    从而生成更丰富的图案。

    参数:
        base_tiles: 原始瓦片列表
        allow_rotate: 是否生成旋转变体（0°, 90°, 180°, 270°）
        allow_flip: 是否生成翻转变体

    返回:
        包含所有变体的瓦片列表
    """
    all_tiles = []
    tile_set = set()

    for tile in base_tiles:
        variants = [tile]

        # 生成旋转变体（90°, 180°, 270°）
        if allow_rotate:
            t = tile
            for _ in range(3):
                t = t.rotate_90()
                variants.append(t)

        # 生成翻转变体
        if allow_flip:
            flipped = tile.flip_horizontal()
            variants.append(flipped)

            # 翻转后的瓦片也可以旋转
            if allow_rotate:
                t = flipped
                for _ in range(3):
                    t = t.rotate_90()
                    variants.append(t)

        # 去重（有些变体可能相同，如对称图案）
        for v in variants:
            sig = (v.edges['up'], v.edges['right'], v.edges['down'], v.edges['left'])
            if sig not in tile_set:
                tile_set.add(sig)
                all_tiles.append(v)

    return all_tiles


def build_adjacency_rules(tiles: list[Tile]) -> dict:
    """
    构建瓦片邻接规则

    计算哪些瓦片可以相邻放置。两个瓦片可以相邻当且仅当
    它们接触的边缘像素完全相同。

    参数:
        tiles: 所有可用的瓦片列表

    返回:
        邻接规则字典，格式为 {瓦片索引: {'up': [...], 'right': [...], ...}}
    """
    rules = {
        i: {'up': [], 'right': [], 'down': [], 'left': []}
        for i in range(len(tiles))
    }

    # 对每一对瓦片，检查它们是否可以相邻
    for i, tile1 in enumerate(tiles):
        for j, tile2 in enumerate(tiles):
            # 检查tile1的右边缘是否匹配tile2的左边缘
            if tile1.edges['right'] == tile2.edges['left']:
                rules[i]['right'].append(j)
                rules[j]['left'].append(i)

            # 检查tile1的下边缘是否匹配tile2的上边缘
            if tile1.edges['down'] == tile2.edges['up']:
                rules[i]['down'].append(j)
                rules[j]['up'].append(i)

    return rules


# ============================================================================
# WFC网格类
# ============================================================================

class WFCGrid:
    """
    Wave Function Collapse 算法的核心网格类

    维护一个网格，其中每个格子都有一组可能的瓦片选项。
    通过不断"塌缩"（选择）格子的瓦片并传播约束，最终生成完整图像。
    """

    def __init__(
        self,
        width: int,
        height: int,
        tiles: list[Tile],
        rules: dict
    ):
        """
        初始化WFC网格

        参数:
            width: 网格宽度（瓦片数量）
            height: 网格高度（瓦片数量）
            tiles: 可用的瓦片列表
            rules: 瓦片邻接规则
        """
        self.width = width
        self.height = height
        self.tiles = tiles
        self.rules = rules
        self.num_tiles = len(tiles)

        # 每个格子的可能选项（初始时每个格子都可以是任意瓦片）
        self.grid = [
            [set(range(self.num_tiles)) for _ in range(width)]
            for _ in range(height)
        ]

        # 已确定的瓦片（None表示未确定）
        self.collapsed = [
            [None for _ in range(width)]
            for _ in range(height)
        ]

    def get_min_entropy_cell(self) -> Optional[tuple[int, int]]:
        """
        找到熵最小的未塌缩格子

        熵表示一个格子的不确定性（可能选项的数量）。
        WFC算法优先选择熵最小的格子进行塌缩。

        返回:
            (x, y) 坐标，如果没有可用格子则返回None
        """
        min_entropy = self.num_tiles + 1
        candidates = []

        for y in range(self.height):
            for x in range(self.width):
                if self.collapsed[y][x] is None:
                    entropy = len(self.grid[y][x])

                    if entropy == 0:
                        return None  # 无解（矛盾）

                    if entropy < min_entropy:
                        min_entropy = entropy
                        candidates = [(x, y)]
                    elif entropy == min_entropy:
                        candidates.append((x, y))

        # 在熵相同的格子中随机选择一个
        return random.choice(candidates) if candidates else None

    def collapse(self, x: int, y: int):
        """
        塌缩指定格子（选择一个瓦片）

        从该格子的可能选项中随机选择一个瓦片，
        然后传播约束到相邻格子。

        参数:
            x: 格子的x坐标
            y: 格子的y坐标
        """
        if self.collapsed[y][x] is not None:
            return  # 已经塌缩过

        if not self.grid[y][x]:
            return  # 无可用选项

        # 随机选择一个瓦片
        tile_idx = random.choice(list(self.grid[y][x]))
        self.collapsed[y][x] = tile_idx
        self.grid[y][x] = {tile_idx}

        # 传播约束到相邻格子
        self.propagate(x, y)

    def propagate(self, start_x: int, start_y: int):
        """
        约束传播算法

        当一个格子被塌缩后，根据邻接规则更新相邻格子的可能选项。
        这个过程会递归地传播到更远的格子。

        参数:
            start_x: 起始格子的x坐标
            start_y: 起始格子的y坐标
        """
        stack = [(start_x, start_y)]

        while stack:
            x, y = stack.pop()

            # 检查四个方向的邻居
            directions = [
                (x, y - 1, 'up', 'down'),      # 上方邻居
                (x + 1, y, 'right', 'left'),   # 右方邻居
                (x, y + 1, 'down', 'up'),      # 下方邻居
                (x - 1, y, 'left', 'right')    # 左方邻居
            ]

            for nx, ny, my_dir, their_dir in directions:
                # 检查邻居是否在网格范围内
                if 0 <= nx < self.width and 0 <= ny < self.height:
                    if self.collapsed[ny][nx] is None:
                        # 计算邻居的有效选项
                        # （必须与当前格子的所有可能瓦片兼容）
                        valid_neighbors = set()
                        for tile_idx in self.grid[y][x]:
                            valid_neighbors.update(self.rules[tile_idx][my_dir])

                        # 更新邻居的选项（取交集）
                        old_options = self.grid[ny][nx]
                        new_options = old_options & valid_neighbors

                        # 如果选项发生变化，继续传播
                        if new_options != old_options:
                            self.grid[ny][nx] = new_options
                            if new_options:  # 只有非空时才继续传播
                                stack.append((nx, ny))

    def is_fully_collapsed(self) -> bool:
        """
        检查是否所有格子都已塌缩

        返回:
            True表示生成完成，False表示还有未确定的格子
        """
        for y in range(self.height):
            for x in range(self.width):
                if self.collapsed[y][x] is None:
                    return False
        return True

    def has_contradiction(self) -> bool:
        """
        检查是否存在矛盾

        矛盾指某个格子没有任何可用的瓦片选项，
        这意味着当前的塌缩序列无法完成。

        返回:
            True表示存在矛盾，需要重新开始
        """
        for y in range(self.height):
            for x in range(self.width):
                if self.collapsed[y][x] is None and not self.grid[y][x]:
                    return True
        return False


# ============================================================================
# WFC算法执行函数
# ============================================================================

def run_wfc(
    width: int,
    height: int,
    tiles: list[Tile],
    rules: dict,
    max_attempts: int = 10
) -> Optional[WFCGrid]:
    """
    运行WFC算法（无显示）

    参数:
        width: 输出宽度（瓦片数量）
        height: 输出高度（瓦片数量）
        tiles: 可用瓦片列表
        rules: 邻接规则
        max_attempts: 最大尝试次数

    返回:
        成功时返回完成的网格，失败返回None
    """
    for attempt in range(max_attempts):
        print(f"\n尝试第 {attempt + 1} 次生成...")
        grid = WFCGrid(width, height, tiles, rules)

        steps = 0
        while not grid.is_fully_collapsed():
            # 检查是否出现矛盾
            if grid.has_contradiction():
                print(f"  ✗ 发现矛盾（步骤 {steps}），重新开始")
                break

            # 找到熵最小的格子
            cell = grid.get_min_entropy_cell()
            if cell is None:
                print(f"  ✗ 无法继续（步骤 {steps}），重新开始")
                break

            # 塌缩该格子
            x, y = cell
            grid.collapse(x, y)
            steps += 1

            # 显示进度
            if Config.SHOW_GENERATION and steps % 50 == 0:
                uncollapsed = sum(
                    1 for row in grid.collapsed
                    for cell in row if cell is None
                )
                print(f"  进度: {steps} 步，剩余 {uncollapsed} 个格子")
                time.sleep(Config.GENERATION_DELAY)

        # 检查是否成功完成
        if grid.is_fully_collapsed():
            print(f"✓ 生成成功！共 {steps} 步")
            return grid

    print(f"\n⚠ {max_attempts} 次尝试后失败")
    return None


def run_wfc_with_display(
    width: int,
    height: int,
    tiles: list[Tile],
    rules: dict,
    display: pygame.Surface,
    scale: float,
    max_attempts: int = 10
) -> Optional[WFCGrid]:
    """
    运行WFC算法，并实时显示生成过程

    参数:
        width: 输出宽度（瓦片数量）
        height: 输出高度（瓦片数量）
        tiles: 可用瓦片列表
        rules: 邻接规则
        display: Pygame显示窗口
        scale: 显示缩放比例
        max_attempts: 最大尝试次数

    返回:
        成功时返回完成的网格，失败返回None
    """
    tile_size = tiles[0].width
    overlap = 1
    original_width = tile_size + (width - 1) * (tile_size - overlap)
    original_height = tile_size + (height - 1) * (tile_size - overlap)
    scaled_width = int(original_width * scale)
    scaled_height = int(original_height * scale)

    # 创建用于实时渲染的surface
    render_surface = pygame.Surface((original_width, original_height), pygame.SRCALPHA)
    clock = pygame.time.Clock()

    for attempt in range(max_attempts):
        print(f"\n尝试第 {attempt + 1} 次生成...")
        grid = WFCGrid(width, height, tiles, rules)
        render_surface.fill((50, 50, 50, 255))  # 灰色背景表示未确定的区域

        steps = 0
        while not grid.is_fully_collapsed():
            # 处理pygame事件（允许用户关闭窗口）
            for event in pygame.event.get():
                if event.type == pygame.QUIT or \
                   (event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE):
                    pygame.quit()
                    return None

            # 检查是否出现矛盾
            if grid.has_contradiction():
                print(f"  ✗ 发现矛盾（步骤 {steps}），重新开始")
                break

            # 找到熵最小的格子
            cell = grid.get_min_entropy_cell()
            if cell is None:
                print(f"  ✗ 无法继续（步骤 {steps}），重新开始")
                break

            # 塌缩该格子
            x, y = cell
            grid.collapse(x, y)
            steps += 1

            # 实时渲染刚刚塌缩的瓦片
            if Config.SHOW_GENERATION:
                tile_idx = grid.collapsed[y][x]
                if tile_idx is not None:
                    tile = tiles[tile_idx]
                    pos_x = x * (tile_size - overlap)
                    pos_y = y * (tile_size - overlap)

                    # 绘制瓦片的所有像素
                    for ty in range(tile_size):
                        for tx in range(tile_size):
                            pixel = tile.pixels[ty][tx]
                            dest_x = pos_x + tx
                            dest_y = pos_y + ty
                            render_surface.set_at((dest_x, dest_y), pixel)

                # 缩放并显示
                scaled_surface = pygame.transform.scale(
                    render_surface,
                    (scaled_width, scaled_height)
                )
                display.fill((255, 255, 255))
                display.blit(scaled_surface, (0, 0))

                # 显示统计信息
                if steps % 10 == 0:
                    uncollapsed = sum(
                        1 for row in grid.collapsed
                        for cell in row if cell is None
                    )
                    progress = (steps / (width * height)) * 100
                    pygame.display.set_caption(
                        f"WFC生成中... 步骤:{steps}/{width*height} "
                        f"进度:{progress:.1f}% 剩余:{uncollapsed}"
                    )

                pygame.display.flip()
                clock.tick(1000)  # 限制最大帧率

                if Config.GENERATION_DELAY > 0:
                    time.sleep(Config.GENERATION_DELAY)

        # 检查是否成功完成
        if grid.is_fully_collapsed():
            print(f"✓ 生成成功！共 {steps} 步")
            return grid

    print(f"\n⚠ {max_attempts} 次尝试后失败")
    return None


# ============================================================================
# 渲染和输出函数
# ============================================================================

def generate_output_filename() -> str:
    """
    生成包含配置信息的输出文件名

    文件名格式: 源图像名_瓦片大小_输出尺寸_变换_时间戳.png
    例如: Water_t4_24x24_RF_20240101_120000.png

    返回:
        生成的文件名
    """
    # 获取源图像名（不含扩展名）
    source_name = os.path.splitext(os.path.basename(Config.ORIGINAL_IMG_PATH))[0]

    # 生成时间戳
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    # 构建变换描述
    rotate_str = "R" if Config.ALLOW_ROTATE else ""
    flip_str = "F" if Config.ALLOW_FLIP else ""
    transform_str = rotate_str + flip_str if (rotate_str or flip_str) else "NoTrans"

    # 构建完整文件名
    filename = (
        f"{source_name}_t{Config.TILE_SIZE}_"
        f"{Config.OUTPUT_WIDTH}x{Config.OUTPUT_HEIGHT}_"
        f"{transform_str}_{timestamp}.png"
    )

    return filename


def render_grid(grid: WFCGrid, tiles: list[Tile], scale: float) -> pygame.Surface:
    """
    渲染网格到Surface，正确处理瓦片重叠

    相邻瓦片共享边缘像素，确保生成的图像无缝连接。

    参数:
        grid: 完成的WFC网格
        tiles: 瓦片列表
        scale: 缩放比例

    返回:
        渲染并缩放后的Surface
    """
    tile_size = tiles[0].width  # 瓦片的完整尺寸
    overlap = 1  # 相邻瓦片共享1个像素的边缘

    # 计算原始图像尺寸
    # 第一个瓦片占 tile_size 像素
    # 后续每个瓦片增加 (tile_size - overlap) 像素
    original_width = tile_size + (grid.width - 1) * (tile_size - overlap)
    original_height = tile_size + (grid.height - 1) * (tile_size - overlap)

    print(f"  瓦片尺寸: {tile_size}x{tile_size} 像素")
    print(f"  网格尺寸: {grid.width}x{grid.height} 个瓦片")
    print(f"  原始图像尺寸: {original_width}x{original_height} 像素")
    print(f"  缩放后尺寸: {int(original_width * scale)}x{int(original_height * scale)} 像素")

    # 创建渲染surface
    surface = pygame.Surface((original_width, original_height), pygame.SRCALPHA)

    # 渲染每个瓦片（相邻瓦片重叠绘制）
    for y in range(grid.height):
        for x in range(grid.width):
            tile_idx = grid.collapsed[y][x]
            if tile_idx is not None:
                tile = tiles[tile_idx]

                # 计算此瓦片在最终图像中的左上角位置
                pos_x = x * (tile_size - overlap)
                pos_y = y * (tile_size - overlap)

                # 绘制完整的瓦片（所有像素）
                for ty in range(tile_size):
                    for tx in range(tile_size):
                        pixel = tile.pixels[ty][tx]
                        dest_x = pos_x + tx
                        dest_y = pos_y + ty
                        surface.set_at((dest_x, dest_y), pixel)

    # 缩放surface
    scaled_width = int(original_width * scale)
    scaled_height = int(original_height * scale)
    scaled_surface = pygame.transform.scale(surface, (scaled_width, scaled_height))

    return scaled_surface


# ============================================================================
# 主函数
# ============================================================================

def main():
    """
    主程序入口

    执行流程:
    1. 加载并提取瓦片
    2. 生成瓦片变体
    3. 构建邻接规则
    4. 创建显示窗口
    5. 运行WFC算法
    6. 渲染并保存结果
    """
    print("=" * 60)
    print("Wave Function Collapse - 图像生成器")
    print("=" * 60)

    start_time = time.time()

    # ========================================================================
    # 第1步: 加载瓦片
    # ========================================================================
    print(f"\n【第1步】加载瓦片...")
    print(f"         模式: {Config.TILE_SOURCE_MODE}")

    if Config.TILE_SOURCE_MODE == "folder":
        # 从文件夹加载PNG瓦片
        print(f"         文件夹路径: {Config.TILE_FOLDER_PATH}")
        base_tiles = load_tiles_from_folder(Config.TILE_FOLDER_PATH)
    elif Config.TILE_SOURCE_MODE == "image":
        # 从图像切割瓦片
        print(f"         源图像: {Config.ORIGINAL_IMG_PATH}")
        print(f"         瓦片大小: {Config.TILE_SIZE}x{Config.TILE_SIZE}")
        base_tiles = load_tiles_from_image(Config.ORIGINAL_IMG_PATH, Config.TILE_SIZE)
    else:
        raise ValueError(f"不支持的瓦片来源模式: {Config.TILE_SOURCE_MODE}，请使用 'image' 或 'folder'")

    # ========================================================================
    # 第2步: 生成变体
    # ========================================================================
    print(f"\n【第2步】生成变体 (旋转={Config.ALLOW_ROTATE}, 翻转={Config.ALLOW_FLIP})...")
    all_tiles = generate_tile_variants(
        base_tiles,
        Config.ALLOW_ROTATE,
        Config.ALLOW_FLIP
    )
    print(f"         总共 {len(all_tiles)} 个瓦片")

    # ========================================================================
    # 第3步: 构建邻接规则
    # ========================================================================
    print("\n【第3步】构建邻接规则...")
    rules = build_adjacency_rules(all_tiles)

    # 统计规则
    total_rules = sum(
        len(r['up']) + len(r['right']) + len(r['down']) + len(r['left'])
        for r in rules.values()
    )
    print(f"         总共 {total_rules} 条邻接规则")

    # 显示前几个瓦片的规则统计
    print("\n         前5个瓦片的邻接规则:")
    for i in range(min(5, len(all_tiles))):
        print(f"           瓦片 {i}: 上={len(rules[i]['up'])}, "
              f"右={len(rules[i]['right'])}, "
              f"下={len(rules[i]['down'])}, "
              f"左={len(rules[i]['left'])}")

    # ========================================================================
    # 第4步: 创建显示窗口
    # ========================================================================
    tile_size = all_tiles[0].width
    overlap = 1
    original_width = tile_size + (Config.OUTPUT_WIDTH - 1) * (tile_size - overlap)
    original_height = tile_size + (Config.OUTPUT_HEIGHT - 1) * (tile_size - overlap)
    window_width = int(original_width * Config.SCALE)
    window_height = int(original_height * Config.SCALE)

    print(f"\n【第4步】创建窗口: {window_width}x{window_height} 像素")
    Config.DISPLAY = pygame.display.set_mode((window_width, window_height))
    pygame.display.set_caption("Wave Function Collapse - 初始化中...")

    # ========================================================================
    # 第5步: 运行WFC算法
    # ========================================================================
    print(f"\n【第5步】开始生成 {Config.OUTPUT_WIDTH}x{Config.OUTPUT_HEIGHT} 的图像...")

    if Config.SHOW_GENERATION:
        grid = run_wfc_with_display(
            Config.OUTPUT_WIDTH,
            Config.OUTPUT_HEIGHT,
            all_tiles,
            rules,
            Config.DISPLAY,
            Config.SCALE,
            max_attempts=20
        )
    else:
        grid = run_wfc(
            Config.OUTPUT_WIDTH,
            Config.OUTPUT_HEIGHT,
            all_tiles,
            rules,
            max_attempts=20
        )

    # 检查生成是否成功
    if grid is None:
        print("\n❌ 生成失败！尝试调整参数：")
        print("   - 减小输出尺寸")
        print("   - 增大瓦片尺寸")
        print("   - 启用旋转/翻转")
        return

    # ========================================================================
    # 第6步: 渲染最终图像
    # ========================================================================
    print("\n【第6步】渲染最终图像...")
    rendered = render_grid(grid, all_tiles, Config.SCALE)

    # 保存图像
    if Config.SAVE_OUTPUT:
        # 确保输出目录存在
        os.makedirs(Config.OUTPUT_DIR, exist_ok=True)

        # 生成包含设置信息的文件名
        filename = generate_output_filename()
        output_path = os.path.join(Config.OUTPUT_DIR, filename)

        pygame.image.save(rendered, output_path)
        print(f"         图像已保存到: {output_path}")

    # ========================================================================
    # 显示总结信息
    # ========================================================================
    elapsed = time.time() - start_time
    print(f"\n{'=' * 60}")
    print(f"✓ 生成完成！总耗时: {elapsed:.2f}秒")
    print(f"{'=' * 60}")

    # ========================================================================
    # 显示循环（等待用户关闭窗口）
    # ========================================================================
    pygame.display.set_caption("Wave Function Collapse - 完成！按ESC退出")
    print("\n窗口已打开，按ESC或关闭窗口退出...")

    clock = pygame.time.Clock()
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False

        Config.DISPLAY.fill((255, 255, 255))
        Config.DISPLAY.blit(rendered, (0, 0))
        pygame.display.flip()
        clock.tick(30)

    pygame.quit()
    print("\n程序结束")


# ============================================================================
# 程序入口
# ============================================================================

if __name__ == "__main__":
    main()