#!/usr/bin/env python3
"""teaser · Boss 换装用的视觉打分工具（img env：$M/envs/img/bin/python，GPU）。被 bossswap_f0.py / bossswap_metrics.py import。

* BossSeg：SAM2.1-hiera-large（transformers Sam2Model，本地 HF_HOME snapshot），提示 = 条件 Boss 掩码（膨胀 10 px）的外接框 +
  身体内核（非武器部件，腐蚀 3 px）最远点采样 4 个正点 + 武器中线上离身体最远的一点；玩家身体内核 2 个负点。
  输出限制在 Boss 掩码膨胀 40 px 内（防止整块地面被分进去）。iou(seg, cond) = 生成画面里 Boss 轮廓与条件掩码的重合度。
  注意：胶囊掩码只是近似轮廓（f0_realchar 实测真实角色 vs 胶囊 IoU 0.39–0.75），所以要和原版 Boss 同一指标比相对值。
* BossClip：CLIP ViT-L/14，Boss 裁图对 16 种外观 + 原版描述做零样本分类，p(目标) 当「认不认得出是这个 Boss」。
"""
from __future__ import annotations

import glob
import os

import numpy as np

os.environ.setdefault("HF_HOME", "/opt/dlami/nvme/zhiyangdeng/.cache/huggingface")
HUB = f"{os.environ['HF_HOME']}/hub"
PLAYER, BOSS = (10, 12, 13, 14, 15, 16), (11, 17, 18, 19, 20, 21)
WEAPONS = (16, 21)


def snap(repo: str) -> str:
    p = sorted(glob.glob(f"{HUB}/models--{repo.replace('/', '--')}/snapshots/*"))
    if not p:
        raise FileNotFoundError(repo)
    return p[-1]


def grow(mask: np.ndarray, r: int) -> np.ndarray:
    import cv2
    if r == 0:
        return mask.copy()
    k = np.ones((2 * abs(r) + 1, 2 * abs(r) + 1), np.uint8)
    f = cv2.dilate if r > 0 else cv2.erode
    return f(mask.astype(np.uint8), k) > 0


def fps_points(core: np.ndarray, k: int = 4) -> list[list[float]]:
    ys, xs = np.nonzero(core)
    P = np.stack([xs, ys], 1).astype(np.float32)
    sel = [P[np.argmin(((P - P.mean(0)) ** 2).sum(1))]]
    for _ in range(k - 1):
        sel.append(P[np.argmax(np.min([((P - q) ** 2).sum(1) for q in sel], 0))])
    return [[float(x), float(y)] for x, y in sel]


def bbox(mask: np.ndarray, pad: float = 0.0, min_side: int = 1, W: int = 832, H: int = 480):
    ys, xs = np.nonzero(mask)
    x0, x1, y0, y1 = xs.min(), xs.max() + 1, ys.min(), ys.max() + 1
    w, h = x1 - x0, y1 - y0
    cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
    w, h = max(min_side, w * (1 + pad)), max(min_side, h * (1 + pad))
    return (int(max(0, cx - w / 2)), int(max(0, cy - h / 2)), int(min(W, cx + w / 2)), int(min(H, cy + h / 2)))


