#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
occluder.py — 遮挡感知（xseg_1.onnx），修 §11.2「手遮脸时手变模糊、透出脸」。

背景：RoiSwapper 用一个**固定的椭圆羽化遮罩**把换出来的脸贴回去，完全不管脸前
有没有东西挡着。手挡在脸前时，那块手的像素被 `mk*fr` 连同脸一起换掉了——成品里
手变糊、里面浮出鼻子嘴巴的轮廓（图证 results/debug_frame/occ_cmp.jpg）。

修法：xseg 在对齐空间检测「哪些像素是真正露出来的脸」，把它并进融合遮罩——
遮挡处遮罩降为 0，保留原像素（手），换脸只作用在真正露出的脸皮肤上。

🔴 接口口径（CLAUDE.md §11.2 已在旧实例上实测验证过，别按 FaceFusion 文档字面猜）：
   · 输入 NHWC (1,256,256,3)，BGR，/255，float32（**不翻通道**——FaceFusion 全程 BGR）
   · 输出 (1,256,256,1)，**高值 = 脸区（该换）**，低值 = 遮挡区
   · 后处理照抄 FaceFusion create_occlusion_mask：blur σ=5 → clip(0.5,1) → (x-0.5)*2，
     把软输出硬化成「脸区≈1 / 遮挡区≈0 + 一圈软边」，避免手的边界留硬切缝。
🔴 ORT 的 CUDA EP 依赖 torch 先把 CUDA 库加载进进程（库在 torch 的 nvidia/cu13 目录，
   不在系统库路径）。调用方（pipeline_rt）必须已 `import torch`，否则这里静默回落 CPU
   （§11.4 坑2：inswapper 128 CPU 346ms vs GPU 6ms，相差 58 倍）。
