# -*- coding: utf-8 -*-
"""
cascade_match.py - 双区级联卡牌检索（上半区召回 + 下半区重排）

查询图 → YOLO裁卡 → PIL letterbox 392 → 上/下半区特征
       → 上半区全库 Top-K 召回(物种级) → 下半区对候选算版本相似度
       → fusion = α·upper_sim + (1-α)·lower_sim 重排 → 最终 card_id

三种模式:
  --image <路径>        单图查询（打印 Top-K）
  --image-dir <目录>    目录批量查询
  --self-test N         全库自检索 sanity（抽 N 张 gallery 图，Top-1 应命中自己 card_id）
  --pairs <json>        版本消歧测试（test_pairs_valid.json：同精灵普通卡 vs 球闪卡），
                        内置 α∈{1.0,0.7,0.5,0.3,0.0} 扫描，输出各权重分开成功率

用法(pytorch 环境):
  python cascade_match.py --self-test 500
  python cascade_match.py --pairs test_pairs_valid.json --alpha 0.5
  python cascade_match.py --image /path/to/card.jpg
"""
import os
import sys
import json
import time
import random
import hashlib
import argparse

sys.stdout.reconfigure(encoding="utf-8")

_THIS = os.path.dirname(os.path.abspath(__file__))
_ROOT = os.path.dirname(_THIS)
sys.path.insert(0, _ROOT)
sys.path.insert(0, _THIS)

import numpy as np
from PIL import Image

import config
from modules.yolo_detector import CardDetector
from modules.feature_extractor_dual import DualCardFeatureExtractor, letterbox392, split_half

ALPHAS = [1.0, 0.7, 0.5, 0.3, 0.0]   # pairs 扫描权重；1.0=纯上半区召回(不重排), 0.0=纯下半区


def md5n(u):
    return hashlib.md5(u.encode()).hexdigest() + ".jpg"


