# -*- coding: utf-8 -*-
"""build_dual_gallery_88384.py - 新卡池(88,384张)双区特征库构建

与 pokemon/build_dual_gallery.py 完全同源的预处理流程（YOLO裁卡 → letterbox392 →
split_half → DualCardFeatureExtractor），**唯一差异**是图片定位方式：
  - 原脚本: card_master_all.csv 的 img_url → md5(url).jpg
  - 本脚本: train_meta_88407.json 的 fname（= img_url 尾段原始文件名）

原脚本一字未改，本文件为新增，互不影响。

输出（写到独立目录，不覆盖任何现有文件）:
  data/gallery_v819/gallery_upper_features.npy
  data/gallery_v819/gallery_lower_features.npy
  data/gallery_v819/gallery_dual_meta.json

用法(pytorch 环境, GPU):
  python build_dual_gallery_88384.py                # 全量
  python build_dual_gallery_88384.py --limit 500    # 冒烟
"""
import os
import sys
import json
import time

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

_THIS = os.path.dirname(os.path.abspath(__file__))
_ROOT = _THIS if os.path.exists(os.path.join(_THIS, "config.py")) else os.path.dirname(_THIS)
sys.path.insert(0, _ROOT)

import argparse
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,
)

META_IN = os.path.join(_ROOT, "data", "train_meta_v0904.json")
GI = config.GALLERY_IMG_DIR
OUT_DIR = os.path.join(_ROOT, "data", "gallery_v0904")
OUT_U = os.path.join(OUT_DIR, "gallery_upper_features.npy")
OUT_L = os.path.join(OUT_DIR, "gallery_lower_features.npy")
OUT_M = os.path.join(OUT_DIR, "gallery_dual_meta.json")
BATCH = 64


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--batch", type=int, default=BATCH)
    args = ap.parse_args()

    os.makedirs(OUT_DIR, exist_ok=True)
    print(f"[INIT] META={META_IN}", flush=True)
    print(f"[INIT] GALLERY_IMG_DIR={GI}", flush=True)
    print(f"[INIT] UPPER_MODEL={'/home/user/顾工交接/wzj/upper_model_output_v0904/best_upper_half_model.pth'}", flush=True)
    print(f"[INIT] LOWER_MODEL={'/home/user/顾工交接/wzj/layer3_bg_model_output_v0904/best_layer3_bottom_model.pth'}", flush=True)
    print(f"[INIT] OUT={OUT_DIR}", flush=True)

    detector = CardDetector(config.YOLO_MODEL_PATH)
    extractor = DualCardFeatureExtractor(
        "/home/user/顾工交接/wzj/upper_model_output_v0904/best_upper_half_model.pth", "/home/user/顾工交接/wzj/layer3_bg_model_output_v0904/best_layer3_bottom_model.pth", batch_size=args.batch,
    )

    cards = json.load(open(META_IN, encoding="utf-8"))["records"]
    if args.limit:
        cards = cards[: args.limit]
    print(f"[INIT] META: {len(cards)} 条待处理\n", flush=True)

    feats_u, feats_l, metas = [], [], []
    buf_u, buf_l, buf_meta = [], [], []
    n_done = n_noimg = n_fail = 0
    t0 = time.time()

    def flush():
        if not buf_u:
            return
        fu = extractor.extract_upper(buf_u)
        fl = extractor.extract_lower(buf_l)
        for i, m in enumerate(buf_meta):
            if np.isnan(fu[i]).any() or np.isnan(fl[i]).any():
                continue
            feats_u.append(fu[i])
            feats_l.append(fl[i])
            metas.append(m)
        buf_u.clear()
        buf_l.clear()
        buf_meta.clear()

    for card in cards:
        n_done += 1
        img_path = os.path.join(GI, card["fname"])
        if not os.path.exists(img_path):
            n_noimg += 1
            continue
        try:
            arr = np.array(Image.open(img_path).convert("RGB"))
            crop = detector.detect_and_crop(arr)
            if crop is None:
                crop = arr
            lb = letterbox392(Image.fromarray(crop))
            up, lo = split_half(lb)
        except Exception as e:
            n_fail += 1
            if n_fail <= 10:
                print(f"  [FAIL] {card['fname']}: {e}", flush=True)
            continue

        buf_u.append(up)
        buf_l.append(lo)
        buf_meta.append({
            "card_id": card.get("card_id", ""),
            "card_name_ch": card.get("card_name_ch", ""),
            "language": card.get("language", ""),
            "year": card.get("year", ""),
            "card_no": card.get("card_no", ""),
            "pg_label": card.get("pg_label", ""),
        })
        if len(buf_u) >= args.batch:
            flush()
        if n_done % 1000 == 0:
            flush()
            print(f"  [{n_done}/{len(cards)}] 有效{len(feats_u)} 无图{n_noimg} 失败{n_fail}  "
                  f"({(time.time()-t0)/60:.1f}min)", flush=True)

    flush()

    feats_u_arr = np.array(feats_u, dtype=np.float32)
    feats_l_arr = np.array(feats_l, dtype=np.float32)
    np.save(OUT_U, feats_u_arr)
    np.save(OUT_L, feats_l_arr)
    with open(OUT_M, "w", encoding="utf-8") as f:
        json.dump({"card_ids": [m["card_id"] for m in metas], "metas": metas},
                  f, ensure_ascii=False)

    print(f"\n[DONE] upper={feats_u_arr.shape} lower={feats_l_arr.shape} "
          f"meta={len(metas)} 无图{n_noimg} 失败{n_fail} "
          f"耗时{(time.time()-t0)/60:.1f}min", flush=True)
    print(f"[DONE] -> {OUT_DIR}", flush=True)


if __name__ == "__main__":
    main()