import cv2
import numpy as np
from PIL import Image


def extract_color_region_multiple(original_image, mask_image, target_colors, margin=10):
    """
    支持多种颜色的版本：提取指定颜色区域，白色背景

    Args:
        original_image: 原始图片 PIL Image 对象
        mask_image: 色块图 PIL Image 对象
        target_colors: 目标颜色列表，例如 [[192, 0, 0], [0, 192, 0], [0, 0, 192]] (红绿蓝)
        margin: 裁剪边距，默认 10

    Returns:
        PIL.Image: 提取的颜色区域图片，失败返回 None
    """

    if original_image is None:
        print("原始图片为空")
        return None

    if mask_image is None:
        print("掩码图片为空")
        return None

    # 将PIL Image转换为numpy数组
    original_rgb = np.array(original_image.convert('RGB'))
    mask_rgb = np.array(mask_image.convert('RGB'))

    # 创建多个目标颜色的组合掩码
    combined_mask = np.zeros(mask_rgb.shape[:2], dtype=np.uint8)

    found_colors = []
    for i, target_color in enumerate(target_colors):
        target_color = np.array(target_color)
        # 为每个颜色创建掩码
        color_mask = np.all(mask_rgb == target_color, axis=2).astype(np.uint8) * 255

        # 统计该颜色的像素数量
        color_pixels = np.sum(color_mask > 0)
        if color_pixels > 0:
            found_colors.append((target_color.tolist(), color_pixels))
            # 将该颜色的掩码添加到组合掩码中
            combined_mask = cv2.bitwise_or(combined_mask, color_mask)

    if len(found_colors) == 0:
        print(f"未找到任何目标颜色区域！")
        print(f"目标颜色列表: {[color for color in target_colors]}")
        return None

    # 形态学处理
    kernel = np.ones((3, 3), np.uint8)
    combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel)
    combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)

    # 调整掩码尺寸
    if combined_mask.shape != original_rgb.shape[:2]:
        combined_mask = cv2.resize(combined_mask, (original_rgb.shape[1], original_rgb.shape[0]))

    # 找边界框并裁剪
    coords = np.column_stack(np.where(combined_mask > 0))
    if len(coords) == 0:
        print(f"处理后未找到有效区域！")
        return None

    y_min, x_min = coords.min(axis=0)
    y_max, x_max = coords.max(axis=0)

    # 添加边距
    y_min = max(0, y_min - margin)
    x_min = max(0, x_min - margin)
    y_max = min(original_rgb.shape[0], y_max + margin)
    x_max = min(original_rgb.shape[1], x_max + margin)

    # 裁剪
    cropped_original = original_rgb[y_min:y_max, x_min:x_max]
    cropped_mask = combined_mask[y_min:y_max, x_min:x_max]

    # 创建白色背景结果
    result = np.full_like(cropped_original, 255, dtype=np.uint8)
    result[cropped_mask > 0] = cropped_original[cropped_mask > 0]

    # 统计信息
    mask_pixels = np.sum(cropped_mask > 0)
    total_pixels = cropped_mask.shape[0] * cropped_mask.shape[1]

    print(f"提取完成")
    print(f"找到的目标颜色:")
    for color, pixels in found_colors:
        print(f"  RGB{color}: {pixels} 像素")
    print(f"裁剪区域尺寸: {result.shape[1]} x {result.shape[0]}")
    print(f"总目标颜色像素数量: {mask_pixels}")
    print(f"占裁剪区域的比例: {mask_pixels / total_pixels * 100:.2f}%")

    # 转换为PIL Image并返回
    pil_image = Image.fromarray(result)
    return pil_image


# 使用示例
if __name__ == "__main__":
    # 从文件加载PIL Image
    original_image_path = r"D:\work\image_search\img1.jpg"
    mask_image_path = r"D:\work\image_search\output.png"

    original_image = Image.open(original_image_path)
    mask_image = Image.open(mask_image_path)

    # 示例1: 提取多种颜色区域
    print("=== 提取多种颜色区域 ===")
    target_colors = [
        [192, 0, 0],  # 红色
        [0, 192, 0],  # 绿色
        [0, 0, 192],  # 蓝色
        [192, 192, 0],  # 黄色
    ]

    result_image_multiple = extract_color_region_multiple(
        original_image,
        mask_image,
        target_colors=target_colors
    )

    if result_image_multiple:
        result_image_multiple.save(r"D:\work\image_search\extracted_multiple.png")
        print("多颜色区域提取成功并已保存")
    else:
        print("多颜色区域提取失败")

    # 示例2: 只提取红色和绿色
    print("\n=== 提取红色和绿色区域 ===")
    red_green_colors = [
        [192, 0, 0],  # 红色
        [0, 192, 0],  # 绿色
    ]

    result_image_rg = extract_color_region_multiple(
        original_image,
        mask_image,
        target_colors=red_green_colors
    )

    if result_image_rg:
        result_image_rg.save(r"D:\work\image_search\extracted_red_green.png")
        print("红绿区域提取成功并已保存")
    else:
        print("红绿区域提取失败")