#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""swappers.py — 可切换的换脸后端（inswapper 之外的模型）。

起因（CLAUDE.md §11.2c③）：inswapper 所属的「属性编码器」谱系**结构与形状始终来自目标图**，
换脸只换五官纹理、换不了脸型/头骨轮廓。user 要求把候选模型都加成选项亲自对比。

🔴 **调用口径全部由本地实测标定**（scratchpad/hifi_probe*.py、hyper_probe.py、swap_all.py），
   猜错**不报错、直接输出彩色乱块**——这是这一类模型最坑的地方。三条通用规律：

   ① **source 向量必须是单位向量**。带 arcface converter 的模型（hififace/ghost/simswap），
      converter 的输出**不是**单位向量（实测模长 3.6~10.2），**必须再做一次 L2 归一化**，
      否则生成器直接崩。这一条让 hififace 从 0.16 跳到 0.48。
      不带 converter 的（hyperswap）直接用 `face.normed_embedding`（本身已 L2）。
      converter 的输入用 `face.embedding`（原始）不是 normed。
   ② 输入/输出的归一化**可以不对称**（[0,1] 与 [-1,1] 各自独立），必须分别扫。
   ③ **不能只按 ArcFace 分数选口径**：分数最高的组合可能是偏暗/过曝的废图（hififace 上
      实测：最高分 0.511 的图严重偏暗发红，视觉正确的只有 0.469）。ArcFace 对亮度色调
      不敏感。所以下面每个 spec 都是**色域偏差合格前提下**的最优，不是裸分最高。

📊 横向实测（本地 CPU，同一对素材，像 persona = ArcFace 余弦）：
     inswapper_128  0.847 ← 仍然最强，保持默认
     hyperswap_1b   0.657（256 输出、CPU 436ms 比 inswapper 快 2.6 倍，视觉上脸型偏宽、有竞争力）
     hyperswap_1c   0.653 · hyperswap_1a 0.642
     simswap_512    0.538（512 输出，脸偏白）
     hififace_256   0.469（未复现论文里的形状转移优势）
     ghost_3_256    0.189（口径未调通，故不列入选项）
"""
import os

import cv2
import numpy as np

R30 = "https://huggingface.co/facefusion/models-3.0.0/resolve/main/"
R31 = "https://huggingface.co/facefusion/models-3.1.0/resolve/main/"
R33 = "https://huggingface.co/facefusion/models-3.3.0/resolve/main/"

# 每个后端的完整口径。src: "normed"=直接用 normed_embedding / "conv_l2"=过 converter 再 L2
SPECS = {
    "hyperswap_1a_256": dict(gen=R33 + "hyperswap_1a_256.onnx", conv=None,
                             src="normed", inn="-11", outn="-11", rgb=True, size=256),
    "hyperswap_1b_256": dict(gen=R33 + "hyperswap_1b_256.onnx", conv=None,
                             src="normed", inn="-11", outn="-11", rgb=True, size=256),
    "hyperswap_1c_256": dict(gen=R33 + "hyperswap_1c_256.onnx", conv=None,
                             src="normed", inn="-11", outn="-11", rgb=True, size=256),
    "hififace_unofficial_256": dict(gen=R31 + "hififace_unofficial_256.onnx",
                                    conv=R31 + "arcface_converter_hififace.onnx",
                                    src="conv_l2", inn="01", outn="01", rgb=True, size=256),
    "simswap_unofficial_512": dict(gen=R30 + "simswap_unofficial_512.onnx",
                                   conv=R30 + "arcface_converter_simswap.onnx",
                                   src="conv_l2", inn="01", outn="01", rgb=False, size=512),
}


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


def _fetch(url, assets_dir, minsize=1_000_000):
    p = os.path.join(assets_dir, os.path.basename(url))
    if not (os.path.exists(p) and os.path.getsize(p) > minsize):
        import urllib.request
        _log("下载", os.path.basename(url), "…")
        urllib.request.urlretrieve(url, p)
        _log(f"  就绪 {os.path.getsize(p)/1e6:.0f}MB")
    return p


class GenericSwapper:
    """按 SPECS 里的口径驱动一个 ONNX 换脸生成器。

    RoiSwapper 依赖的三件事：
      · `input_size` —— 对齐空间边长，据此建羽化遮罩
      · `prepare_latent(src_face)` —— 算一次缓存
      · `swap_aligned(aligned_bgr, latent)` —— 吃对齐脸出对齐脸（uint8 BGR）
    """

    def __init__(self, name, spec, assets_dir, providers=None):
        import onnxruntime as ort
        providers = providers or ["CUDAExecutionProvider", "CPUExecutionProvider"]
        self.name, self.spec = name, spec
        self.gen = ort.InferenceSession(_fetch(spec["gen"], assets_dir), providers=providers)
        if "CUDAExecutionProvider" not in self.gen.get_providers():
            _log(f"⚠️ {name} 未跑在 CUDA 上:", self.gen.get_providers())
        gi = {i.name: i for i in self.gen.get_inputs()}
        self.go = [o.name for o in self.gen.get_outputs()]
        # 🔴 按 shape 维数认输入，不按名字：各模型的输入名顺序不一致（hififace 是
        #    target,source；hyperswap 是 source,target），按名字写死必然接反。
        self.tgt_n = next(n for n, i in gi.items() if len(i.shape) == 4)
        self.src_n = next(n for n, i in gi.items() if len(i.shape) == 2)
        self.conv = None
        if spec.get("conv"):
            self.conv = ort.InferenceSession(_fetch(spec["conv"], assets_dir), providers=providers)
            self.conv_in = self.conv.get_inputs()[0].name
        self.input_size = (spec["size"], spec["size"])
        self.last_mask = None

    def prepare_latent(self, src_face):
        if self.spec["src"] == "conv_l2":
            e = np.asarray(src_face.embedding, np.float32).reshape(1, 512)
            v = self.conv.run(None, {self.conv_in: e})[0].astype(np.float32)
        else:                                     # "normed"：本身已是单位向量
            v = np.asarray(src_face.normed_embedding, np.float32).reshape(1, 512)
        n = float(np.linalg.norm(v))              # 🔴 无论哪条路，最终都必须是单位向量
        return (v / (n if n > 1e-6 else 1.0)).astype(np.float32)

    def swap_aligned(self, aligned_bgr, latent):
        S = self.input_size[0]
        a = aligned_bgr if aligned_bgr.shape[0] == S else cv2.resize(aligned_bgr, (S, S))
        x = (a[:, :, ::-1] if self.spec["rgb"] else a).astype(np.float32) / 255.0
        if self.spec["inn"] == "-11":
            x = x * 2.0 - 1.0
        out = self.gen.run(self.go, {self.tgt_n: x.transpose(2, 0, 1)[None],
                                     self.src_n: latent})
        if len(out) > 1:
            self.last_mask = out[1]
        y = out[0][0].transpose(1, 2, 0)
        if self.spec["outn"] == "-11":
            y = (y + 1.0) / 2.0
        y = np.clip(y * 255.0, 0, 255).astype(np.uint8)
        return y[:, :, ::-1] if self.spec["rgb"] else y


def list_models():
    """网页下拉用：默认 inswapper_128 排第一。"""
    return ["inswapper_128"] + list(SPECS)


def build_swapper(name, assets_dir, providers=None):
    """按名字造后端。返回 None = 用默认 inswapper（由调用方处理）。"""
    if name in (None, "", "inswapper_128"):
        return None
    if name not in SPECS:
        raise ValueError("未知换脸模型: " + str(name))
    return GenericSwapper(name, SPECS[name], assets_dir, providers)
