# -*- coding: utf-8 -*-
"""
merge_fixed.py — 原图 + AI修改图 无缝融合工具
============================================
背景：ComfyUI inpainting 移除合影中的其他人物时，即使指定"非修改区域不得像素漂移"，
     全图清晰度仍不可避免损失。本工具把：
       - 原图的"未修改区域"（清晰度完好）
       - 修改图的"修改区域"（人物被干净移除）
     自动合并，过渡自然无痕迹。

用法：
  python merge_fixed.py <原图目录> <修改图目录> <输出目录> [--feather 25] [--threshold 20]

匹配规则：
  - 原图 IMG_0463.jpg（任意后缀 jpg/png/webp/jpeg）
  - 修改图 IMG_0463*.png（以原图文件名开头、.png 结尾、中间任意字符）
  - 一张原图可有多张修改图，每张都独立合并
  - 输出：{原图stem}_{修改图中间部分}.png

算法：
  1. 修改图 resize 到原图尺寸
  2. Lab 色彩空间差异图 → 高斯模糊去噪 → 阈值 → 闭运算+膨胀 = 修改区 mask
  3. 羽化混合：result = 原图*(1-alpha) + 修改图*alpha（alpha 由 mask 高斯模糊得到）
"""
import argparse
import os
import sys
from pathlib import Path

import cv2
import numpy as np

sys.stdout.reconfigure(encoding="utf-8", errors="replace")

# 复用 face_checker 的人脸检测（目标人物面部保护）
sys.path.insert(0, str(Path(__file__).resolve().parent))
import face_checker as fc

IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
_fc_det = None


def get_face_detector():
    global _fc_det
    if _fc_det is None:
        _fc_det = fc.FaceDetector()
    return _fc_det


def protect_faces(alpha, mod_rgb, expand=1.35, margin=40):
    """
    面部强制保护（保险）：检测修改图中残留的人脸（=目标人物，被移除者已被 inpaint 掉），
    把这些人脸区域的 alpha 置 0（100% 用原图像素，零改变）。
    通用规则（有限修改→原图）已覆盖大部分，此项兜底确保面部绝对不变。
    """
    det = get_face_detector()
    faces = fc._detect_faces_fast(mod_rgb, det)
    if not faces:
        return alpha
    for f in faces:
        x, y, w, h = f[0], f[1], f[2], f[3]
        cw, ch = w * expand, h * expand
        x0 = max(0, int(x - (cw - w) / 2) - margin)
        y0 = max(0, int(y - (ch - h) / 2) - margin)
        x1 = min(alpha.shape[1], int(x + w + (cw - w) / 2) + margin)
        y1 = min(alpha.shape[0], int(y + h + (ch - h) / 2) + margin)
        alpha[y0:y1, x0:x1] = 0.0
    return alpha


