# -*- coding: utf-8 -*-
"""
relight_fbc.py — 方案A：用 IC-Light FBC 出一张「这个人站在这个背景光线下该是什么样」的
参考照，从中蒸馏出一个逐帧零成本的颜色变换，实时只套用不重新生成。

2026-08-15 三步离线验证通过（scratchpad/fbc_phase1.py / fbc_phase2.py / fbc_phase2b.py，
产物 results/dark_diag/fbc_*.jpg，user 已过目拍板）：
  Phase1：FBC 出图身份可辨、环境光正确反弹到衣服/皮肤（offset合并686个张量，
          与项目里已验证的 FC 模型数字一致）。
  Phase2：LAB 逐通道仿射拟合（不是精细 3D LUT——单帧样本只覆盖"皮肤+浅色衣服"两类
          颜色，用更复杂的模型会在没见过的颜色区间外插失控，复杂度必须匹配数据量），
          套到 3 帧未参与拟合的帧上，ΔLAB 位移高度一致（a 通道 18.2~18.8，
          b 通道 28.3~28.9）——证明是稳定可泛化的变换，不是拟合巧合。
  Phase2b：满强度过头（白T恤变饱和橙色，像"衣服被染色"）。强度扫描后
          user 拍板默认 0.3（做成可调，见 pipeline_rt.py 的 relight_strength）。

架构位置：本模块只做「重（扩散生成，~1-2s）→ 轻（LAB仿射，纯数组运算 <1ms）」的蒸馏，
不在实时路径里跑扩散模型——完全符合 §2 架构铁律"重打光必须离线，实时环节不做调和"，
只是这次"重打光"离线求的不是背景本身，而是一个作用于人像的颜色变换。
"""
import time

import cv2
import numpy as np
import torch

from validated_impl import patch_unet_for_iclight_fbc, resize_and_center_crop, rmbg_cutout

WORK_SIZE = (768, 448)          # IC-Light 工作分辨率，源自 offline_bg 既有实测值，/64 整除
SD_MODEL = "stablediffusionapi/realistic-vision-v51"      # 与 offline_bg 同一底模，一致的美学先验
FBC_REPO, FBC_FILE = "lllyasviel/ic-light", "iclight_sd15_fbc.safetensors"
PROMPT = "a person standing indoors, photograph, best quality"
NEG_PROMPT = "lowres, bad quality, distorted"
STEPS, CFG_SCALE = 25, 2.0

_SHARED = None
_LOCK = None      # 懒建：threading 在下面 shared() 里 import，避免改动模块顶部 import 顺序


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


