"""expr_restorer.py — LivePortrait 表情迁移（FaceFusion `expression_restorer` 的逐步复刻，§11.13/§11.26 已离线实测：
张嘴幅度从换脸后的 0.63 恢复到 0.94，锐度不降，身份 ArcFace −0.03，12 帧连续抖动不升）。

做什么：用**原始帧**的脸当驱动，把它的 21×3 隐式关键点表情量按系数迁移到**换脸后**的脸上再重新生成——换脸模型
（inswapper/DFM）压掉的眨眼、嘴型、眉毛动作被补回来，身份不变。放在换脸之后、人脸修复之前。

口径逐字对照 facefusion master 2026-09-07：
  · 裁剪 arcface_128 模板 @512（= insightface norm_crop2(…,512)），512 → 256 INTER_AREA，BGR→RGB，/255，NCHW
  · feature_extractor(换脸裁剪) → feature_volume；motion_extractor 各跑一次：原始裁剪取 expression，换脸裁剪取 pitch/yaw/roll/scale/translation/expression/motion_points
  · R = scipy Rotation.from_euler('xyz',[pitch,yaw,roll],degrees=True)；tgt = tgt*f + tmp*(1−f)，再 clip(EXPRESSION_MIN, EXPRESSION_MAX)
  · pts = scale*(mp@R.T + expr) + translation；🔴 generator 输入 `source`=原始帧的点、`target`=换脸帧的点（反了嘴会更闭，已验）
  · 输出 512 → 盒状遮罩 blur 0.3 ∧ 遮挡遮罩 → ROI 回贴
模型：facefusion/models-3.0.0 的 live_portrait_{feature_extractor 3.4MB, motion_extractor 113MB, generator 222MB}.onnx，放 assets/，首次启用时下载。
"""
import os
import time
import urllib.request

import cv2
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
ASSETS = os.path.join(HERE, "assets")
HF = "https://huggingface.co/facefusion/models-3.0.0/resolve/main/"
MODELS = {"feat": "live_portrait_feature_extractor.onnx",
          "motion": "live_portrait_motion_extractor.onnx",
          "gen": "live_portrait_generator.onnx"}
MIN_SIZE = {"feat": 1_000_000, "motion": 50_000_000, "gen": 100_000_000}
UPPER = [1, 2, 6, 10, 11, 12, 13, 15, 16]      # FaceFusion restrict_expression_areas 的索引组
LOWER = [3, 7, 14, 17, 18, 19, 20]
FIXED = [0, 4, 5, 8, 9]
EXPRESSION_MIN = np.array([[
    [-2.88067125e-02, -8.12731311e-02, -1.70541159e-03], [-4.88598682e-02, -3.32196616e-02, -1.67431499e-04],
    [-6.75425082e-02, -4.28681746e-02, -1.98950816e-04], [-7.23103955e-02, -3.28503326e-02, -7.31324719e-04],
    [-3.87073644e-02, -6.01546466e-02, -5.50269964e-04], [-6.38048723e-02, -2.23840728e-01, -7.13261834e-04],
    [-3.02710701e-02, -3.93195450e-02, -8.24086510e-06], [-2.95799859e-02, -5.39318882e-02, -1.74219604e-04],
    [-2.92359516e-02, -1.53050944e-02, -6.30460854e-05], [-5.56493877e-03, -2.34344602e-02, -1.26858242e-04],
    [-4.37593013e-02, -2.77768299e-02, -2.70503685e-02], [-1.76926646e-02, -1.91676542e-02, -1.15090821e-04],
    [-8.34268332e-03, -3.99775570e-03, -3.27481248e-05], [-3.40162888e-02, -2.81868968e-02, -1.96679524e-04],
    [-2.91855410e-02, -3.97511162e-02, -2.81230678e-05], [-1.50395725e-02, -2.49494594e-02, -9.42573533e-05],
    [-1.67938769e-02, -2.00953931e-02, -4.00750607e-04], [-1.86435618e-02, -2.48535164e-02, -2.74416432e-02],
    [-4.61211195e-03, -1.21660791e-02, -2.93173041e-04], [-4.10017073e-02, -7.43824020e-02, -4.42762971e-02],
    [-1.90370996e-02, -3.74363363e-02, -1.34740388e-02]]], np.float32)
