#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
offline_bg.iclight — ② 用 IC-Light 按指定方向给背景重打光。

🔴 三个要点缺一不可（v8 连错两轮才对，细节见 CLAUDE.md §4）：
   ① 权重是 OFFSET，必须加到原 UNet 上（直接替换 → 纯噪声）—— 走 validated.patch_unet_for_iclight
   ② 方向靠「线性渐变图作 img2img 起点」注入，**完全不靠 prompt**
      实测：只写 prompt 左/右亮度差 +20.90 vs +20.51（无效）
            用渐变图     +126.07 vs −120.99（符号翻转，生效）
   ③ conv_in 先 4→8 扩展；lat = add_noise(渐变latent)；每步 scale_model_input
速度：1.0–1.7s/张 @768×448 RTX5090。
"""
from typing import Optional

import cv2
import numpy as np
import torch

from .config import DEFAULT, OfflineBGConfig
from .validated import gradient_light_map, patch_unet_for_iclight

_SHARED = None


def radial_light_map(w: int, h: int, falloff: float = 1.0) -> np.ndarray:
    """中心亮四周暗的径向图，用于 front（无方向光）时把背景压成均匀布光。
    ⚠️ 未实测验证——config.front_strategy 默认是 skip，要用请自行先跑一次校验。"""
    yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
    d = np.hypot((xx - w / 2) / (w / 2), (yy - h / 2) / (h / 2))
    v = np.clip(1.0 - falloff * d / np.sqrt(2), 0, 1) * 255.0
    return np.stack([v] * 3, -1).astype(np.uint8)


def light_map(w: int, h: int, side: str) -> np.ndarray:
    """方向 → 光照注入图。四个方向走已验证的线性渐变，front 走径向。"""
    if side in ("left", "right", "top", "bottom"):
        return gradient_light_map(w, h, side)
    if side == "front":
        return radial_light_map(w, h)
    raise ValueError(f"未知光照方向: {side}")


class ICLightRelighter:
    """SD1.5 + IC-Light(fc) 重打光引擎。模型加载约 30–60s，进程内常驻复用。"""

    def __init__(self, cfg: OfflineBGConfig = DEFAULT, verbose: bool = True):
        from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline
        from huggingface_hub import hf_hub_download

        self.cfg = cfg
        self.device = cfg.device
        self.dtype = torch.float16
        icl = hf_hub_download(cfg.iclight_repo, cfg.iclight_file)
        self.pipe = StableDiffusionPipeline.from_pretrained(
            cfg.sd_model, torch_dtype=self.dtype, safety_checker=None).to(self.device)
        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
        self._sched_cls = DPMSolverMultistepScheduler
        self.unet, self.vae = self.pipe.unet, self.pipe.vae
        merged = patch_unet_for_iclight(self.unet, icl, self.device, self.dtype)   # 🔴 相加不是替换
        if verbose:
            print(f"[iclight] offset 合并 {merged} 个张量（v8 实测应为 686）", flush=True)
        self.merged_keys = merged

    @classmethod
    def shared(cls, cfg: OfflineBGConfig = DEFAULT) -> "ICLightRelighter":
        global _SHARED
        if _SHARED is None:
            _SHARED = cls(cfg)
        return _SHARED

    # ---------------- 内部 ----------------
    @torch.no_grad()
    def _encode(self, bgr: np.ndarray) -> torch.Tensor:
        x = torch.from_numpy(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 127.5 - 1)
        x = x.permute(2, 0, 1)[None].to(self.device, self.dtype)
        return self.vae.encode(x).latent_dist.mode() * self.vae.config.scaling_factor

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

    # ---------------- 对外 ----------------
    @torch.no_grad()
    def relight(self, img_bgr: np.ndarray, side: str, prompt: Optional[str] = None,
                strength: Optional[float] = None, steps: Optional[int] = None,
                cfg_scale: Optional[float] = None, seed: int = 0) -> np.ndarray:
        """把 img_bgr 按 side 方向重打光。输入输出同尺寸，尺寸须为 64 的倍数。"""
        c = self.cfg
        h, w = img_bgr.shape[:2]
        if w % 64 or h % 64:
            raise ValueError(f"IC-Light 输入尺寸须为 64 的倍数，当前 {w}×{h}（见 config.work_size）")
        prompt = prompt or c.prompt_for(side)
        strength = c.strength if strength is None else strength
        steps = c.steps if steps is None else steps
        cfg_scale = c.cfg_scale if cfg_scale is None else cfg_scale

        fg = self._encode(img_bgr)                       # 原图内容 → concat 条件（保内容）
        init = self._encode(light_map(w, h, side))       # 方向图     → img2img 起点（注方向）
        pe, ne = self.pipe.encode_prompt(prompt, self.device, 1, True, c.negative_prompt)

        sch = self._sched_cls.from_config(self.pipe.scheduler.config)   # 每次新建：DPMSolver 有内部多步状态
        sch.set_timesteps(steps, device=self.device)
        ts_all = sch.timesteps
        k = max(int(len(ts_all) * strength), 1)
        ts = ts_all[len(ts_all) - k:]

        gen = torch.Generator(device=self.device).manual_seed(int(seed))
        noise = torch.randn(init.shape, generator=gen, device=self.device, dtype=init.dtype)
        lat = sch.add_noise(init, noise, ts[:1])
        for tt in ts:
            li = sch.scale_model_input(lat, tt)
            inp = torch.cat([li, fg], 1)                 # 🔴 8 通道：4 噪声 + 4 原图条件
            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)
