"""colorfit.py — 局部（低频）光照与色彩迁移：把「原始脸」的明暗分布搬到「换出来的脸」上，只动低频、保留换脸细节。

为什么需要（§11.30，user 截图：DFM 整脸替换后是一块均匀受光的浅色平面，和有阴影的脖子接不上）：
全局 LAB 匹配只能对齐均值与对比度，搬不了"额头亮、下颌暗、左亮右暗"这种空间分布；FaceFusion 给 deep_swapper 用的
`vision.match_frame_color` 就是这类多尺度局部匹配——它在 16×16 粗网格上取（原始−换脸）的差再插值回全图叠加，
后面两级更细的网格只是对同一个平滑场重采样，实际起作用的就是粗网格这一级。这里做等价实现并加遮罩加权：
网格内只用脸像素求差（头发/背景不参与），几乎没有脸像素的网格不迁移。
"""
import cv2
import numpy as np


def _skin_weight(img, sigma, thresh=22.0):
    """"皮肤权重"：比邻域明显更暗的像素（眉毛、胡子、鼻孔、发际）权重→0。
    🔴 为什么要它（user 2026-09-07：开局部光照后胡子眉毛变少）：光照场是从（原始−换脸）的粗网格差估出来的，
    模型脸上有胡子/浓眉而你的脸上没有时，那些格子的差是"亮−暗"的正值，叠上去等于把模型的胡子眉毛漂白。
    光照本来就该从皮肤上估，把两张图里的深色结构都排除在统计之外，场只反映皮肤明暗；叠加时仍作用于整脸。"""
    L = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(np.float32)
    base = cv2.GaussianBlur(L, (0, 0), sigma)
    return 1.0 - np.clip((base - L) / float(thresh), 0.0, 1.0)


def local_color_transfer(orig, swapped, mask=None, strength=1.0, cells=(16,), min_cover=0.05, skin_only=True):
    """orig / swapped：同尺寸 uint8 BGR；mask：float32 [0,1] 同尺寸（None = 全 1）；strength 0~1；
    cells：粗网格边长（格数）序列，逐级细化（默认只有 16，与 FaceFusion 等效）；
    skin_only：光照场只从皮肤像素估计（排除眉毛/胡子等深色结构，见 _skin_weight）。返回 uint8 BGR。"""
    if strength <= 0 or orig is None or swapped is None or orig.shape != swapped.shape:
        return swapped
    h, w = swapped.shape[:2]
    o = orig.astype(np.float32)
    out = swapped.astype(np.float32)
    m = np.ones((h, w), np.float32) if mask is None else np.clip(mask.astype(np.float32), 0, 1)
    m3 = m[..., None]
    wgt = m
    if skin_only:
        sigma = max(2.0, min(h, w) / float(cells[0]))
        wgt = m * _skin_weight(orig, sigma) * _skin_weight(swapped, sigma)
    w3 = wgt[..., None]
    for c in cells:
        size = (int(max(2, min(c, w))), int(max(2, min(c, h))))
        num = cv2.resize((o - out) * w3, size, interpolation=cv2.INTER_AREA)   # 皮肤加权的差
        den = cv2.resize(wgt, size, interpolation=cv2.INTER_AREA)
        diff = num / np.maximum(den, 1e-3)[..., None]
        diff[den < min_cover] = 0.0                                              # 网格里几乎没有皮肤像素 → 不迁移
        up = cv2.resize(diff, (w, h), interpolation=cv2.INTER_CUBIC)
        out = out + up * float(strength) * m3                                    # 叠加仍作用于整个脸区
    return np.clip(out, 0, 255).astype(np.uint8)


if __name__ == "__main__":       # 自测：给"换脸"加一条左亮右暗的光照，看迁移后是否复现原始的明暗比
    import sys
    sys.stdout.reconfigure(encoding="utf-8")
    base = np.full((256, 256, 3), 150, np.uint8)
    grad = np.linspace(1.3, 0.6, 256, dtype=np.float32)[None, :, None]
    orig = np.clip(base.astype(np.float32) * grad, 0, 255).astype(np.uint8)   # 原始：左亮右暗
    out = local_color_transfer(orig, base, None, 1.0)
    lr = lambda a: float(a[:, :64].mean() / a[:, -64:].mean())
    print("左右亮度比 原始 %.2f  换脸(均匀) %.2f  迁移后 %.2f" % (lr(orig), lr(base), lr(out)))
    assert abs(lr(out) - lr(orig)) < 0.08, "迁移后没有复现原始的明暗分布"
    print("OK")