EXPRESSION_MAX = np.array([[
    [4.46682945e-02, 7.08772913e-02, 4.08344204e-04], [2.14308221e-02, 6.15894832e-02, 4.85319615e-05],
    [3.02363783e-02, 4.45043296e-02, 1.28298725e-05], [3.05869691e-02, 3.79812494e-02, 6.57040102e-04],
    [4.45670523e-02, 3.97259220e-02, 7.10966764e-04], [9.43699256e-02, 9.85926315e-02, 2.02551950e-04],
    [1.61131397e-02, 2.92906128e-02, 3.44733417e-06], [5.23825921e-02, 1.07065082e-01, 6.61510974e-04],
    [2.85718683e-03, 8.32320191e-03, 2.39314613e-04], [2.57947259e-02, 1.60935968e-02, 2.41853559e-05],
    [4.90833223e-02, 3.43903080e-02, 3.22353356e-02], [1.44766076e-02, 3.39248963e-02, 1.42291479e-04],
    [8.75749043e-04, 6.82212645e-03, 2.76097053e-05], [1.86958015e-02, 3.84016186e-02, 7.33085908e-05],
    [2.01714113e-02, 4.90544215e-02, 2.34028921e-05], [2.46518422e-02, 3.29151377e-02, 3.48571630e-05],
    [2.22457591e-02, 1.21796541e-02, 1.56396593e-04], [1.72109623e-02, 3.01626958e-02, 1.36556877e-02],
    [1.83460284e-02, 1.61141958e-02, 2.87440169e-04], [3.57594155e-02, 1.80554688e-01, 2.75554154e-02],
    [2.17450950e-02, 8.66811201e-02, 3.34241726e-02]]], np.float32)


def _log(*a):
    s = "[expr] " + " ".join(str(x) for x in a)
    print(s, flush=True)
    try:
        with open(os.path.join(HERE, "pipeline.log"), "a", encoding="utf-8") as f:
            f.write(time.strftime("%H:%M:%S ") + s + "\n")
    except Exception:
        pass


def models_ready(assets_dir=ASSETS):
    return all(os.path.isfile(os.path.join(assets_dir, n)) and os.path.getsize(os.path.join(assets_dir, n)) > MIN_SIZE[k]
               for k, n in MODELS.items())


def ensure_models(assets_dir=ASSETS, progress=None):
    """缺哪个下哪个。失败直接抛（不静默）。progress(name, done, total)。"""
    os.makedirs(assets_dir, exist_ok=True)
    for k, n in MODELS.items():
        p = os.path.join(assets_dir, n)
        if os.path.isfile(p) and os.path.getsize(p) > MIN_SIZE[k]:
            continue
        _log("下载", n)
        req = urllib.request.Request(HF + n, headers={"User-Agent": "Mozilla/5.0"})
        tmp = p + ".part"
        with urllib.request.urlopen(req, timeout=60) as r, open(tmp, "wb") as f:
            total = int(r.headers.get("Content-Length") or 0)
            done, last = 0, 0.0
            while True:
                b = r.read(1 << 20)
                if not b:
                    break
                f.write(b)
                done += len(b)
                if progress and time.time() - last > 0.5:
                    progress(n, done, total)
                    last = time.time()
        if os.path.getsize(tmp) <= MIN_SIZE[k]:
            os.remove(tmp)
            raise RuntimeError("%s 下载不完整" % n)
        os.replace(tmp, p)
    return True