class DualCascadeMatcher:
    def __init__(self, alpha=None, top_k_recall=None, device=None, verbose=True, backend=None,
                 detector_device=None, milvus_alias=None):
        """device: DINOv2 特征提取器设备（torch.device/str/None=自动）。
        detector_device: YOLO 分割模型设备（None=ultralytics 自动=CVD 掩码后第一张；
            多 worker 分卡时传 torch.device(f"cuda:N")——ultralytics select_device 对
            torch.device 输入直接返回、不改写 CUDA_VISIBLE_DEVICES）。
        milvus_alias: pymilvus 连接 alias（None=默认连接）。多 worker 并发时各自
            独立 alias，避免共享默认连接；单例服务不传即保持旧行为。"""
        self.alpha = config.CASCADE_ALPHA if alpha is None else alpha
        self.top_k_recall = config.CASCADE_TOP_K_RECALL if top_k_recall is None else top_k_recall
        self.verbose = verbose
        self.backend = backend or getattr(config, "CASCADE_BACKEND", "npy")
        assert self.backend in ("npy", "milvus", "gpu"), \
            f"未知 backend: {self.backend!r}（只支持 npy/milvus/gpu）"

        with open(config.GALLERY_DUAL_META_PATH, "r", encoding="utf-8") as f:
            m = json.load(f)
        self.card_ids = m["card_ids"]
        self.metas = m["metas"]

        if self.backend == "gpu" and device is None:
            # gpu 后端必须有确定设备：调用方没传时自动选（与 extractor 同规则）
            try:
                import torch as _t
                device = _t.device("cuda") if _t.cuda.is_available() else None
            except Exception:
                device = None
            if device is None:
                print("[Cascade] ⚠ backend=gpu 但无可用 CUDA，回退 npy", flush=True)
                self.backend = "npy"

        if self.backend == "npy":
            if self.verbose:
                print(f"[Cascade] 加载双区特征库(npy) ...", flush=True)
            self.gallery_upper = np.load(config.GALLERY_UPPER_FEATURES_PATH).astype(np.float32)
            self.gallery_lower = np.load(config.GALLERY_LOWER_FEATURES_PATH).astype(np.float32)
            assert self.gallery_upper.shape[0] == self.gallery_lower.shape[0] == len(self.card_ids), \
                "双区特征库行数与 meta 不一致!"
        elif self.backend == "gpu":
            # 双区特征库常驻 GPU 显存（fp32）：单卡召回 torch 矩阵乘实测 0.4ms，
            # 替代 milvus FLAT 扫描+网络往返的 ~40ms/卡（2026-09-03 bench：Top-5 与
            # milvus 3/3 精确一致；fp16 图库无额外收益且有 tie 抖动，故用 fp32）。
            # 每 worker 显存 +724MB（2×88384×1024×4B），V100 16G 无压力。
            import torch
            if self.verbose:
                print(f"[Cascade] 加载双区特征库(GPU 常驻, device={device}) ...", flush=True)
            self.gallery_upper = torch.from_numpy(
                np.load(config.GALLERY_UPPER_FEATURES_PATH).astype(np.float32)).to(device)
            self.gallery_lower = torch.from_numpy(
                np.load(config.GALLERY_LOWER_FEATURES_PATH).astype(np.float32)).to(device)
            assert self.gallery_upper.shape[0] == self.gallery_lower.shape[0] == len(self.card_ids), \
                "双区特征库行数与 meta 不一致!"
        else:
            alias = milvus_alias or "default"
            if self.verbose:
                print(f"[Cascade] 连接 Milvus {config.MILVUS_HOST}:{config.MILVUS_PORT} "
                      f"(alias={alias}) ...", flush=True)
            from pymilvus import connections, Collection
            if not connections.has_connection(alias):
                connections.connect(alias=alias, host=config.MILVUS_HOST, port=config.MILVUS_PORT)
            self.upper_coll = Collection(config.MILVUS_UPPER_COLLECTION, using=alias)
            self.lower_coll = Collection(config.MILVUS_LOWER_COLLECTION, using=alias)

        if self.verbose:
            print(f"[Cascade] 图库 {len(self.card_ids)} 行  K_recall={self.top_k_recall}  "
                  f"α={self.alpha}  backend={self.backend}", flush=True)

        self.detector = CardDetector(config.YOLO_MODEL_PATH, device=detector_device)
        self.extractor = DualCardFeatureExtractor(
            config.UPPER_MODEL_PATH, config.LOWER_MODEL_PATH, device=device,
        )

    def _features_of(self, img_path, want_label=False, want_crop=False):
        """读图 → 裁卡 → letterbox → (上半区特征, 下半区特征[, 卡牌类型label][, 裁卡RGB])。
        want_label=True 时第三个返回值是 YOLO 检出的最大实例类别名（如 'pokemon'），供 card_type。
        want_crop=True 时额外返回裁卡 RGB numpy 数组（供 element 检测等下游）。"""
        try:
            arr = np.array(Image.open(img_path).convert("RGB"))
        except Exception:
            return (None, None, None, None) if (want_label and want_crop) else \
                   (None, None, None) if want_label else (None, None)
        if want_label:
            crop, card_type = self.detector.detect_and_crop(arr, return_label=True)
        else:
            crop = self.detector.detect_and_crop(arr)
            card_type = None
        if crop is None:
            crop = arr
        lb = letterbox392(Image.fromarray(crop))
        up, lo = split_half(lb)
        q_u = self.extractor.extract_upper([up])[0]
        q_l = self.extractor.extract_lower([lo])[0]
        if np.isnan(q_u).any() or np.isnan(q_l).any():
            return (None, None, None, None) if (want_label and want_crop) else \
                   (None, None, None) if want_label else (None, None)
        ret = (q_u, q_l)
        if want_label:
            ret = ret + (card_type,)
        if want_crop:
            ret = ret + (crop,)
        return ret

    def _query_raw_npy(self, q_u, q_l):
        sims_u = self.gallery_upper @ q_u                 # (N,)
        K = min(self.top_k_recall, sims_u.shape[0])
        cand = np.argpartition(-sims_u, K - 1)[:K]        # 上半区召回 K 个候选 idx
        cand = cand[np.argsort(-sims_u[cand])]            # 按 upper_sim 降序
        su = sims_u[cand]
        sl = self.gallery_lower[cand] @ q_l               # 候选的下半区相似度
        cand_ids = [self.card_ids[int(c)] for c in cand]
        return cand_ids, su, sl

    def _query_raw_milvus(self, q_u, q_l):
        res = self.upper_coll.search(
            data=[q_u.tolist()], anns_field="vector",
            param={"metric_type": "COSINE", "params": {}},
            limit=self.top_k_recall, output_fields=["card_id"],
        )
        cand_ids = [hit.entity.get("card_id") for hit in res[0]]
        su = np.array([hit.distance for hit in res[0]], dtype=np.float32)
        id_list = ",".join(f'"{c}"' for c in cand_ids)
        lower_rows = self.lower_coll.query(expr=f"card_id in [{id_list}]", output_fields=["card_id", "vector"])
        lower_map = {r["card_id"]: np.array(r["vector"], dtype=np.float32) for r in lower_rows}
        sl = np.array([float(np.dot(lower_map[c], q_l)) if c in lower_map else -1.0 for c in cand_ids])
        return cand_ids, su, sl

    def _query_raw_gpu(self, q_u, q_l):
        """GPU 常驻图库召回（backend='gpu'）：与 _query_raw_npy 同数学（归一化向量
        点积=COSINE），只是库在显存、查询向量不过网。2026-09-03 bench：Top-5 与
        milvus 精确一致、单卡 0.4ms（milvus 40.9ms）。"""
        import torch
        dev = self.gallery_upper.device
        q_u_t = torch.from_numpy(q_u).to(dev)
        K = min(self.top_k_recall, self.gallery_upper.shape[0])
        tv = torch.topk(self.gallery_upper @ q_u_t, K)
        cand = tv.indices.cpu().numpy()
        cand_ids = [self.card_ids[int(c)] for c in cand]
        su = tv.values.detach().cpu().numpy().astype(np.float32)
        q_l_t = torch.from_numpy(q_l).to(dev)
        cand_t = torch.from_numpy(cand).to(dev)
        sl = (self.gallery_lower[cand_t] @ q_l_t).detach().cpu().numpy().astype(np.float32)
        return cand_ids, su, sl

    def _query_raw(self, q_u, q_l):
        """按 backend 分发检索。"""
        if self.backend == "npy":
            return self._query_raw_npy(q_u, q_l)
        if self.backend == "gpu":
            return self._query_raw_gpu(q_u, q_l)
        return self._query_raw_milvus(q_u, q_l)

    def query_raw(self, img_path, want_crop=False):
        """返回上半区召回的候选 card_id + 各候选的 upper_sim/lower_sim（供扫 α）。
        额外带 card_type（YOLO 检出的卡牌类型 label）；want_crop=True 时额外带 crop（裁卡RGB）。"""
        # _features_of 返回元组长度随请求变化(2/3/4)，按 want_crop 分别解包
        if want_crop:
            q_u, q_l, card_type, crop = self._features_of(img_path, want_label=True, want_crop=True)
        else:
            q_u, q_l, card_type = self._features_of(img_path, want_label=True, want_crop=False)
            crop = None
        if q_u is None:
            return None
        cand_ids, su, sl = self._query_raw(q_u, q_l)
        ret = {"cand_ids": cand_ids, "upper_sim": su, "lower_sim": sl, "card_type": card_type}
        if want_crop:
            ret["crop"] = crop
        return ret

    def _features_from_crop(self, crop):
        """已裁卡 RGB → (q_u, q_l)。失败返回 (None, None)。"""
        if crop is None:
            return None, None
        lb = letterbox392(Image.fromarray(crop))
        up, lo = split_half(lb)
        q_u = self.extractor.extract_upper([up])[0]
        q_l = self.extractor.extract_lower([lo])[0]
        if np.isnan(q_u).any() or np.isnan(q_l).any():
            return None, None
        return q_u, q_l

    def query_all(self, img_path, alpha=None, top_n=5):
        """多卡查询：YOLO 一次检出全部实例，按阅读顺序各自检索。
        返回 list[dict]，每项结构同 query()（含 crop/box）。未检出返回 []。
        后台 /match 仍走 query() 单卡，本方法仅给小程序 /match_fields。
        特征提取按实例批量（上下半区各合一批前向）：多卡图 GPU 利用率显著更高
        （FP32 bs=2 实测每 half 18.1→11.7ms），批量与逐张前向数值差 ~1e-6 量级，
        不影响检索结果。"""
        a = self.alpha if alpha is None else alpha
        try:
            arr = np.array(Image.open(img_path).convert("RGB"))
        except Exception:
            return []
        instances = self.detector.detect_and_crop_all(arr)
        if not instances:
            return []
        ups, los = [], []
        for inst in instances:
            up, lo = split_half(letterbox392(Image.fromarray(inst["crop"])))
            ups.append(up)
            los.append(lo)
        q_us = self.extractor.extract_upper(ups)      # (N,1024)，失败行 NaN
        q_ls = self.extractor.extract_lower(los)
        out = []
        for i, inst in enumerate(instances):
            q_u, q_l = q_us[i], q_ls[i]
            if q_u is None or q_l is None or np.isnan(q_u).any() or np.isnan(q_l).any():
                continue
            cand_ids, su, sl = self._query_raw(q_u, q_l)
            raw = {"cand_ids": cand_ids, "upper_sim": su, "lower_sim": sl,
                   "card_type": inst.get("label"), "crop": inst["crop"]}
            ranked = self.rank(raw, a)
            if not ranked:
                continue
            out.append({
                "predicted_card_id": ranked[0][0],
                "fusion_score": ranked[0][1],
                "upper_sim": ranked[0][2],
                "lower_sim": ranked[0][3],
                "card_type": inst.get("label"),
                "top_k": [{"card_id": r[0], "fusion": round(r[1], 4),
                           "upper": round(r[2], 4), "lower": round(r[3], 4)}
                          for r in ranked[:top_n]],
                "crop": inst["crop"],
                "box": inst.get("box"),
            })
        return out

    def rank(self, raw, alpha):
        """给定 raw 和 α，返回 [(card_id, fusion, upper_sim, lower_sim), ...] 按 fusion 降序。"""
        cand_ids = raw["cand_ids"]
        fusion = alpha * raw["upper_sim"] + (1 - alpha) * raw["lower_sim"]
        order = np.argsort(-fusion)
        out = []
        for j in order:
            out.append((cand_ids[j], float(fusion[j]),
                        float(raw["upper_sim"][j]), float(raw["lower_sim"][j])))
        return out

    def query(self, img_path, alpha=None, top_n=5, want_crop=False):
        """单次查询（用 self.alpha），返回 dict（含 card_type=YOLO 卡牌类型 label）。
        want_crop=True 时额外返回 crop（裁卡 RGB numpy 数组，供 element 检测等）。"""
        a = self.alpha if alpha is None else alpha
        raw = self.query_raw(img_path, want_crop=want_crop)
        if raw is None:
            return None
        ranked = self.rank(raw, a)
        ret = {
            "predicted_card_id": ranked[0][0],
            "fusion_score": ranked[0][1],
            "upper_sim": ranked[0][2],
            "lower_sim": ranked[0][3],
            "card_type": raw.get("card_type"),
            "top_k": [{"card_id": r[0], "fusion": round(r[1], 4),
                       "upper": round(r[2], 4), "lower": round(r[3], 4)} for r in ranked[:top_n]],
        }
        if want_crop:
            ret["crop"] = raw.get("crop")
        return ret