class FBCRelighter:
    """SD1.5 + IC-Light(fbc) 重打光引擎。模型加载约 30-40s，进程内常驻复用（同
    offline_bg.ICLightRelighter.shared() 的单例手法）。"""

    def __init__(self):
        from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline
        from huggingface_hub import hf_hub_download

        t0 = time.time()
        icl_path = hf_hub_download(FBC_REPO, FBC_FILE)
        self.pipe = StableDiffusionPipeline.from_pretrained(
            SD_MODEL, torch_dtype=torch.float16, safety_checker=None).to("cuda")
        self._sched_cfg = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config).config
        self._sched_cls = DPMSolverMultistepScheduler
        self.unet, self.vae = self.pipe.unet, self.pipe.vae
        merged = patch_unet_for_iclight_fbc(self.unet, icl_path, "cuda", torch.float16)
        _log("就绪 %.1fs，offset合并 %d 个张量（应为 686）" % (time.time() - t0, merged))

    @classmethod
    def shared(cls):
        """进程内单例。🔴 加锁：节点启动的预加载线程与用户点「生成校正」的拟合线程可能同时首次调用，
        不加锁会加载两份模型（显存翻倍、慢一倍）。"""
        global _SHARED, _LOCK
        import threading
        if _LOCK is None:
            _LOCK = threading.Lock()
        with _LOCK:
            if _SHARED is None:
                _SHARED = cls()
            return _SHARED

    @torch.no_grad()
    def _encode(self, bgr):
        x = torch.from_numpy(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 127.5 - 1)
        x = x.permute(2, 0, 1)[None].to("cuda", torch.float16)
        return self.vae.encode(x).latent_dist.mode() * self.vae.config.scaling_factor

    @torch.no_grad()
    def _decode(self, lat):
        o = self.vae.decode(lat / self.vae.config.scaling_factor).sample
        a = ((o[0].permute(1, 2, 0).float().cpu().numpy() + 1) * 127.5).clip(0, 255).astype(np.uint8)
        return cv2.cvtColor(a, cv2.COLOR_RGB2BGR)

    @torch.no_grad()
    def relight(self, person_bgr, bg_bgr, alpha01, seed=0):
        """person_bgr=换脸+修复后的人像帧；alpha01=同帧的人像置信度(0-1 float, HxW)；
        bg_bgr=目标虚拟背景。返回 (relit_bgr, person_ws, alpha_ws)——后两个是工作分辨率
        对齐后的版本，供 fit_transform 用同一坐标系。"""
        W, H = WORK_SIZE
        person_ws = resize_and_center_crop(person_bgr, W, H)
        alpha_ws = cv2.resize(alpha01.astype(np.float32), (W, H))
        bg_ws = resize_and_center_crop(bg_bgr, W, H)
        fg_cut = rmbg_cutout(person_ws, alpha_ws)

        fg_lat, bg_lat = self._encode(fg_cut), self._encode(bg_ws)
        concat_conds = torch.cat([fg_lat, bg_lat], dim=1)          # 8ch：4前景+4背景
        pe, ne = self.pipe.encode_prompt(PROMPT, "cuda", 1, True, NEG_PROMPT)

        sch = self._sched_cls.from_config(self._sched_cfg)         # 每次新建：多步调度器有内部状态
        sch.set_timesteps(STEPS, device="cuda")
        gen = torch.Generator(device="cuda").manual_seed(int(seed))
        lat = torch.randn((1, 4, H // 8, W // 8), generator=gen, device="cuda",
                          dtype=torch.float16) * sch.init_noise_sigma
        for tt in sch.timesteps:
            li = sch.scale_model_input(lat, tt)
            inp = torch.cat([li, concat_conds], dim=1)             # 12ch：4噪声+8条件
            a = self.unet(inp, tt, encoder_hidden_states=pe).sample
            b = self.unet(inp, tt, encoder_hidden_states=ne).sample
            lat = sch.step(b + CFG_SCALE * (a - b), tt, lat).prev_sample
        return self._decode(lat), person_ws, alpha_ws


def fit_transform(person_bgr, alpha01, relit_bgr, trust_thresh=0.6):
    """从 (原图, 置信度高的人像区) → (IC-Light出图) 拟合 LAB 逐通道仿射。
    返回 [(k,b)_L, (k,b)_a, (k,b)_b] 或 None（样本太少时）。"""
    trust = alpha01 > trust_thresh
    if int(trust.sum()) < 500:
        return None, int(trust.sum())
    lab_o = cv2.cvtColor(person_bgr, cv2.COLOR_BGR2LAB).astype(np.float32)
    lab_r = cv2.cvtColor(relit_bgr, cv2.COLOR_BGR2LAB).astype(np.float32)
    coefs = []
    for c in range(3):
        k, b = np.polyfit(lab_o[:, :, c][trust], lab_r[:, :, c][trust], 1)
        coefs.append((float(k), float(b)))
    return coefs, int(trust.sum())


def apply_relight_lab(bgr_u8, alpha01, coefs, strength, pad=24):
    """逐帧套用。

    🔴 2026-08-15 实测纠错（两轮，第二轮是 user 真机实测抓出来的真 bug，不是"还能更快"
       的优化空间）：
       ① 最初"纯数组运算=零成本"是没测过的假设——1080p 整帧跑两次 cv2 LAB 转换实测
          34.5ms，把总耗时从 63ms 推到 98ms，超预算。改成只裁到**人像包围盒**（同 §3
          的 ROI 回贴手法：换脸全帧回贴 65ms→ROI回贴 17.6ms 是本项目最大的单项性能
          教训）降到 17.49ms——这一步是对的，裁剪不改变分辨率，纯粹省去背景那部分
          不需要处理的像素。
       ② 曾在 ROI 内又加了一层"降采样到 280px 做颜色数学、再放大贴回"，想进一步省到
          11ms，自己判断"颜色变换是全局平滑的，降采样不会有可感知差异"——**这个判断
          没有验证过，而且是错的**：user 开启后反馈"人物和背景都变模糊"，真机实测坐实。
          根因：代码把「降采样→放大」之后的整个 ROI **整块贴回**，不是只把「颜色差异」
          贴回去——即便某处 alpha≈0（该保持原样的背景边角），那里的像素也被换成了
          "从 280px 放大回来"的模糊版本，因为 lab/corr 两个变量都来自同一个降采样源，
          alpha 混合只是在两个同样模糊的版本之间插值，混不出清晰的原图。
          已改回**不降采样、只裁 ROI**——17.49ms 是量化正确性代价后该付的价，
          不是"还没榨干"的余地。以后类似"降分辨率省时间"的优化，套用前必须先放大
          看边缘/纹理有没有变软，不能只凭"是全局变换应该没事"的直觉判断。
    包围盒从 alpha 的非零区求，留 pad 像素余量防止边缘裁切。
    strength=0 原样；strength=1 拟合时的满强度（Phase2 实测过头，user 拍板默认 0.3）。
    按 alpha 加权混合回原像素——只在置信度高的人像区生效，边缘自然过渡，不碰背景。
    """
    if not coefs or strength <= 0:
        return bgr_u8
    ys, xs = np.where(alpha01 > 0.05)
    if ys.size == 0:
        return bgr_u8
    h, w = bgr_u8.shape[:2]
    y0, y1 = max(int(ys.min()) - pad, 0), min(int(ys.max()) + pad, h)
    x0, x1 = max(int(xs.min()) - pad, 0), min(int(xs.max()) + pad, w)
    roi = bgr_u8[y0:y1, x0:x1]
    a_roi = alpha01[y0:y1, x0:x1]

    lab = cv2.cvtColor(roi, cv2.COLOR_BGR2LAB).astype(np.float32)
    corr = lab.copy()
    for c, (k, b) in enumerate(coefs):
        k_eff = 1.0 + (k - 1.0) * strength
        b_eff = b * strength
        corr[:, :, c] = lab[:, :, c] * k_eff + b_eff
    corr = np.clip(corr, 0, 255)
    wgt = a_roi[:, :, None].astype(np.float32)
    out_lab = (lab * (1 - wgt) + corr * wgt).astype(np.uint8)
    roi_out = cv2.cvtColor(out_lab, cv2.COLOR_LAB2BGR)

    out = bgr_u8.copy()
    out[y0:y1, x0:x1] = roi_out
    return out


def preload_async():
    """节点启动时后台预加载引擎（5090 上约 30–40s；机器上第一次还要下 2GB+ 权重）。
    之后点「按当前背景生成校正」只剩 1–2s 推理——user 2026-09-07 反馈"校正为什么那么久"，根因就是每次服务重启后首次点击都在加载模型。"""
    import threading

    def run():
        try:
            FBCRelighter.shared()
            _log("预加载完成，之后生成校正只需 1–2s")
        except Exception as e:
            _log("预加载失败: %s: %s（点生成时会再试一次）" % (type(e).__name__, str(e)[:160]))

    t = threading.Thread(target=run, daemon=True, name="relight-preload")
    t.start()
    return t