def _box_mask(size, blur=0.3):
    b = int(size * 0.5 * blur)
    a = max(b // 2, 1)
    m = np.ones((size, size), np.float32)
    m[:a, :] = 0; m[-a:, :] = 0; m[:, :a] = 0; m[:, -a:] = 0
    return cv2.GaussianBlur(m, (0, 0), b * 0.25) if b > 0 else m


class ExpressionRestorer:
    def __init__(self, assets_dir=ASSETS, providers=None):
        import onnxruntime as ort
        import scipy.spatial.transform
        from insightface.utils import face_align
        self._rot, self._fa = scipy.spatial.transform.Rotation, face_align
        providers = providers or ["CUDAExecutionProvider", "CPUExecutionProvider"]
        so = ort.SessionOptions()
        self.feat = ort.InferenceSession(os.path.join(assets_dir, MODELS["feat"]), so, providers=providers)
        self.motion = ort.InferenceSession(os.path.join(assets_dir, MODELS["motion"]), so, providers=providers)
        self.gen = ort.InferenceSession(os.path.join(assets_dir, MODELS["gen"]), so, providers=providers)
        if "CUDAExecutionProvider" not in self.gen.get_providers():
            _log("⚠️ 表情迁移未跑在 CUDA 上:", self.gen.get_providers())
        self.size = 512
        self.box = _box_mask(self.size)
        self.pad = 12
        self.last_ms = 0.0
        self.ema = 0.0                      # 表情量时域平滑（0=关）。运动大时会滞后，默认关，供 A/B
        self._prev = None
        _log("表情迁移就绪 providers=%s" % self.gen.get_providers()[0])

    @staticmethod
    def _prep(crop):
        x = cv2.resize(crop, (256, 256), interpolation=cv2.INTER_AREA)[:, :, ::-1] / 255.0
        return np.expand_dims(x.transpose(2, 0, 1), 0).astype(np.float32)

    def restore(self, temp_img, orig_img, kps, factor=0.8, areas="all", occ128=None):
        """temp_img=换脸后整帧（就地修改并返回），orig_img=原始整帧（驱动脸），kps=5 点，factor 0~1.2，
        areas: all / lower / upper，occ128: arcface-128 空间的遮挡遮罩（1=脸区），None=不用。"""
        S = self.size
        kps = np.asarray(kps, np.float32)
        tcrop, _ = self._fa.norm_crop2(orig_img, kps, S)
        pcrop, M = self._fa.norm_crop2(temp_img, kps, S)
        t = time.perf_counter()
        fv = self.feat.run(None, {"input": self._prep(pcrop)})[0]
        tgt_expr = self.motion.run(None, {"input": self._prep(tcrop)})[5]
        pitch, yaw, roll, scale, trans, tmp_expr, mp = self.motion.run(None, {"input": self._prep(pcrop)})
        R = self._rot.from_euler("xyz", [float(pitch), float(yaw), float(roll)], degrees=True).as_matrix().astype(np.float32)
        tgt = tgt_expr.copy()
        if areas == "lower":
            tgt[:, UPPER] = tmp_expr[:, UPPER]
        elif areas == "upper":
            tgt[:, LOWER] = tmp_expr[:, LOWER]
        tgt[:, FIXED] = tmp_expr[:, FIXED]
        tgt = tgt * float(factor) + tmp_expr * (1.0 - float(factor))
        if self.ema > 0 and self._prev is not None and self._prev.shape == tgt.shape:
            tgt = self.ema * self._prev + (1 - self.ema) * tgt
        self._prev = tgt
        tgt = np.clip(tgt, EXPRESSION_MIN, EXPRESSION_MAX)
        tgt_mp = (scale * (mp @ R.T + tgt) + trans).astype(np.float32)
        tmp_mp = (scale * (mp @ R.T + tmp_expr) + trans).astype(np.float32)
        out = self.gen.run(None, {"feature_volume": fv, "source": tgt_mp, "target": tmp_mp})[0][0]
        self.last_ms = (time.perf_counter() - t) * 1000
        out = (out.transpose(1, 2, 0).clip(0, 1) * 255.0).astype(np.uint8)[:, :, ::-1]
        mask = self.box
        if occ128 is not None:
            mask = mask * cv2.resize(np.asarray(occ128, np.float32), (S, S)).clip(0, 1)   # 同模板 4 倍缩放，几何精确
        # ROI 回贴（§3：全帧回贴是 65ms 大坑）
        h, w = temp_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)
        if x1 - x0 <= 2 or y1 - y0 <= 2:
            return temp_img
        IMr = IM.copy(); IMr[0, 2] -= x0; IMr[1, 2] -= y0
        fr = cv2.warpAffine(out, IMr, (x1 - x0, y1 - y0), borderMode=cv2.BORDER_REPLICATE)
        mk = cv2.warpAffine(mask, IMr, (x1 - x0, y1 - y0), borderValue=0.0).clip(0, 1)[..., None]
        roi = temp_img[y0:y1, x0:x1].astype(np.float32)
        temp_img[y0:y1, x0:x1] = (mk * fr + (1 - mk) * roi).astype(np.uint8)
        return temp_img


if __name__ == "__main__":
    import sys
    sys.stdout.reconfigure(encoding="utf-8")
    ensure_models(progress=lambda n, d, t: print("\r  %s %5.1f%%" % (n, 100.0 * d / max(t, 1)), end=""))
    print("\n模型就绪:", models_ready())