# ===================== 模式：单图 / 目录 =====================
def mode_images(matcher, paths, alpha):
    for p in paths:
        r = matcher.query(p, alpha=alpha)
        if r is None:
            print(f"[FAIL] {p}", flush=True)
            continue
        print(f"{os.path.basename(p):<40} -> {r['predicted_card_id']}  "
              f"fusion={r['fusion_score']:.3f} (u={r['upper_sim']:.3f} l={r['lower_sim']:.3f})", flush=True)


# ===================== 模式：自检索 sanity =====================
def mode_self_test(matcher, n):
    import csv
    cards = list(csv.DictReader(open(config.CARD_MASTER_ALL_CSV, encoding="utf-8-sig")))
    # 只取有图、且 card_id 在当前(可能 limit)图库里的
    in_lib = set(matcher.card_ids)
    cands = [c for c in cards if c.get("card_id") in in_lib
             and os.path.exists(os.path.join(config.GALLERY_IMG_DIR, md5n(c.get("img_url", ""))))]
    random.seed(42)
    random.shuffle(cands)
    cands = cands[:n]
    print(f"[SelfTest] 抽样 {len(cands)} 张 gallery 图做自检索（Top-1 应命中同 card_id）", flush=True)

    hit = 0
    t0 = time.time()
    for i, c in enumerate(cands, 1):
        img_path = os.path.join(config.GALLERY_IMG_DIR, md5n(c["img_url"]))
        r = matcher.query(img_path)
        ok = r and r["predicted_card_id"] == c["card_id"]
        hit += int(bool(ok))
        if i % 100 == 0 or i == len(cands):
            print(f"  [{i}/{len(cands)}] 命中率 {hit/i:.2%}  ({(time.time()-t0)/i:.2f}s/张)", flush=True)
    report = {"mode": "self_test", "n": len(cands), "hit": hit, "top1_rate": hit / max(1, len(cands))}
    print(f"\n[SelfTest] Top-1 自命中: {hit}/{len(cands)} = {hit/max(1,len(cands)):.2%}", flush=True)
    _save_report(report, "self_test")
    return report


