"""
集成脚本：card_seg_v2 分割+warp → yolo26s_card_no 检测编号 → 输出可视化+位置
环境：pytorch
输入：IN_DIR（平铺目录，默认 ~/顾工交接/wzj/_infer_in）
输出：OUT_DIR/vis_num（可视化图）+ OUT_DIR/labels_num（YOLO格式，含conf）
用法：
  export IN_DIR=/path/to/images
  export OUT_DIR=/path/to/output
  CUDA_VISIBLE_DEVICES=$UUID ~/miniconda3/envs/pytorch/bin/python flat_card_no.py
"""
import os, sys, json, time, datetime
from pathlib import Path
sys.stdout.reconfigure(encoding="utf-8")
sys.path.insert(0, "/home/martin/顾工交接/wzj")

import cv2
import numpy as np
import config
from modules.yolo_detector import CardDetector
from modules.rectifier import CardRectifier

IN_DIR = os.environ.get("IN_DIR", "/home/martin/顾工交接/wzj/_infer_in")
OUT_DIR = os.environ.get("OUT_DIR", "/home/martin/顾工交接/wzj/_infer_out")
VIS_DIR = os.path.join(OUT_DIR, "vis_num")
LABEL_DIR = os.path.join(OUT_DIR, "labels_num")

os.makedirs(VIS_DIR, exist_ok=True)
os.makedirs(LABEL_DIR, exist_ok=True)

# 模型路径（绝对路径）
SEG_MODEL = "/home/martin/顾工交接/wzj/ultralytics/runs/segment/card_seg_pn_v2/weights/best.pt"
CARD_NO_MODEL = "/home/martin/顾工交接/wzj/ultralytics/runs/detect/yolo26s_card_no/weights/best.pt"

print("[INIT] 加载 card_seg_v2 分割模型 ...")
_t = time.perf_counter()
seg_detector = CardDetector(SEG_MODEL, conf=0.25)
rectifier = CardRectifier(use_perspective=True)
t_seg = time.perf_counter() - _t

print(f"[INIT] 加载 card_no 检测模型 ...")
from modules.ultralytics_compat import import_yolo
YOLO = import_yolo()
_t = time.perf_counter()
card_no_model = YOLO(CARD_NO_MODEL, task="detect")
t_load = time.perf_counter() - _t

print(f"[INIT] card_seg_v2 {t_seg:.2f}s, card_no {t_load:.2f}s\n")

imgs = []
for ext in ("*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"):
    imgs += list(Path(IN_DIR).glob(ext))
imgs = sorted(set(imgs))
print(f"[RUN] {len(imgs)} 张图片\n")

n_seg_ok, n_seg_fail, n_det_hit, n_det_multi, n_det_none = 0, 0, 0, 0, 0
t_all = time.perf_counter()
for i, imgp in enumerate(imgs, 1):
    name = imgp.name
    try:
        img_bgr = cv2.imread(str(imgp))
        if img_bgr is None:
            n_seg_fail += 1
            continue
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        # ① card_seg_v2 分割 + mask4点 warp 摆正
        box, mask, orig_rgb = seg_detector.detect_box_mask(img_rgb)
        if box is None:
            crop_rgb = seg_detector.detect_and_crop(img_rgb)
            if crop_rgb is None:
                n_seg_fail += 1
                continue
        else:
            crop_rgb = rectifier.rectify(orig_rgb, mask=mask, box=box, angle=0)
        if crop_rgb is None or crop_rgb.size == 0:
            n_seg_fail += 1
            continue
        n_seg_ok += 1

        crop_bgr = cv2.cvtColor(crop_rgb, cv2.COLOR_RGB2BGR)

        # ② yolo26s_card_no 检测编号
        res = card_no_model.predict(crop_bgr, conf=0.25, imgsz=640, verbose=False)[0]
        h, w = crop_bgr.shape[:2]

        boxes = res.boxes
        n = 0 if boxes is None else len(boxes)
        if n == 0:
            n_det_none += 1
        elif n == 1:
            n_det_hit += 1
        else:
            n_det_multi += 1

        # ③ 可视化 + YOLO标签
        vis = crop_bgr.copy()
        label_lines = []
        if boxes is not None:
            xyxy = boxes.xyxy.cpu().numpy()
            confs = boxes.conf.cpu().numpy()
            clss = boxes.cls.cpu().numpy()
            for (x1, y1, x2, y2), conf, cls in zip(xyxy, confs, clss):
                cv2.rectangle(vis, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
                cv2.putText(vis, f"card_no {conf:.2f}", (int(x1), max(0, int(y1) - 6)),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, cv2.LINE_AA)
                xc = (x1 + x2) / 2 / w
                yc = (y1 + y2) / 2 / h
                bw = (x2 - x1) / w
                bh = (y2 - y1) / h
                label_lines.append(f"{int(cls)} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f} {conf:.4f}")

        cv2.imwrite(os.path.join(VIS_DIR, name), vis)
        with open(os.path.join(LABEL_DIR, f"{os.path.splitext(name)[0]}.txt"), "w") as f:
            f.write("\n".join(label_lines))

    except Exception as e:
        n_seg_fail += 1
        print(f"  [ERR] {name}: {e}")

    if i % 200 == 0 or i == len(imgs):
        el = time.perf_counter() - t_all
        rate = i / el if el > 0 else 0
        eta = (len(imgs) - i) / rate if rate > 0 else 0
        print(f"  [{i}/{len(imgs)}] 分割ok={n_seg_ok} fail={n_seg_fail} "
              f"检出1框={n_det_hit} 多框={n_det_multi} 无={n_det_none} "
              f"速率={rate:.1f}/s 剩余≈{eta/60:.1f}min", flush=True)

n = max(len(imgs), 1); t_all = time.perf_counter() - t_all
print(f"\n[RESULT] {n} 张")
print(f"  分割成功: {n_seg_ok}, 分割失败: {n_seg_fail}")
print(f"  card_no 检出1框: {n_det_hit}, 多框: {n_det_multi}, 无检出: {n_det_none}")
print("=" * 50)
print(f"⏱ flat_card_no (card_seg_v2 + yolo26s_card_no) | {n} 张")
print(f"  加载 card_seg_v2 {t_seg:.2f}s, card_no {t_load:.2f}s")
print(f"  总耗时 {t_all:.2f}s ({n/t_all:.1f}张/s)")
print("=" * 50)
with open(os.path.join(OUT_DIR, "timing.log"), "a", encoding="utf-8") as lf:
    lf.write(f"\n[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] === flat_card_no.py (card_seg_v2 + yolo26s_card_no) ===\n")
    lf.write(f"  样本 {n} 张\n")
    lf.write(f"  分割成功 {n_seg_ok}, 分割失败 {n_seg_fail}\n")
    lf.write(f"  card_no 检出1框 {n_det_hit}, 多框 {n_det_multi}, 无检出 {n_det_none}\n")
    lf.write(f"  加载 card_seg_v2 {t_seg:.2f}s, card_no {t_load:.2f}s\n")
    lf.write(f"  总耗时 {t_all:.2f}s (均 {t_all/n*1000:.1f} ms/张, 吞吐 {n/t_all:.1f} 张/s)\n")
print(f"✓ 可视化: {VIS_DIR} | 标签: {LABEL_DIR} | 计时→ timing.log")