#!/usr/bin/env python3
"""
证件照换头工具 v2 - 先擦除目标头部，再贴源头部
用法: python3 head_swap.py <源图> <目标图> <输出图>
"""
import sys
import os
import argparse
import cv2
import numpy as np
from insightface.app import FaceAnalysis


def init_app():
    app = FaceAnalysis(
        name="buffalo_l",
        root=os.path.join(os.path.dirname(os.path.abspath(__file__)), "models", "insightface"),
        providers=["CPUExecutionProvider"],
        allowed_modules=["detection", "landmark_2d_106"],
    )
    app.prepare(ctx_id=0, det_size=(640, 640))
    return app


def get_affine_matrix(src_kps, dst_kps):
    src = src_kps.astype(np.float32)
    dst = dst_kps.astype(np.float32)
    m, _ = cv2.estimateAffinePartial2D(src, dst, method=cv2.LMEDS)
    return m


def make_head_mask(kps, img_shape, scale=2.0):
    """基于关键点生成整颗头(含头发)的大椭圆mask"""
    h, w = img_shape[:2]
    cx = kps[:, 0].mean()
    cy = kps[:, 1].mean()
    face_w = kps[:, 0].max() - kps[:, 0].min()
    face_h = kps[:, 1].max() - kps[:, 1].min()

    # 椭圆半径
    rx = face_w * scale / 2
    ry = face_h * scale * 1.1 / 2  # 纵向稍大
    # 中心向上偏移(头顶头发区域)
    cy_shift = ry * 0.08

    mask = np.zeros((h, w), dtype=np.uint8)
    center = (int(cx), int(cy - cy_shift))
    axes = (int(rx), int(ry))
    cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)

    # 羽化边缘
    blur_k = max(31, int(min(rx, ry) * 0.4) | 1)
    mask = cv2.GaussianBlur(mask, (blur_k, blur_k), blur_k // 3)
    return mask


def swap_head(app, src_path, dst_path, out_path):
    src = cv2.imread(src_path)
    dst = cv2.imread(dst_path)
    if src is None:
        print(f"错误: 无法读取源图 {src_path}"); sys.exit(1)
    if dst is None:
        print(f"错误: 无法读取目标图 {dst_path}"); sys.exit(1)

    print(f"源图: {src_path} ({src.shape[1]}x{src.shape[0]})")
    print(f"目标: {dst_path} ({dst.shape[1]}x{dst.shape[0]})")

    src_faces = app.get(src)
    dst_faces = app.get(dst)

    if not src_faces:
        print("错误: 源图未检测到人脸"); sys.exit(1)
    if not dst_faces:
        print("错误: 目标图未检测到人脸"); sys.exit(1)

    src_face = src_faces[0]
    dst_face = dst_faces[0]
    src_kps = src_face.kps  # (5,2)
    dst_kps = dst_face.kps
    print(f"源人脸 bbox: {src_face.bbox.astype(int)}")
    print(f"目标人脸 bbox: {dst_face.bbox.astype(int)}")

    # 1) 仿射矩阵: 源关键点 → 目标关键点
    M = get_affine_matrix(src_kps, dst_kps)
    if M is None:
        print("错误: 无法计算仿射变换"); sys.exit(1)

    # 2) 生成源图的头部mask(含头发), 然后warp到目标图坐标系
    src_head_mask = make_head_mask(src_kps, src.shape, scale=2.0)

    # 3) 把源图warp到目标图尺寸
    h, w = dst.shape[:2]
    src_warped = cv2.warpAffine(src, M, (w, h),
                                 flags=cv2.INTER_LINEAR,
                                 borderMode=cv2.BORDER_CONSTANT,
                                 borderValue=(0, 0, 0))
    mask_warped = cv2.warpAffine(src_head_mask, M, (w, h),
                                  flags=cv2.INTER_LINEAR,
                                  borderMode=cv2.BORDER_CONSTANT,
                                  borderValue=0)

    # 4) 在目标图上擦除头部区域, 然后贴上源头部
    # 擦除mask = 源头部覆盖的区域(不透明部分)
    erase_mask = (mask_warped > 20).astype(np.uint8) * 255
    # 轻微膨胀确保完全覆盖
    kernel = np.ones((7, 7), np.uint8)
    erase_mask = cv2.dilate(erase_mask, kernel, iterations=2)

    # 用目标图背景填充擦除区域(简单方法: 用周围像素扩展)
    # 更好的方法: 用telea/inpaint算法填充
    erase_mask_binary = (erase_mask > 128).astype(np.uint8) * 255
    dst_erased = cv2.inpaint(dst, erase_mask_binary, inpaintRadius=10,
                              flags=cv2.INPAINT_TELEA)

    # 5) 把源头部(带羽化alpha)贴到擦除后的目标图上
    alpha = mask_warped.astype(np.float32) / 255.0
    alpha = alpha[:, :, np.newaxis]

    result = dst_erased.astype(np.float32) * (1 - alpha) + \
             src_warped.astype(np.float32) * alpha

    # 6) 简单色调匹配
    # 取目标图mask边缘外的像素均值 vs 源图mask内的均值
    outer_ring = cv2.dilate(erase_mask, np.ones((20, 20), np.uint8)) - erase_mask
    if outer_ring.sum() > 0:
        outer_px = dst[outer_ring > 0].astype(np.float32).mean(axis=0)
        inner_mask = (mask_warped > 128)
        if inner_mask.sum() > 0:
            inner_px = src_warped[inner_mask].astype(np.float32).mean(axis=0)
            # 用泊松混合的思路: 调整色偏
            color_shift = (outer_px - inner_px) * 0.35
            src_adjusted = np.clip(src_warped.astype(np.float32) + color_shift, 0, 255)
            result = dst_erased.astype(np.float32) * (1 - alpha) + \
                     src_adjusted * alpha

    result = np.clip(result, 0, 255).astype(np.uint8)

    cv2.imwrite(out_path, result, [cv2.IMWRITE_JPEG_QUALITY, 95])
    print(f"输出: {out_path}")
    print("完成!")


def main():
    parser = argparse.ArgumentParser(description="证件照换头v2(含头发)")
    parser.add_argument("src", help="源图(提供头的人)")
    parser.add_argument("dst", help="目标图(要换头的照片)")
    parser.add_argument("out", help="输出图")
    args = parser.parse_args()

    print("初始化 InsightFace...")
    app = init_app()
    print("模型加载完成")
    swap_head(app, args.src, args.dst, args.out)


if __name__ == "__main__":
    main()