# ===================== 模式：版本消歧 pairs（含 α 扫描）=====================
def mode_pairs(matcher, pairs_json):
    from modules import image_downloader
    pairs = json.load(open(pairs_json, encoding="utf-8"))
    print(f"[Pairs] {len(pairs)} 对，每对查 url_a(普通)+url_b(球闪)，扫 α={ALPHAS}", flush=True)

    # 每个查询缓存 raw（扫 α 复用），记录期望 card_id
    queries = []   # (expect_id, raw_or_None, side, name)
    for p in pairs:
        name = p.get("name", "")
        for side, prefix in [("a", "id_a"), ("b", "id_b")]:
            cid = p.get(prefix)
            url = p.get(f"url_{side}")
            if not cid or not url:
                continue
            local = image_downloader.download_card_image(url, config.QUERY_IMG_DIR)
            raw = matcher.query_raw(local) if local else None
            queries.append((cid, raw, side, name, url))
            tag = "OK" if raw else "FAIL_DL"
            print(f"  [{name}/{side}] expect={cid} {tag}", flush=True)

    valid = [q for q in queries if q[1] is not None]
    print(f"[Pairs] 有效查询 {len(valid)}/{len(queries)}（下载失败 {len(queries)-len(valid)}）", flush=True)

    # 扫 α：对每个 α，统计 Top-1 命中期望、混淆到对家、其他错
    table = []
    per_query_detail = []
    for alpha in ALPHAS:
        hit = confuse = other = 0
        for cid, raw, side, name, url in valid:
            ranked = matcher.rank(raw, alpha)
            pred = ranked[0][0]
            # 对家 card_id：a 的对家是同 pair 的 b
            pair = next(pp for pp in pairs if cid in (pp.get("id_a"), pp.get("id_b")))
            other_id = pair["id_b"] if cid == pair.get("id_a") else pair["id_a"]
            if pred == cid:
                hit += 1
            elif pred == other_id:
                confuse += 1
            else:
                other += 1
            if alpha == matcher.alpha:
                per_query_detail.append({"name": name, "side": side, "expect": cid,
                                         "pred": pred, "other": other_id,
                                         "upper": round(ranked[0][2], 4), "lower": round(ranked[0][3], 4)})
        n = len(valid)
        table.append({"alpha": alpha, "hit": hit, "confuse_a_b": confuse, "other_wrong": other,
                      "hit_rate": round(hit / n, 4) if n else None,
                      "confuse_rate": round(confuse / n, 4) if n else None})
        print(f"  α={alpha:.1f}: 正确 {hit}/{n}={hit/max(1,n):.2%}  混淆a↔b {confuse}  其他错 {other}", flush=True)

    report = {"mode": "pairs", "n_pairs": len(pairs), "n_valid_queries": len(valid),
              "alpha_scan": table, "detail_at_alpha": per_query_detail}
    _save_report(report, "pairs")
    return report