"""
import os

import cv2
import numpy as np

# HF resolve 与 GitHub releases 双 fallback（68MB）
XSEG_URLS = [
    "https://huggingface.co/facefusion/models-3.1.0/resolve/main/xseg_1.onnx",
    "https://github.com/facefusion/facefusion-assets/releases/download/models-3.0.0/xseg_1.onnx",
]


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


def _download(path):
    """下载原始 xseg_1.onnx。实例是临时的、销毁即没，所以下载逻辑必须常驻代码里。"""
    if os.path.exists(path) and os.path.getsize(path) > 1_000_000:
        return path
    import urllib.request
    for u in XSEG_URLS:
        try:
            _log("下载 xseg …", u.split("/")[2])
            urllib.request.urlretrieve(u, path)
            if os.path.getsize(path) > 1_000_000:
                _log(f"xseg 就绪 {os.path.getsize(path)/1e6:.0f}MB")
                return path
        except Exception as e:
            _log("  失败:", e)
    raise RuntimeError("xseg_1.onnx 下载失败，所有源都不可达")


def _patch_to_gpu(src, dst):
    """把原始模型的 6 个 ConvTranspose 从非对称 pad 改写成「对称 pad + Slice 切末行末列」。

    🔴 为什么（实测）：xseg 解码层 ConvTranspose 的 pads=[0,0,1,1] 是非对称裁剪，ORT
       CUDA EP 不支持 → 那 6 层回落 CPU + 插入 12 个 GPU↔CPU memcpy 节点，单次推理
       11.79ms 且在线管线里与 RVM 争抢飙到 ~26ms 尖峰、直接掉帧。
    改写是**数值精确等价**：pads=[0,0,1,1] 等于「无裁剪的满输出(2n+1) 再裁掉末行末列」，
       所以设 pads=[0,0,0,0] + Slice ends=[-1,-1] axes=[2,3] 得到完全相同结果，但对称 pad
       与 Slice 都是 CUDA 原生算子 → 全图上 GPU，单次 3.15ms、memcpy 归零。
    与原模型输出最大差 0.023（纯 GPU/CPU 浮点累加差异，非改写误差），软遮罩可忽略。
    """
    import onnx
    from onnx import TensorProto, helper
    m = onnx.load(src)
    g = m.graph
    new_nodes, patched = [], 0
    for n in g.node:
        pads = next((list(a.ints) for a in n.attribute if a.name == "pads"), None) \
            if n.op_type == "ConvTranspose" else None
        if pads == [0, 0, 1, 1]:
            for a in n.attribute:
                if a.name == "pads":
                    del a.ints[:]; a.ints.extend([0, 0, 0, 0])
            out = n.output[0]
            n.output[0] = out + "_full"
            new_nodes.append(n)
            g.initializer.extend([
                helper.make_tensor(out + "_st", TensorProto.INT64, [2], [0, 0]),
                helper.make_tensor(out + "_en", TensorProto.INT64, [2], [-1, -1]),
                helper.make_tensor(out + "_ax", TensorProto.INT64, [2], [2, 3])])
            new_nodes.append(helper.make_node(
                "Slice", [out + "_full", out + "_st", out + "_en", out + "_ax"],
                [out], name=out + "_slice"))
            patched += 1
        else:
            new_nodes.append(n)
    del g.node[:]; g.node.extend(new_nodes)
    onnx.checker.check_model(m)
    onnx.save(m, dst)
    _log(f"xseg 已图手术改造为全 GPU（改写 {patched} 个 ConvTranspose）→ {os.path.basename(dst)}")
    return dst


def ensure_xseg(path):
    """返回一个**可用**的 xseg 模型路径，优先给出全 GPU 改造版。

    流程：下载原始模型 → 图手术产出 xxx_gpu.onnx（缓存）→ 返回 GPU 版。
    图手术失败（无 onnx / 异常）时退回原始模型——功能不受影响，只是部分算子回落 CPU 慢一些
    （此时应把 occ_every 调大降频，见 pipeline_rt / validated_impl）。
    """
    gpu = os.path.splitext(path)[0] + "_gpu.onnx"
    if os.path.exists(gpu) and os.path.getsize(gpu) > 1_000_000:
        return gpu
    orig = _download(path)
    try:
        return _patch_to_gpu(orig, gpu)
    except Exception as e:
        _log("⚠️ xseg 图手术失败，退回原始模型（部分算子回落 CPU，建议调大 occ_every）:",
             type(e).__name__, e)
        return orig


class XSegOccluder:
    """输入一张**已对齐的人脸 crop**（任意尺寸 BGR uint8），返回同尺寸的
    「可换区置信度」浮点遮罩 [0,1]：1 = 露出的脸（该换），0 = 被遮挡（保留原像素）。

    harden=True 时套 FaceFusion 的硬化后处理；设 False 返回原始软输出（用于 debug 观察）。
    """

    def __init__(self, model_path, providers=None, size=256, harden=True, invert=False):
        import onnxruntime as ort
        providers = providers or ["CUDAExecutionProvider", "CPUExecutionProvider"]
        self.sess = ort.InferenceSession(model_path, providers=providers)
        real = self.sess.get_providers()
        if "CUDAExecutionProvider" not in real:
            _log("⚠️ xseg 未跑在 CUDA 上，现有 provider:", real)
        self.in_name = self.sess.get_inputs()[0].name
        self.out_name = self.sess.get_outputs()[0].name
        self.size = size
        self.harden = harden
        # invert=True：模型极性相反时（高值=遮挡区）翻转，使 face_mask 恒为「可换区置信度」。
        # 出厂 False（依 CLAUDE.md 已验证口径：高=脸区）；实测反了就置 True，一处生效。
        self.invert = invert

    def _raw(self, aligned_bgr):
        """跑一次 xseg，返回 256×256 float32 [0,1]（未硬化、未 resize）。"""
        x = cv2.resize(aligned_bgr, (self.size, self.size)).astype(np.float32) / 255.0
        x = x[None]                                    # NHWC (1,256,256,3)
        y = self.sess.run([self.out_name], {self.in_name: x})[0]
        m = np.squeeze(y).astype(np.float32)           # (256,256)，兼容 NHWC/NCHW 输出
        m = np.clip(m, 0.0, 1.0)
        return 1.0 - m if self.invert else m

    def face_mask(self, aligned_bgr, out_size=None):
        """返回可换区遮罩，float32 [0,1]，尺寸 out_size=(w,h)（默认与输入同尺寸）。"""
        if out_size is None:
            h, w = aligned_bgr.shape[:2]
            out_size = (w, h)
        m = self._raw(aligned_bgr)
        if self.harden:
            # 🔴 照抄 FaceFusion：软边保留一圈过渡（不留硬切缝），中间硬化到 0/1
            m = (np.clip(cv2.GaussianBlur(m, (0, 0), 5), 0.5, 1.0) - 0.5) * 2.0
        if out_size != (self.size, self.size):
            m = cv2.resize(m, out_size, interpolation=cv2.INTER_LINEAR)
        return m