class BossSeg:
    def __init__(self, device="cuda"):
        import torch
        from transformers import Sam2Model, Sam2Processor
        p = snap("facebook/sam2.1-hiera-large")
        self.torch, self.dev = torch, device
        self.proc = Sam2Processor.from_pretrained(p)
        self.model = Sam2Model.from_pretrained(p, dtype=torch.float32).to(device).eval()

    def prompts(self, ids: np.ndarray):
        mk = np.isin(ids, BOSS)
        if mk.sum() < 300:
            return None
        body = mk & ~np.isin(ids, WEAPONS)
        core = grow(body, -3)
        core = core if core.sum() >= 10 else (body if body.sum() >= 10 else mk)
        box = list(bbox(grow(mk, 10)))
        pts = fps_points(core)
        lab = [1] * len(pts)
        wp = grow(mk & np.isin(ids, WEAPONS), -2)
        if wp.sum() >= 60:
            ys, xs = np.nonzero(wp)
            c = np.array(pts[0])
            k = int(np.argmax((xs - c[0]) ** 2 + (ys - c[1]) ** 2))
            pts.append([float(xs[k]), float(ys[k])])
            lab.append(1)
        pc = grow(np.isin(ids, PLAYER) & ~np.isin(ids, WEAPONS), -3) & ~grow(mk, 4)
        if pc.sum() >= 30:
            for q in fps_points(pc, 2):
                pts.append(q)
                lab.append(0)
        return mk, box, pts, lab

    def __call__(self, rgbs: list[np.ndarray], idss: list[np.ndarray]) -> list[np.ndarray | None]:
        """每帧一个 Boss 掩码（看不到 Boss 的帧 = None）。批处理。"""
        outs: list = [None] * len(rgbs)
        todo = []
        for i, (rgb, ids) in enumerate(zip(rgbs, idss)):
            pr = self.prompts(ids)
            if pr is not None:
                todo.append((i, rgb, pr))
        for s in range(0, len(todo), 8):
            ch = todo[s:s + 8]
            inp = self.proc(images=[c[1] for c in ch], input_boxes=[[c[2][1]] for c in ch],
                            input_points=[[c[2][2]] for c in ch], input_labels=[[c[2][3]] for c in ch],
                            return_tensors="pt").to(self.dev)
            with self.torch.no_grad():
                res = self.model(**inp, multimask_output=False)
            ms = self.proc.post_process_masks(res.pred_masks.cpu(), inp["original_sizes"])
            for (i, _, pr), m in zip(ch, ms):
                outs[i] = (m[0, 0].numpy() > 0) & grow(pr[0], 40)
        return outs


def iou(a: np.ndarray, b: np.ndarray) -> float:
    u = (a | b).sum()
    return float((a & b).sum() / u) if u else 0.0


class BossClip:
    def __init__(self, labels: dict[str, str], device="cuda"):
        """labels：键 → 文本描述（含原版 Boss）。"""
        import torch
        from transformers import CLIPModel, CLIPProcessor
        p = snap("openai/clip-vit-large-patch14")
        self.torch, self.dev = torch, device
        self.model = CLIPModel.from_pretrained(p, dtype=torch.float16).to(device).eval()
        self.proc = CLIPProcessor.from_pretrained(p)
        self.keys = list(labels)
        txt = [f"a top-down video game screenshot of {labels[k]}" for k in self.keys]
        with torch.no_grad():
            t = self.proc(text=txt, return_tensors="pt", padding=True, truncation=True).to(device)
            e = self.model.get_text_features(**t)
            e = e if isinstance(e, torch.Tensor) else e.pooler_output
            self.temb = torch.nn.functional.normalize(e.float(), dim=-1)

    def probs(self, crops: list[np.ndarray]) -> np.ndarray:
        """crops：RGB uint8 → [N, K] softmax 概率（CLIP logit scale 100）。"""
        from PIL import Image
        torch = self.torch
        out = []
        for s in range(0, len(crops), 32):
            im = [Image.fromarray(c) for c in crops[s:s + 32]]
            with torch.no_grad():
                x = self.proc(images=im, return_tensors="pt").to(self.dev)
                e = self.model.get_image_features(pixel_values=x["pixel_values"].half())
                e = e if isinstance(e, torch.Tensor) else e.pooler_output
                e = torch.nn.functional.normalize(e.float(), dim=-1)
                out.append(torch.softmax(100 * e @ self.temb.T, -1).cpu().numpy())
        return np.concatenate(out, 0)


def boss_crop(rgb: np.ndarray, mask: np.ndarray, pad: float = 0.25, min_side: int = 96) -> np.ndarray:
    x0, y0, x1, y1 = bbox(mask, pad=pad, min_side=min_side)
    return np.ascontiguousarray(rgb[y0:y1, x0:x1])


def lab(img: np.ndarray) -> np.ndarray:
    import cv2
    return cv2.cvtColor(img, cv2.COLOR_RGB2LAB).astype(np.float32) * np.array([100 / 255, 1, 1], np.float32)


FFMPEG = f"/opt/dlami/nvme/zhiyangdeng/maliang/.venv/lib/python3.12/site-packages/imageio_ffmpeg/binaries/ffmpeg-linux-x86_64-v7.0.2"


def read_video(path: str, w: int = 832, h: int = 480) -> np.ndarray:
    """mp4 → RGB uint8 [F,h,w,3]（ffmpeg 解码；尺寸须已知，不依赖 PyAV）。"""
    import subprocess
    p = subprocess.run([FFMPEG, "-loglevel", "error", "-i", path, "-f", "rawvideo", "-pix_fmt", "rgb24", "-"],
                       capture_output=True, check=True)
    return np.frombuffer(p.stdout, np.uint8).reshape(-1, h, w, 3)