def _save_report(report, tag):
    out = os.path.join(config.DATA_DIR, f"cascade_test_report_{tag}.json")
    with open(out, "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)
    print(f"[Report] 已保存 {out}", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", type=str, default=None)
    ap.add_argument("--image-dir", type=str, default=None)
    ap.add_argument("--self-test", type=int, default=0)
    ap.add_argument("--pairs", type=str, default=None)
    ap.add_argument("--alpha", type=float, default=None)
    ap.add_argument("--top-k-recall", type=int, default=None)
    ap.add_argument("--backend", choices=["npy", "milvus"], default=None,
                     help="检索后端，默认取 config.CASCADE_BACKEND(=npy)")
    args = ap.parse_args()

    config.ensure_dirs()
    matcher = DualCascadeMatcher(alpha=args.alpha, top_k_recall=args.top_k_recall, backend=args.backend)

    if args.self_test:
        mode_self_test(matcher, args.self_test)
    elif args.pairs:
        mode_pairs(matcher, args.pairs)
    elif args.image or args.image_dir:
        paths = [args.image] if args.image else sorted(
            os.path.join(args.image_dir, f) for f in os.listdir(args.image_dir)
            if f.lower().endswith((".jpg", ".jpeg", ".png")))
        mode_images(matcher, paths, args.alpha)
    else:
        ap.error("需指定 --image / --image-dir / --self-test / --pairs 之一")


if __name__ == "__main__":
    main()