import numpy as np
from PIL import Image
from skimage import exposure
from matplotlib import colors

def match_image_hsv(source_path, reference_path, output_path):
    # 1. 读取图片并归一化到 0-1 (matplotlib 的 hsv 转换需要 0-1)
    src_rgb = np.array(Image.open(source_path)) / 255.0
    ref_rgb = np.array(Image.open(reference_path)) / 255.0

    # 2. 将图片从 RGB 转换为 HSV
    src_hsv = colors.rgb_to_hsv(src_rgb)
    ref_hsv = colors.rgb_to_hsv(ref_rgb)

    # 3. 分离 HSV 通道
    # src_h: 色调, src_s: 饱和度, src_v: 亮度
    s_h, s_s, s_v = src_hsv[:,:,0], src_hsv[:,:,1], src_hsv[:,:,2]
    r_h, r_s, r_v = ref_hsv[:,:,0], ref_hsv[:,:,1], ref_hsv[:,:,2]

    # 4. 对 V (亮度) 和 S (饱和度) 通道进行直方图匹配
    # 我们使用参考图的 V 和 S 分布来调整源图
    matched_h = exposure.match_histograms(s_h, r_h)
    matched_v = exposure.match_histograms(s_v, r_v)
    matched_s = exposure.match_histograms(s_s, r_s)

    # 5. 合并通道
    # 使用原始的 H (色调)，加上匹配后的 S 和 V
    # V1 不调整H版本
    matched_hsv = np.stack([s_h, matched_s, matched_v], axis=2)
    # V2 调整H版本
    # matched_hsv = np.stack([matched_h, matched_s, matched_v], axis=2)

    # 6. 转换回 RGB 并保存
    # 转换回 RGB 后需要 clip 到 0-1 范围，防止数值溢出
    matched_rgb = np.clip(colors.hsv_to_rgb(matched_hsv), 0, 1)
    # 将 0-1 转换回 0-255 的整数并保存
    result_image = Image.fromarray((matched_rgb * 255).astype(np.uint8))
    result_image.save(output_path)
    print(f"HSV 处理完成，图片已保存至: {output_path}")

def match_image_appearance(source_path, reference_path, output_path):
    # 1. 读取图片
    # source: 需要调整的图片 (第一张, 偏暗)
    # reference: 目标风格图片 (第二张, 正常)
    src = np.array(Image.open(source_path))
    ref = np.array(Image.open(reference_path))

    # 2. 进行直方图匹配
    # channel_axis=-1 表示对 RGB 每个通道分别进行匹配
    matched = exposure.match_histograms(src, ref, channel_axis=-1)

    # 3. 保存结果
    result_image = Image.fromarray(matched.astype(np.uint8))
    result_image.save(output_path)
    print(f"处理完成，图片已保存至: {output_path}")

# 使用示例
source_file = "./去雾图像-北航合作/2025-07-02_084220_VID002.mp4_20251027_001308.661.png"
reference_file = "./去雾图像-北航合作-Result_Baidu/2025-07-02_084220_VID002.mp4_20251027_001308.661.png"
# V1
# match_image_appearance(source_file, reference_file, "adjusted_image_rgb.png")
# V2
match_image_hsv(source_file, reference_file, "adjusted_image_hsv.png")