#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
enhancer.py — 换脸后人脸修复（GPEN / GFPGAN），修 §11 的换脸糊化。

背景（实测）：inswapper 输出只有 128×128，贴回 ~240px 的脸要放大 ~2×→糊。实测换脸后
脸部 Laplacian 锐度只有原脸的 25%（8.1 vs 90.8，对齐 512 空间）。身份转移是好的，就是软。

修法：换脸后接一个人脸修复器，把糊掉的脸拉回清晰。实测（RTX 5090，对齐 512 空间）：
  · gpen_bfr_256  锐度 8.1→20.7（2.6×）  4.92ms  身份忠实、自然         ← 默认（实时够快）
  · gfpgan_1.4    锐度 8.1→114.7（14×）  13.33ms 最锐但轻微理想化改身份   ← 质量优先可切
  · gpen_bfr_512  25ms 太慢，不用

🔴 两个必须做对的点：
  1. **ROI 回贴**：修复器在 512 对齐空间产出整张脸，回贴必须只在人脸包围盒内 warp+blend，
     绝不能全帧 warp（§3 的教训：全帧回贴 65ms）。
  2. **遮挡遮罩要并进回贴 mask**：修复器会把挡在脸前的手也「修复」成皮肤，等于把 §11.2
     刚修好的遮挡又破坏掉。所以回贴 mask = 椭圆羽化 × 遮挡遮罩（occ），手区不被修复覆盖。
