import os
from PIL import Image
import numpy as np


def get_edge_opaque_pixels_with_limit_filtered_colors(image_path):
    """
    提取图片中每一行和每一列最左、最右、最上、最下的完全不透明像素（alpha == 255），
    并根据位置限制过滤，去掉位置和alpha通道。

    Args:
        image_path (str): 图片文件路径。

    Returns:
        dict: 包含四个边的颜色信息列表。
    """
    with Image.open(image_path).convert("RGBA") as img:
        width, height = img.size
        pixels = img.load()
        limit_left = width / 3
        limit_right = 2 * width / 3
        limit_top = height / 3
        limit_bottom = 2 * height / 3

        left_colors = []
        right_colors = []
        top_colors = []
        bottom_colors = []

        # 左边
        for y in range(height):
            for x in range(width):
                r, g, b, a = pixels[x, y]
                if a == 255 and x <= limit_left:
                    left_colors.append((r, g, b))
                    break

        # 右边
        for y in range(height):
            for x in range(width - 1, -1, -1):
                r, g, b, a = pixels[x, y]
                if a == 255 and x >= limit_right:
                    right_colors.append((r, g, b))
                    break

        # 上边
        for x in range(width):
            for y in range(height):
                r, g, b, a = pixels[x, y]
                if a == 255 and y <= limit_top:
                    top_colors.append((r, g, b))
                    break

        # 下边
        for x in range(width):
            for y in range(height - 1, -1, -1):
                r, g, b, a = pixels[x, y]
                if a == 255 and y >= limit_bottom:
                    bottom_colors.append((r, g, b))
                    break

    return {
        "left": left_colors,
        "right": right_colors,
        "top": top_colors,
        "bottom": bottom_colors
    }


def get_edge_opaque_pixels_with_limit_filtered_colors_numpy(image_path):
    """
    使用NumPy提取图片中每一行和每一列最左、最右、最上、最下的完全不透明像素（alpha == 255），
    并根据位置限制过滤，去掉位置和alpha通道。

    Args:
        image_path (str): 图片文件路径。

    Returns:
        dict: 包含四个边的颜色信息列表。
    """
    with Image.open(image_path).convert("RGBA") as img:
        data = np.array(img)
        height, width, _ = data.shape
        limit_left = width / 3
        limit_right = 2 * width / 3
        limit_top = height / 3
        limit_bottom = 2 * height / 3

        left_colors = []
        right_colors = []
        top_colors = []
        bottom_colors = []

        # 左边
        for y in range(height):
            row = data[y, :, :]
            opaque_indices = np.where(row[:, 3] == 255)[0]
            valid_indices = opaque_indices[opaque_indices <= limit_left]
            if valid_indices.size > 0:
                x = valid_indices[0]
                r, g, b, _ = row[x]
                left_colors.append((int(r), int(g), int(b)))

        # 右边
        for y in range(height):
            row = data[y, :, :]
            opaque_indices = np.where(row[:, 3] == 255)[0]
            valid_indices = opaque_indices[opaque_indices >= limit_right]
            if valid_indices.size > 0:
                x = valid_indices[-1]
                r, g, b, _ = row[x]
                right_colors.append((int(r), int(g), int(b)))

        # 上边
        for x in range(width):
            column = data[:, x, :]
            opaque_indices = np.where(column[:, 3] == 255)[0]
            valid_indices = opaque_indices[opaque_indices <= limit_top]
            if valid_indices.size > 0:
                y = valid_indices[0]
                r, g, b, _ = column[y]
                top_colors.append((int(r), int(g), int(b)))

        # 下边
        for x in range(width):
            column = data[:, x, :]
            opaque_indices = np.where(column[:, 3] == 255)[0]
            valid_indices = opaque_indices[opaque_indices >= limit_bottom]
            if valid_indices.size > 0:
                y = valid_indices[-1]
                r, g, b, _ = column[y]
                bottom_colors.append((int(r), int(g), int(b)))

    return {
        "left": left_colors,
        "right": right_colors,
        "top": top_colors,
        "bottom": bottom_colors
    }


def process_directory_filtered(directory_path, use_numpy=False):
    """
    遍历指定文件夹下所有图片文件，提取每张图片的边缘不透明像素颜色信息。

    Args:
        directory_path (str): 文件夹路径。
        use_numpy (bool): 是否使用NumPy加速处理。

    Returns:
        dict: 每个文件对应的四个方向的颜色信息列表。
    """
    supported_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff')
    result_map = {}

    for filename in os.listdir(directory_path):
        if filename.lower().endswith(supported_extensions):
            file_path = os.path.join(directory_path, filename)
            if use_numpy:
                edge_colors = get_edge_opaque_pixels_with_limit_filtered_colors_numpy(file_path)
            else:
                edge_colors = get_edge_opaque_pixels_with_limit_filtered_colors(file_path)
            result_map[filename] = edge_colors

    return result_map


if __name__ == "__main__":
    directory_path = "output4"  # 替换为你的文件夹路径
    use_numpy = True  # 设置为True使用NumPy优化，False使用PIL
    result = process_directory_filtered(directory_path, use_numpy)

    for filename, edges in result.items():
        print(f"文件: {filename}")

        # 左边
        print("  左边的最左不透明像素颜色:")
        for color in edges["left"]:
            print(f"    颜色: {color}")

        # 右边
        print("  右边的最右不透明像素颜色:")
        for color in edges["right"]:
            print(f"    颜色: {color}")

        # 上边
        print("  上边的最上不透明像素颜色:")
        for color in edges["top"]:
            print(f"    颜色: {color}")

        # 下边
        print("  下边的最下不透明像素颜色:")
        for color in edges["bottom"]:
            print(f"    颜色: {color}")

        print("\n")