#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""单图多卡批量推理 → XML（含 YOLO 分割框 + 每卡 Top-5 card_id）+ 裁剪图。

与生产 query_all 同逻辑（batch 提特征 + backend 检索 + rank），额外透传 conf。
用法（249）:
  CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=1 \
    ~/miniconda3/envs/pytorch/bin/python batch_multicard_xml.py \
    --images /tmp/mc_batch/in --out /tmp/mc_batch/out
"""
import os
import sys
import time
import argparse
import traceback
from xml.sax.saxutils import escape

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 numpy as np
import cv2
from PIL import Image

import config
import torch
from scripts.cascade_match import DualCascadeMatcher
from modules.feature_extractor_dual import letterbox392, split_half

IMG_EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".webp")


def analyze(matcher, path):
    """单图 → list[card]；与 query_all 同管线，额外带 conf/box/crop。"""
    arr = np.array(Image.open(path).convert("RGB"))
    instances = matcher.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 = matcher.extractor.extract_upper(ups)
    q_ls = matcher.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 = matcher._query_raw(q_u, q_l)
        ranked = matcher.rank({"cand_ids": cand_ids, "upper_sim": su, "lower_sim": sl},
                              matcher.alpha)
        if not ranked:
            continue
        out.append({
            "box": inst["box"], "label": inst.get("label"), "conf": inst.get("conf"),
            "crop": inst["crop"],
            "top_k": [(r[0], round(r[1], 4)) for r in ranked[:5]],
        })
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--images", required=True)
    ap.add_argument("--out", required=True)
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    crops_dir = os.path.join(args.out, "crops")
    os.makedirs(crops_dir, exist_ok=True)

    files = sorted(f for f in os.listdir(args.images)
                   if f.lower().endswith(IMG_EXTS))
    print(f"[Batch] {len(files)} 张图", flush=True)

    matcher = DualCascadeMatcher(device=torch.device("cuda:0"),
                                 detector_device=torch.device("cuda:0"),
                                 milvus_alias="batch")

    xml_path = os.path.join(args.out, "多卡识别结果.xml")
    n_cards_total = 0
    n_err = 0
    t0 = time.time()
    with open(xml_path, "w", encoding="utf-8") as fx:
        fx.write('<?xml version="1.0" encoding="utf-8"?>\n')
        fx.write(f'<results source="full_images" total="{len(files)}">\n')
        for idx, fname in enumerate(files, 1):
            path = os.path.join(args.images, fname)
            stem = os.path.splitext(fname)[0]
            try:
                cards = analyze(matcher, path)
                n_cards_total += len(cards)
            except Exception as e:
                n_err += 1
                print(f"  [ERR] {fname}: {e}", flush=True)
                traceback.print_exc()
                fx.write(f'  <image name="{escape(fname)}" error="{escape(str(e)[:200])}"/>\n')
                continue
            fx.write(f'  <image name="{escape(fname)}" cards="{len(cards)}">\n')
            for ci, c in enumerate(cards, 1):
                # 裁剪图（cv2 质量 100；文件名 = 图名__c序号）
                crop_name = f"{stem}__c{ci}.jpg"
                ok_img = cv2.imwrite(os.path.join(crops_dir, crop_name),
                                     cv2.cvtColor(c["crop"], cv2.COLOR_RGB2BGR),
                                     [int(cv2.IMWRITE_JPEG_QUALITY), 100])
                x1, y1, x2, y2 = c["box"]
                fx.write(f'    <card index="{ci}" crop="{escape(crop_name)}" saved="{int(bool(ok_img))}"'
                         f' card_type="{escape(str(c["label"]))}" conf="{c["conf"]:.3f}">\n')
                fx.write(f'      <box x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}"/>\n')
                fx.write(f'      <top1 card_id="{c["top_k"][0][0]}"/>\n')
                fx.write('      <matches>\n')
                for r, (cid, fus) in enumerate(c["top_k"], 1):
                    fx.write(f'        <match rank="{r}" card_id="{cid}"'
                             f' fusion="{fus}" match_rate="{int(round(fus * 100))}"/>\n')
                fx.write('      </matches>\n')
                fx.write('    </card>\n')
            fx.write('  </image>\n')
            if idx % 20 == 0 or idx == len(files):
                rate = idx / (time.time() - t0)
                print(f"  [{idx}/{len(files)}] {rate:.2f} 张/s  卡数累计 {n_cards_total}", flush=True)
        fx.write('</results>\n')

    print(f"\n[Done] {len(files)} 图 / {n_cards_total} 卡 / {n_err} 错误  "
          f"总耗时 {time.time()-t0:.0f}s", flush=True)
    print(f"[Done] XML: {xml_path}", flush=True)
    print(f"[Done] 裁剪图: {crops_dir}", flush=True)


if __name__ == "__main__":
    main()