# -*- coding: utf-8 -*-
"""
双区 DINOv2-large 特征提取模块（上半区·精灵识别 + 下半区·Layer3 版本/特效）。

两个半区模型完全同构：facebook/dinov2-large，取 CLS token，L2 归一化，
**无 BN Neck**，forward 必须开 interpolate_pos_encoding=True（输入 196×392 非正方形）。

预处理严格对齐训练侧（stage1_upper_half.py 的 letterbox + train_dinov2_*_half.py 的 transform_clean）：
    原图/裁切图 → CardDetector.detect_and_crop(RGB 数组) → PIL letterbox 392×392
                → 上半区 rows[0:196] / 下半区 rows[196:392]（196×392）
                → ImageNet 归一化 → backbone(interpolate_pos_encoding=True) → CLS → L2

建库与查询必须用同一套预处理（本模块），否则相似度不可比。
"""
import os
import numpy as np
from PIL import Image

os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

HF_MODEL_ID = "facebook/dinov2-large"
IMG_SIZE = 392
HALF = 196           # 392 // 2，与 stage1 / Layer3 stage3 完全一致
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
FEATURE_DIM = 1024   # dinov2-large CLS 维度


# ===================== 预处理（与训练严格对齐，纯函数，不依赖 torch）=====================
def letterbox392(pil_img):
    """与 stage1_upper_half.py / Layer3 stage3_finalize_bottom_crop_norot.letterbox 完全一致。
    PIL BICUBIC、min-ratio、黑色(0,0,0)填充、居中、输出 392×392。"""
    img = pil_img.convert("RGB")
    w, h = img.size
    r = min(IMG_SIZE / h, IMG_SIZE / w)
    nw, nh = int(w * r), int(h * r)
    img = img.resize((nw, nh), Image.BICUBIC)
    canvas = Image.new("RGB", (IMG_SIZE, IMG_SIZE), (0, 0, 0))
    canvas.paste(img, ((IMG_SIZE - nw) // 2, (IMG_SIZE - nh) // 2))
    return canvas


def split_half(pil_letterboxed):
    """从 392×392 letterbox 图取上下半区（各 196×392）。
    上半区 crop((0,0,392,196)) = rows[0:196]；下半区 crop((0,196,392,392)) = rows[196:392]。
    两者来自同一 letterbox，card_id 一一对应。"""
    upper = pil_letterboxed.crop((0, 0, IMG_SIZE, HALF))
    lower = pil_letterboxed.crop((0, HALF, IMG_SIZE, IMG_SIZE))
    return upper, lower


def half_to_chw(pil_half):
    """PIL 196×392 半区 → ImageNet 归一化 CHW float32（对齐 A.Resize(196,392)+Normalize+ToTensorV2）。
    若尺寸非 196×392，用 INTER_CUBIC 补 resize（与训练 transform 一致）。"""
    import cv2
    arr = np.array(pil_half.convert("RGB"))
    if arr.shape[:2] != (HALF, IMG_SIZE):
        arr = cv2.resize(arr, (IMG_SIZE, HALF), interpolation=cv2.INTER_CUBIC)
    arr = arr.astype(np.float32) / 255.0
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return arr.transpose(2, 0, 1)  # CHW


class DualCardFeatureExtractor:
    """双区 DINOv2-large 卡牌特征提取器：上半区认精灵、下半区认版本/特效。"""

    def __init__(self, upper_pth, lower_pth, device=None, batch_size=64, hf_model_id=HF_MODEL_ID):
        import torch
        self.torch = torch
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        self.batch_size = batch_size
        self.feature_dim = FEATURE_DIM

        print(f"[DualExtractor] 加载上半区模型: {upper_pth}", flush=True)
        self.upper = self._load_half_net(upper_pth, hf_model_id)
        print(f"[DualExtractor] 加载下半区模型: {lower_pth}", flush=True)
        self.lower = self._load_half_net(lower_pth, hf_model_id)
        print(f"[DualExtractor] device={self.device}  batch_size={batch_size}", flush=True)

    def _load_half_net(self, pth_path, hf_model_id):
        """加载单个半区 dinov2-large backbone 并装入训练好的权重。
        训练模型结构为 self.backbone=Dinov2Model，故 .pth 的 key 带 'backbone.' 前缀，这里剥离后装入裸 Dinov2Model。"""
        import torch
        from transformers import Dinov2Model
        backbone = Dinov2Model.from_pretrained(hf_model_id)
        state = torch.load(pth_path, map_location="cpu")
        state = {k[len("backbone."):]: v for k, v in state.items() if k.startswith("backbone.")}
        missing, unexpected = backbone.load_state_dict(state, strict=False)
        if missing:
            print(f"  [warn] missing keys(前5): {list(missing)[:5]}", flush=True)
        if unexpected:
            print(f"  [warn] unexpected keys(前5): {list(unexpected)[:5]}", flush=True)
        backbone.to(self.device).eval()
        return backbone

    # ---------- 静态预处理入口（建库/查询共用）----------
    @staticmethod
    def letterbox(pil_img):
        return letterbox392(pil_img)

    @staticmethod
    def split_half(pil_letterboxed):
        return split_half(pil_letterboxed)

    @staticmethod
    def crop_to_halves(crop_rgb):
        """detect_and_crop 返回的 RGB numpy 数组 → (上半区PIL, 下半区PIL)。"""
        lb = letterbox392(Image.fromarray(crop_rgb))
        return split_half(lb)

    # ---------- 批量提特征 ----------
    def _forward(self, backbone, pil_halves):
        """对一批 PIL 半区图提 CLS 并 L2 归一化，返回 (N, 1024) float32，失败行 NaN。"""
        import torch.nn.functional as F
        torch = self.torch
        out = np.full((len(pil_halves), self.feature_dim), np.nan, dtype=np.float32)
        tensors, valid = [], []
        for i, p in enumerate(pil_halves):
            if p is None:
                continue
            try:
                tensors.append(torch.from_numpy(half_to_chw(p)).float())
                valid.append(i)
            except Exception:
                continue
        if not tensors:
            return out
        for s in range(0, len(tensors), self.batch_size):
            chunk = torch.stack(tensors[s:s + self.batch_size]).to(self.device)
            with torch.no_grad():
                # transformers>=4.5x 的 Dinov2Model.forward 已去掉 interpolate_pos_encoding 参数
                # （非标准分辨率时内部自动插值）；硬传会 TypeError。
                try:
                    cls = backbone(chunk, interpolate_pos_encoding=True).last_hidden_state[:, 0, :]
                except TypeError:
                    cls = backbone(chunk).last_hidden_state[:, 0, :]
                feat = F.normalize(cls, p=2, dim=1).cpu().numpy()
            for k, vi in enumerate(valid[s:s + self.batch_size]):
                out[vi] = feat[k]
        return out

    def extract_upper(self, pil_uppers):
        """pil_uppers: list[PIL 196×392]（或 None）→ (N, 1024) L2 归一化。"""
        return self._forward(self.upper, pil_uppers)

    def extract_lower(self, pil_lowers):
        """pil_lowers: list[PIL 196×392]（或 None）→ (N, 1024) L2 归一化。"""
        return self._forward(self.lower, pil_lowers)