"""
import os

import cv2
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
from validated_impl import build_feather_mask  # noqa: E402

# FaceFusion 修复器 ONNX 都在 models-3.0.0（3.1.0 没有，返回 Entry not found）
_B = "https://huggingface.co/facefusion/models-3.0.0/resolve/main/"
ENH = {
    "codeformer":   (512, _B + "codeformer.onnx"),      # ← 默认：真实感最佳（SOTA，保真+自然）
    "gpen_bfr_512": (512, _B + "gpen_bfr_512.onnx"),    # 最大细节的备选
    "gpen_bfr_256": (256, _B + "gpen_bfr_256.onnx"),    # 最快、最省，偏软
    "gfpgan_1.4":   (512, _B + "gfpgan_1.4.onnx"),      # 过度理想化，不推荐真实场景
}


def _log(*a):
    print("[enhancer]", *a, flush=True)


def ensure_enhancer(model, assets_dir):
    """下载修复器模型（实例临时、销毁即没，下载逻辑常驻代码）。返回 (路径, 输入尺寸)。"""
    if model not in ENH:
        raise ValueError(f"未知修复器 {model}，可选 {list(ENH)}")
    size, url = ENH[model]
    path = os.path.join(assets_dir, model + ".onnx")
    if not (os.path.exists(path) and os.path.getsize(path) > 1_000_000):
        import urllib.request
        _log(f"下载 {model} …")
        urllib.request.urlretrieve(url, path)
        _log(f"{model} 就绪 {os.path.getsize(path)/1e6:.0f}MB")
    return path, size


class FaceEnhancer:
    """换脸后人脸修复。__call__(img, face, occ=None) 原地把 img 里的脸修清晰后返回。

    align_size=512：GPEN/GFPGAN 都在 512 对齐空间工作（gpen_256 内部再缩到 256 推理）。
    blend<1：与修复前的脸线性混合，压一压修复器的「理想化」倾向、保住身份（默认 0.8）。
    """

    def __init__(self, model_path, in_size, providers=None, blend=0.8, align_size=512,
                 weight=0.8):
        import onnxruntime as ort
        providers = providers or ["CUDAExecutionProvider", "CPUExecutionProvider"]
        self.sess = ort.InferenceSession(model_path, providers=providers)
        if "CUDAExecutionProvider" not in self.sess.get_providers():
            _log("⚠️ 修复器未跑在 CUDA 上:", self.sess.get_providers())
        self.inputs = self.sess.get_inputs()          # codeformer 有第二个 fidelity 权重输入
        self.in_size = in_size
        self.align = align_size
        self.blend = blend
        # codeformer 的 fidelity 权重：0=最大修复(可能失真) / 1=最大保真。0.8 偏保身份。
        # 对单输入模型（gpen/gfpgan）无效，忽略。
        self.weight = weight
        self.pad = 12
        # 椭圆羽化（512 对齐空间），只把脸内修复结果贴回、边界藏进明暗过渡，避免锐度断层缝
        self.feather = build_feather_mask(align_size, shape="ellipse")
        from insightface.utils import face_align
        self.fa = face_align

    def _run(self, aligned_bgr):
        a = cv2.resize(aligned_bgr, (self.in_size, self.in_size)) if self.in_size != self.align else aligned_bgr
        x = (a[:, :, ::-1].astype(np.float32) / 255.0 - 0.5) / 0.5   # RGB [-1,1]
        x = x.transpose(2, 0, 1)[None]
        feeds = {}
        for inp in self.inputs:
            if len(inp.shape) == 4:                 # 图像输入
                feeds[inp.name] = x
            else:                                    # codeformer 的 fidelity 权重标量
                dt = np.float64 if "double" in inp.type else np.float32
                feeds[inp.name] = np.array([self.weight], dt)
        y = np.squeeze(self.sess.run(None, feeds)[0])
        if y.shape[0] == 3:
            y = y.transpose(1, 2, 0)
        y = np.clip((np.clip(y, -1, 1) + 1) / 2 * 255, 0, 255)[:, :, ::-1]   # → BGR [0,255]
        y = np.clip(y, 0, 255).astype(np.uint8)
        return cv2.resize(y, (self.align, self.align)) if y.shape[0] != self.align else y

    def __call__(self, img, face, occ=None):
        S = self.align
        al, M = self.fa.norm_crop2(img, face.kps, S)
        enh = self._run(al)
        # 🔴 亮度保持（2026-08-17，user 真机「摄像头暗时换脸后人更黑」归因后加）：
        #    修复器在暗脸上有系统性调子漂移（gpen_bfr_256 本地实测：×0.4 暗档 L −2.0，
        #    与链路 −3.5 叠加成肉眼可见的"更黑"）。把输出的 L 均值/标准差拉回**它自己
        #    输入**的统计——这是对自身的恒等约束（修锐度不该改曝光），不是 §11.12 被禁的
        #    场景级统计对齐（那个是拿房间均值去改虚拟背景）。只动 L、不碰 ab（肤色仍由
        #    swap 层的 color_match 负责）；gain 钳制 [0.85,1.18] 防极端帧。
        lab_e = cv2.cvtColor(enh, cv2.COLOR_BGR2LAB).astype(np.float32)
        lab_a = cv2.cvtColor(al, cv2.COLOR_BGR2LAB).astype(np.float32)
        ctr = slice(S // 4, S * 3 // 4)                    # 中心半幅=脸内部，避开黑边/背景
        Le, La = lab_e[ctr, ctr, 0], lab_a[ctr, ctr, 0]
        g = float(np.clip(La.std() / max(Le.std(), 1e-6), 0.85, 1.18))
        b = float(La.mean() - g * Le.mean())
        lab_e[..., 0] = np.clip(lab_e[..., 0] * g + b, 0, 255)
        enh = cv2.cvtColor(lab_e.astype(np.uint8), cv2.COLOR_LAB2BGR)
        if self.blend < 1.0:
            enh = cv2.addWeighted(enh, self.blend, al, 1.0 - self.blend, 0)
        # 回贴 mask = 椭圆羽化 × 遮挡遮罩（occ 传入时；否则纯椭圆）
        mask = self.feather
        if occ is not None:
            o = cv2.resize(occ, (S, S)) if occ.shape[:2] != (S, S) else occ
            mask = self.feather * o                                   # 手区 →0，不被修复覆盖

        # 🔴 ROI 回贴：只在人脸包围盒内 warp+blend（全帧回贴是 §3 的 65ms 大坑）
        h, w = img.shape[:2]
        IM = cv2.invertAffineTransform(M)
        corners = np.array([[[0, 0], [S, 0], [S, S], [0, S]]], np.float32)
        d = cv2.transform(corners, IM).reshape(-1, 2)
        x0 = max(int(d[:, 0].min()) - self.pad, 0); x1 = min(int(d[:, 0].max()) + self.pad, w)
        y0 = max(int(d[:, 1].min()) - self.pad, 0); y1 = min(int(d[:, 1].max()) + self.pad, h)
        rw, rh = x1 - x0, y1 - y0
        if rw <= 2 or rh <= 2:
            return img
        IMr = IM.copy(); IMr[0, 2] -= x0; IMr[1, 2] -= y0
        fr = cv2.warpAffine(enh, IMr, (rw, rh), borderValue=0.0)
        mk = (cv2.warpAffine(mask, IMr, (rw, rh), borderValue=0.0) / 255.)[..., None]
        roi = img[y0:y1, x0:x1].astype(np.float32)
        img[y0:y1, x0:x1] = (mk * fr + (1 - mk) * roi).astype(np.uint8)
        return img