def build_alpha_from_mask(mask_rgb, feather=20):
    """
    从 ComfyUI inpaint 蒙版生成 alpha：白色区域 = 被修改区（100% 修改图），黑色 = 原图。
    蒙版是最准确的"修改区"标记（用户在 ComfyUI 画的就是要移除的人物），
    比 diff 检测可靠 100 倍——紧贴人物的残影完美解决。
    feather: 蒙版边缘羽化像素
    """
    gray = cv2.cvtColor(mask_rgb, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0
    # 蒙版白色→1（修改图），黑色→0（原图）
    alpha = gray
    if feather > 0:
        k = feather * 2 + 1
        alpha = cv2.GaussianBlur(alpha, (k, k), 0)
    return alpha


def find_mask(orig_path, mask_dir):
    """按文件名匹配蒙版：{stem}*mask*.png 或 {stem}*.png（取 mask 关键字优先）"""
    stem = Path(orig_path).stem
    cands = sorted(p for p in mask_dir.iterdir()
                   if p.suffix.lower() == ".png" and p.name.startswith(stem))
    # 优先带 mask 关键字的
    for p in cands:
        if "mask" in p.name.lower():
            return p
    return cands[0] if cands else None


def build_alpha(orig_rgb, mod_rgb, blur_kernel=15, threshold=4.0, min_area_ratio=0.005,
                softness=1.5, dilate=60, body_extend=2.2):
    """
    融合权重 alpha（0=原图像素, 1=修改图像素）——双图人脸差集定位被移除人物：
      1. 原图人脸 - 修改图人脸 = 被移除人物（合影中被清掉的人）
      2. 被移除人物区域（脸框 + 向下 body_extend 倍覆盖身体）= 修改区
      3. 目标人物（修改图残留的最大人脸）紧贴框 = 保留区（强制原图）
      4. diff 高值连通块（排除保留区）兜底其他修改
      5. 膨胀覆盖边缘残影
    紧贴合影场景：被移除人物与目标紧贴时，靠"人脸差集"精确定位，不靠 diff 幅度。
    """
    det = get_face_detector()
    orig_faces = fc._detect_faces_fast(orig_rgb, det)
    mod_faces = fc._detect_faces_fast(mod_rgb, det)
    h, w = orig_rgb.shape[:2]

    def matched(of, mfaces):
        ox, oy = of[0] + of[2] / 2, of[1] + of[3] / 2
        return any(abs(ox - (mf[0] + mf[2] / 2)) < of[2] * 0.8 and abs(oy - (mf[1] + mf[3] / 2)) < of[3] * 0.8
                   for mf in mfaces)

    removed_faces = [f for f in orig_faces if not matched(f, mod_faces)]
    # 过滤误检小脸（背景物体/花纹被当脸）：只保留足够大的人脸
    # 真实人物脸 ≥ 图面积 0.05% 或短边 ≥ 40px
    min_face_area = max(1, int(h * w * 0.0005))
    removed_faces = [f for f in removed_faces
                     if f[2] * f[3] >= min_face_area and min(f[2], f[3]) >= 40]

    # 被移除人物区（脸 + 身体向下扩展）
    mod_mask = np.zeros((h, w), np.uint8)
    for f in removed_faces:
        x, y, fw, fh = [int(v) for v in f[:4]]
        x0 = max(0, int(x - fw * 0.4)); x1 = min(w, int(x + fw * 1.4))
        y0 = max(0, int(y - fh * 0.3)); y1 = min(h, int(y + fh * body_extend))
        mod_mask[y0:y1, x0:x1] = 1

    # 保留区：目标人物最大脸，紧贴框（左右不扩，防侵入紧贴的被移除人物）
    keep_mask = np.zeros((h, w), np.uint8)
    if mod_faces:
        big = max(mod_faces, key=lambda f: f[2] * f[3])
        x, y, fw, fh = [int(v) for v in big[:4]]
        x0, x1 = max(0, x), min(w, x + fw)
        y0, y1 = max(0, int(y - fh * 0.35)), min(h, int(y + fh * 1.2))
        keep_mask[y0:y1, x0:x1] = 1

    # diff 兜底连通块（排除保留区）。无被移除人物时用高阈值（保守，防降质误判）
    o_lab = cv2.cvtColor(orig_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
    m_lab = cv2.cvtColor(mod_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
    diff_s = cv2.GaussianBlur(np.abs(o_lab - m_lab).mean(axis=2), (blur_kernel, blur_kernel), 0)
    eff_threshold = threshold if removed_faces else max(threshold, 15.0)
    base = ((diff_s > eff_threshold) & ~keep_mask.astype(bool)).astype(np.uint8)
    num, labels, stats, _ = cv2.connectedComponentsWithStats(base, 8)
    min_area = max(1, int(min_area_ratio * h * w))
    filtered = np.zeros_like(base)
    for i in range(1, num):
        if stats[i, cv2.CC_STAT_AREA] >= min_area:
            filtered[labels == i] = 1

    core = np.maximum(filtered, mod_mask)
    if dilate > 0:
        core = cv2.dilate(core, np.ones((dilate, dilate), np.uint8))
    alpha = core.astype(np.float32)
    if softness > 0:
        k = int(softness * 6) * 2 + 1
        alpha = cv2.GaussianBlur(alpha, (k, k), 0)
    return alpha
    """
    融合权重 alpha（0=原图像素, 1=修改图像素）：
      - 固定阈值判定明显修改区（人物移除 diff 高，有限降质 diff 低，区分度 ~20x）
      - diff > 阈值 → alpha≈1（100% 修改图，修改区内部零残影）
      - diff < 阈值 → alpha≈0（100% 原图，目标人物/有限修改区零改变）
      - 修改区膨胀 dilate 像素（默认 75）：覆盖人物边缘 10-40px 的浅残影带
        （残影带 diff 接近未修改区，无法用阈值检测，必须靠膨胀）
      - 陡 sigmoid 边缘过渡防硬边
    """
    o = cv2.cvtColor(orig_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
    m = cv2.cvtColor(mod_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
    diff = np.abs(o - m).mean(axis=2)
    diff_s = cv2.GaussianBlur(diff, (blur_kernel, blur_kernel), 0)

    alpha = 1.0 / (1.0 + np.exp(-(diff_s - threshold) / softness))

    # 修改区膨胀：覆盖人物边缘残影带
    if dilate > 0:
        kernel = np.ones((dilate, dilate), np.uint8)
        core = (alpha > 0.5).astype(np.uint8)
        core = cv2.dilate(core, kernel)
        alpha = np.maximum(alpha, core.astype(np.float32))
    return alpha


def seamless_merge(orig_rgb, mod_rgb, alpha):
    """
    连续 alpha 混合：result = orig*(1-alpha) + mod*alpha。
    alpha 已是平滑的浮点图（sigmoid 过渡），无需额外羽化。
    """
    a = alpha[..., None].astype(np.float32)
    result = orig_rgb.astype(np.float32) * (1.0 - a) + mod_rgb.astype(np.float32) * a
    return np.clip(result, 0, 255).astype(np.uint8)


def find_modified(orig_path, mod_dir):
    """按文件名匹配修改图：{stem}*.png"""
    stem = Path(orig_path).stem
    return sorted(p for p in mod_dir.iterdir()
                  if p.suffix.lower() == ".png" and p.name.startswith(stem))


def process(orig_dir, mod_dir, out_dir, alpha_threshold=8.0, softness=1.5, protect=True,
            mask_dir=None, feather=20, verbose=True):
    """
    alpha_threshold: 无蒙版时的 diff 阈值（默认 8；误判时调）
    softness: 过渡带宽度
    protect: 面部强制保护（无蒙版 fallback 时的保险）
    mask_dir: ComfyUI inpaint 蒙版目录（可选，有则优先用，蒙版白色=被移除区，最准确）
    feather: 蒙版边缘羽化
    """
    orig_dir, mod_dir, out_dir = Path(orig_dir), Path(mod_dir), Path(out_dir)
    if mask_dir:
        mask_dir = Path(mask_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    originals = sorted(p for p in orig_dir.iterdir()
                       if p.suffix.lower() in IMG_EXTS and not p.name.startswith("."))

    if not originals:
        print(f"[WARN] 原图目录没有图片: {orig_dir}")
        return

    total = 0
    for op in originals:
        mods = find_modified(op, mod_dir)
        if not mods:
            if verbose:
                print(f"[跳过] {op.name}: 无匹配修改图")
            continue

        from PIL import Image
        orig_img = Image.open(op).convert("RGB")
        ow, oh = orig_img.size
        orig_rgb = np.array(orig_img)

        # 尝试匹配蒙版
        mask_path = find_mask(op, mask_dir) if mask_dir else None

        for mp in mods:
            try:
                mod_img = Image.open(mp).convert("RGB")
                if mod_img.size != (ow, oh):
                    mod_img = mod_img.resize((ow, oh), Image.LANCZOS)
                mod_rgb = np.array(mod_img)

                if mask_path and mask_path.exists():
                    # 蒙版模式：最准确
                    mask_img = Image.open(mask_path).convert("RGB")
                    if mask_img.size != (ow, oh):
                        mask_img = mask_img.resize((ow, oh), Image.LANCZOS)
                    mask_rgb = np.array(mask_img)
                    alpha = build_alpha_from_mask(mask_rgb, feather=feather)
                    mode = f"蒙版模式({mask_path.name})"
                else:
                    # 双图人脸差集模式（自动定位被移除人物 + 保留目标脸）
                    alpha = build_alpha(orig_rgb, mod_rgb, threshold=alpha_threshold, softness=softness)
                    mode = "差集模式"

                area_ratio = float((alpha > 0.5).mean())
                if area_ratio < 0.0005:
                    if verbose:
                        print(f"[警告] {op.name} <- {mp.name}: 明显修改区仅 {area_ratio*100:.2f}%%，可能未检测到修改")
                    mask_area_note = f"明显修改区 {area_ratio*100:.2f}% (偏小?)"
                else:
                    mask_area_note = f"明显修改区 {area_ratio*100:.1f}%"

                result = seamless_merge(orig_rgb, mod_rgb, alpha)
                mid = mp.stem[len(Path(op).stem):] or "_mod"
                out_name = f"{Path(op).stem}{mid}.png"
                Image.fromarray(result).save(out_dir / out_name)
                total += 1
                if verbose:
                    print(f"[OK] {out_name} | {mode} | {mask_area_note}")
            except Exception as e:
                if verbose:
                    print(f"[失败] {op.name} <- {mp.name}: {e}")

    print(f"\n完成：共输出 {total} 张融合图 -> {out_dir}")


def main():
    ap = argparse.ArgumentParser(description="原图 + AI修改图 无缝融合工具（支持 ComfyUI 蒙版）")
    ap.add_argument("orig_dir", help="原图目录")
    ap.add_argument("mod_dir", help="修改图目录（文件名以原图名开头、.png 结尾）")
    ap.add_argument("out_dir", help="输出目录")
    ap.add_argument("--mask-dir", default=None,
                    help="ComfyUI inpaint 蒙版目录（可选，有则优先用；蒙版白色=被移除区，最准确解决紧贴残影）")
    ap.add_argument("--alpha-threshold", type=float, default=8.0,
                    help="无蒙版时的 diff 阈值（默认 8）")
    ap.add_argument("--softness", type=float, default=1.5, help="过渡带宽度（默认 1.5）")
    ap.add_argument("--feather", type=int, default=20, help="蒙版边缘羽化（默认 20）")
    ap.add_argument("--no-protect-face", action="store_true", help="关闭面部强制保护（无蒙版时的保险）")
    args = ap.parse_args()

    process(args.orig_dir, args.mod_dir, args.out_dir,
            alpha_threshold=args.alpha_threshold, softness=args.softness,
            protect=not args.no_protect_face, mask_dir=args.mask_dir, feather=args.feather)


if __name__ == "__main__":
    main()