#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""gpu 后端 A/B 等价测试：同一进程先后用 milvus / gpu 后端跑同一批图，
输出（predicted_card_id + Top-1 的 id/fusion/upper/lower）必须逐字段一致。"""
import os
import sys

sys.stdout.reconfigure(encoding="utf-8")
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, _ROOT)
sys.path.insert(0, os.path.join(_ROOT, "scripts"))

import torch
from scripts.cascade_match import DualCascadeMatcher

IMGS = [
    "data/audit_imgs/multicard_2up.jpg",
    "data/audit_imgs/raspi_front_2026-09-01_13-22-48_529284.jpg",
    "data/audit_imgs/raspi_front_2026-09-01_11-24-01_965039.jpg",
]

def run(backend):
    dev = torch.device("cuda:0")
    m = DualCascadeMatcher(backend=backend, device=dev, detector_device=dev,
                           milvus_alias="ab_" + backend)
    out = {}
    for p in IMGS:
        path = p if os.path.isabs(p) else os.path.join(_ROOT, p)
        r = m.query_all(path)
        out[p] = [(x["predicted_card_id"],
                   [(t["card_id"], t["fusion"], t["upper"], t["lower"]) for t in x["top_k"]])
                  for x in r]
        print(f"[{backend}] {os.path.basename(p)}: {out[p]}", flush=True)
    del m
    torch.cuda.empty_cache()
    return out

a = run("milvus")
b = run("gpu")
same = a == b
print(f"\n[A/B] milvus vs gpu 输出全等: {same}", flush=True)
if not same:
    for p in a:
        if a[p] != b.get(p):
            print(f"  差异 {p}:\n    milvus={a[p]}\n    gpu   ={b.get(p)}", flush=True)
sys.exit(0 if same else 1)