import os
from PIL import Image
import numpy as np

from .common import calculate_depth_mean_pil
from tqdm import tqdm 

def depth_apply_mask(image_path):
    """
    基于掩码图像对另一张图像进行掩码处理，背景区域设为黑色

    Args:
        image_path (str): 输入图像的路径
        mask_path (str): 掩码图像的路径
        output_path (str): 保存处理后图像的路径
    """
    # 解析各文件名
    image_name = os.path.splitext(os.path.basename(image_path))[0].split('-')[-1]
    mask_path = f"data/mask_image/mask-{image_name}.png"
    output_path = f"data/mask_depth_image/mask-depth-{image_name}.png"

    # 读取图像和掩码图像
    image = Image.open(image_path)  # 假设已经是单通道
    mask = Image.open(mask_path)    # 假设已经是单通道

    # 将图像和掩码转换为 NumPy 数组
    image_array = np.array(image)
    mask_array = np.array(mask)

    # 创建一个全黑的图像
    masked_image_array = np.zeros_like(image_array)

    # 应用掩码：保留掩码区域的像素，背景设为黑色
    masked_image_array[mask_array > 0] = image_array[mask_array > 0]

    # 将处理后的数组转换回图像
    masked_image = Image.fromarray(masked_image_array)

    # 保存处理后的图像
    masked_image.save(output_path)
    return masked_image

def depth_apply_mask_pilimg(image, mask, image_path):
    """
    基于掩码图像对另一张图像进行掩码处理，背景区域设为黑色

    Args:
        image_path (str): 输入图像的路径
        mask_path (str): 掩码图像的路径
        output_path (str): 保存处理后图像的路径
    """
    # # 解析各文件名
    image_name = os.path.splitext(os.path.basename(image_path))[0]
    # mask_path = f"data/mask_image/mask-{image_name}.png"
    output_path = f"data/mask_depth_image/mask-depth-{image_name}.png"

    # # 读取图像和掩码图像
    # image = Image.open(image_path)  # 假设已经是单通道
    # mask = Image.open(mask_path)    # 假设已经是单通道

    # 将图像和掩码转换为 NumPy 数组
    image_array = np.array(image)
    mask_array = np.array(mask)

    # 创建一个全黑的图像
    masked_image_array = np.zeros_like(image_array)

    # 应用掩码：保留掩码区域的像素，背景设为黑色
    masked_image_array[mask_array > 0] = image_array[mask_array > 0]

    # 将处理后的数组转换回图像
    masked_image = Image.fromarray(masked_image_array)

    # 保存处理后的图像
    masked_image.save(output_path)
    return masked_image

def depth_apply_mask_pilimg_list(images, masks, output_paths):
    """
    对图像列表和掩码列表应用掩码处理，并保存处理后的图像

    Args:
        images (list): 输入图像列表
        masks (list): 掩码图像列表
        output_paths (list): 保存处理后图像的路径列表
    """
    masked_depths = []
    depth_means = []
    for image, mask, output_path in tqdm(zip(images, masks, output_paths)):
        masked_image = depth_apply_mask_pilimg(image, mask, output_path)
        depth_mean = calculate_depth_mean_pil(masked_image)
        masked_depths.append(masked_image)
        depth_means.append(depth_mean)

    return masked_depths, depth_means




# 使用示例
if __name__ == "__main__":
    input_image_path = "depth_image.png"  # 输入图像路径
    mask_image_path = "rmbgmask_image.png"      # 掩码图像路径
    output_image_path = "masked_output_image11.png"

    apply_mask(input_image_path)