#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Stage2 黑三角 YOLO —— 在 Maix 设备上用本地图片测试（与线上 target_roi_yolo.try_black_triangle_boxes_work 完全一致）。

不在 PC 上跑 NPU；需把脚本与 config / target_roi_yolo.py 同步到设备，并在设备上执行。

典型用法
--------
  # 输入已是 Stage1 裁切（与你保存的 stage2_roi_*.jpg 一致）
  python test/test_stage2_black_yolo_device.py /root/phot/stage2_roi_xxx.jpg

  # 输入为整幅相机图，手动给出 Stage1 环靶 ROI（与线上日志 ring全图=[rx0,ry0,rx1,ry1] 一致）
  python test/test_stage2_black_yolo_device.py /root/phot/full.jpg --roi 197,196,507,461

  # 对比 native / letterbox 坐标映射（排查 contain 训练与推理对齐）
  python test/test_stage2_black_yolo_device.py ./crop.jpg --compare-coord

  # 覆盖置信度、模型路径（仍读其余项自 config）
  python test/test_stage2_black_yolo_device.py ./crop.jpg --conf 0.25 -m /maixapp/apps/t11/model_270648.mud

  # 只看 NPU 原始框（映射前）：判断坐标是 ~224 网络空间还是归一化 0~1
  python test/test_stage2_black_yolo_device.py ./crop.jpg --conf 0.05 --dump-raw 15

依赖：MaixPy（maix.nn）、OpenCV（cv2）、numpy；项目根须在 sys.path（本脚本已插入上级目录）。
"""

from __future__ import annotations

import argparse
import os
import sys

_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if _ROOT not in sys.path:
    sys.path.insert(0, _ROOT)


def _parse_roi(s: str) -> tuple[int, int, int, int]:
    parts = [p.strip() for p in s.replace(" ", "").split(",")]
    if len(parts) != 4:
        raise ValueError("ROI 需要 4 个整数：x0,y0,x1,y1")
    return tuple(int(x) for x in parts)  # type: ignore[return-value]


def _load_rgb_numpy(path: str) -> "object":
    import cv2
    import numpy as np

    bgr = cv2.imread(path, cv2.IMREAD_COLOR)
    if bgr is None:
        raise FileNotFoundError(f"cv2.imread 失败: {path}")
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    return np.ascontiguousarray(rgb, dtype=np.uint8)


def _draw_boxes_on_crop(
    slab_rgb,
    boxes: list[tuple[int, int, int, int]],
    labels: list[str] | None = None,
):
    """slab_rgb: H×W×3 RGB uint8；boxes 为扩 margin 后的 Stage2 子框（与线上绿框一致）。"""
    import cv2

    vis = slab_rgb.copy()
    bgr = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)
    rh, rw = bgr.shape[:2]
    for i, (bx0, by0, bx1, by1) in enumerate(boxes):
        x0, y0 = int(bx0), int(by0)
        x1, y1 = int(bx1) - 1, int(by1) - 1
        x1 = max(x0, min(x1, rw - 1))
        y1 = max(y0, min(y1, rh - 1))
        cv2.rectangle(bgr, (x0, y0), (x1, y1), (0, 255, 0), 2)
        tag = labels[i] if labels and i < len(labels) else f"s2_{i}"
        cv2.putText(
            bgr,
            tag,
            (x0, max(0, y0 - 4)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 255, 0),
            1,
            cv2.LINE_AA,
        )
    return bgr


class _PrintLogger:
    def info(self, msg):
        print(msg)

    def warning(self, msg):
        print(msg)

    def error(self, msg):
        print(msg)


def _run_once(yroi_mod, img_rgb, roi_xyxy, logger):
    boxes = yroi_mod.try_black_triangle_boxes_work(img_rgb, roi_xyxy, logger)
    rx0, ry0, rx1, ry1 = roi_xyxy
    slab = img_rgb[ry0:ry1, rx0:rx1].copy()
    return boxes, slab


def _copy_dump_raw_rows(yroi_mod, objs):
    """把 Maix detect 返回对象拷贝成基础类型，避免 native 对象跨下一次 detect 存活。"""
    rows = []
    for o in objs:
        cid = yroi_mod._det_obj_class_id(o)
        try:
            sc = float(getattr(o, "score", 0.0))
        except (TypeError, ValueError):
            sc = 0.0
        rows.append((cid, sc, float(o.x), float(o.y), float(o.w), float(o.h)))
    return rows


def _dump_raw_and_hard_exit(det, yroi_mod, slab_for_det, rw_s, rh_s, net_w, net_h, conf_th, iou_th, limit):
    """
    MaixPy 某些版本在 YOLO detect 返回对象正常析构时会 SIGSEGV/pure virtual。
    raw dump 是诊断路径，打印完成后硬退出，绕过 Python/native 析构链。
    """
    from maix import image as maix_image

    roi_maix = maix_image.cv2image(slab_for_det, False, False)
    raw = det.detect(roi_maix, conf_th=conf_th, iou_th=iou_th)
    objs = yroi_mod._normalize_objs(raw if raw is not None else [])
    dump_rows = _copy_dump_raw_rows(yroi_mod, objs)
    raw_count = len(dump_rows)
    print(
        f"[DUMP-RAW] slab={rw_s}×{rh_s} net={net_w}×{net_h} "
        f"conf={conf_th} iou={iou_th} → NMS 后 raw 框数={raw_count}（与 coord_mode 无关）"
    )
    npr = min(int(limit), raw_count)
    for i in range(npr):
        cid, sc, x, y, ww, hh = dump_rows[i]
        print(f"  #{i} cls={cid} score={sc:.4f} xywh=({x:.3f},{y:.3f},{ww:.3f},{hh:.3f})")
    if dump_rows:
        xs = [r[2] for r in dump_rows]
        ws = [r[4] for r in dump_rows]
        print(
            f"[DUMP-RAW]  hint: x 范围≈[{min(xs):.2f},{max(xs):.2f}] "
            f"w 范围≈[{min(ws):.2f},{max(ws):.2f}] — "
            f"若整体在 0~{net_w} 量级多为网络画布坐标→应用 letterbox；"
            f"若 x,w 多在 0~1→可能是归一化，需在代码里乘 net 尺寸"
        )
    print("[INFO] --dump-raw 已完成；为规避 MaixPy YOLO native 析构崩溃，测试进程将直接退出。")
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(0)


def main():
    ap = argparse.ArgumentParser(
        description="Stage2 黑三角 YOLO 设备本地图测试",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    ap.add_argument("image", help="本地图片路径（设备上的路径）")
    ap.add_argument(
        "--roi",
        default="",
        metavar="x0,y0,x1,y1",
        help="可选。若填写：image 为整幅图，在此图上取 Stage1 ROI 再跑 Stage2；"
        "留空：image 本身就是 Stage1 裁切图（默认）",
    )
    ap.add_argument("-o", "--output", default="", help="输出可视化路径；默认 原名_stage2_vis.jpg")
    ap.add_argument("-m", "--model", default="", help="覆盖 config.TRIANGLE_BLACK_YOLO_MODEL_PATH")
    ap.add_argument("--conf", type=float, default=None, help="覆盖 TRIANGLE_BLACK_YOLO_CONF_TH")
    ap.add_argument("--iou", type=float, default=None, help="覆盖 TRIANGLE_BLACK_YOLO_IOU_TH")
    ap.add_argument(
        "--coord",
        choices=["native", "letterbox"],
        default="",
        help="覆盖 TRIANGLE_BLACK_YOLO_COORD_MODE；默认用 config",
    )
    ap.add_argument(
        "--compare-coord",
        action="store_true",
        help="各跑一次 native 与 letterbox，输出两张图 *_stage2_native.jpg / *_stage2_letterbox.jpg",
    )
    ap.add_argument(
        "--fresh-detector",
        action="store_true",
        help="清掉 YOLO 缓存再测（换模型或排查缓存时用）",
    )
    ap.add_argument(
        "--allow-save-roi",
        action="store_true",
        help="不强制关闭 TRIANGLE_BLACK_YOLO_SAVE_ROI_CROP（默认测试时会关掉以免写满相册目录）",
    )
    ap.add_argument(
        "--dump-raw",
        type=int,
        default=0,
        metavar="N",
        help="打印前 N 个 detect 原始框 x,y,w,h,score,cls（coord 映射前；native/letterbox 共用同一批 raw）",
    )
    args = ap.parse_args()

    img_path = os.path.abspath(args.image)
    if not os.path.isfile(img_path):
        print(f"[ERR] 找不到图片: {img_path}")
        sys.exit(1)

    try:
        import config as cfg
        import target_roi_yolo as yroi
    except ImportError as e:
        print(f"[ERR] 无法导入 config / target_roi_yolo: {e}")
        sys.exit(1)

    if args.fresh_detector:
        yroi.reset_yolo_detector_cache()

    # 备份并临时覆盖 config（单进程顺序跑）
    bak: dict[str, object] = {}

    def _patch(key: str, val: object):
        if key not in bak:
            bak[key] = getattr(cfg, key, None)
        setattr(cfg, key, val)

    def _restore():
        for k, v in bak.items():
            setattr(cfg, k, v)

    try:
        _patch("TRIANGLE_BLACK_YOLO_ENABLE", True)
        if not args.allow_save_roi:
            _patch("TRIANGLE_BLACK_YOLO_SAVE_ROI_CROP", False)
        if args.model.strip():
            _patch("TRIANGLE_BLACK_YOLO_MODEL_PATH", args.model.strip())
        if args.conf is not None:
            _patch("TRIANGLE_BLACK_YOLO_CONF_TH", float(args.conf))
        if args.iou is not None:
            _patch("TRIANGLE_BLACK_YOLO_IOU_TH", float(args.iou))
        if args.coord and not args.compare_coord:
            _patch("TRIANGLE_BLACK_YOLO_COORD_MODE", args.coord)

        mp = getattr(cfg, "TRIANGLE_BLACK_YOLO_MODEL_PATH", "") or ""
        if not os.path.isfile(mp):
            print(f"[ERR] 模型文件不存在: {mp}")
            sys.exit(1)

        img_rgb = _load_rgb_numpy(img_path)
        h, w = int(img_rgb.shape[0]), int(img_rgb.shape[1])

        if args.roi.strip():
            roi_xyxy = _parse_roi(args.roi.strip())
            rx0, ry0, rx1, ry1 = [int(round(float(v))) for v in roi_xyxy]
            if rx1 <= rx0 or ry1 <= ry0:
                print("[ERR] ROI 无效：需满足 x1>x0 且 y1>y0")
                sys.exit(1)
            # 与 target_roi_yolo.try_black_triangle_boxes_work 相同的 clip
            rx0 = max(0, min(rx0, w - 1))
            ry0 = max(0, min(ry0, h - 1))
            rx1 = max(rx0 + 1, min(rx1, w))
            ry1 = max(ry0 + 1, min(ry1, h))
            ring_roi = (rx0, ry0, rx1, ry1)
            print(f"[INFO] 模式=整图+ROI ring={ring_roi} image={w}×{h}")
        else:
            ring_roi = (0, 0, w, h)
            print(f"[INFO] 模式=已是 Stage1 裁切 crop={w}×{h}")

        logger = _PrintLogger()
        det = yroi._get_detector(mp)
        if det is None:
            print("[ERR] 无法加载 nn.YOLOv5（检查模型路径与 Maix 环境）")
            sys.exit(1)
        net_w = int(det.input_width())
        net_h = int(det.input_height())
        print(f"[INFO] model={mp} net_in={net_w}×{net_h}")

        rx0, ry0, rx1, ry1 = ring_roi
        import numpy as np

        slab_for_det = np.ascontiguousarray(img_rgb[ry0:ry1, rx0:rx1], dtype=np.uint8).copy()
        rh_s, rw_s = int(slab_for_det.shape[0]), int(slab_for_det.shape[1])

        modes = ["native", "letterbox"] if args.compare_coord else [
            (args.coord or getattr(cfg, "TRIANGLE_BLACK_YOLO_COORD_MODE", "native"))
        ]

        base, ext = os.path.splitext(img_path)
        ext = ext if ext else ".jpg"

        for mode in modes:
            _patch("TRIANGLE_BLACK_YOLO_COORD_MODE", mode)
            cur_coord = getattr(cfg, "TRIANGLE_BLACK_YOLO_COORD_MODE", mode)
            print(f"[INFO] --- TRIANGLE_BLACK_YOLO_COORD_MODE={cur_coord} ---")

            boxes, slab = _run_once(yroi, img_rgb, ring_roi, logger)
            print(
                f"[INFO] 子框数量={len(boxes)} conf={getattr(cfg, 'TRIANGLE_BLACK_YOLO_CONF_TH', '?')} "
                f"coord={cur_coord}"
            )
            for i, b in enumerate(boxes):
                print(f"  s2_{i}: {b}")

            if args.compare_coord:
                out_path = f"{base}_stage2_{mode}{ext}"
            elif args.output.strip():
                out_path = args.output.strip()
            else:
                out_path = base + "_stage2_vis" + ext

            import cv2

            bgr = _draw_boxes_on_crop(slab, boxes)
            cv2.imwrite(out_path, bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 92])
            print(f"[OK] saved: {out_path}")

        if args.compare_coord:
            print(
                "[HINT] contain 训练时若 letterbox 对齐更好，请将 config 里 "
                "TRIANGLE_BLACK_YOLO_COORD_MODE 设为 letterbox"
            )

        if args.dump_raw > 0:
            conf_th = float(getattr(cfg, "TRIANGLE_BLACK_YOLO_CONF_TH", 0.5))
            iou_th = float(getattr(cfg, "TRIANGLE_BLACK_YOLO_IOU_TH", 0.45))
            print("\n[INFO] --dump-raw 放在最后执行，避免 raw native 对象影响 compare-coord 流程。")
            _dump_raw_and_hard_exit(
                det,
                yroi,
                slab_for_det,
                rw_s,
                rh_s,
                net_w,
                net_h,
                conf_th,
                iou_th,
                args.dump_raw,
            )

    finally:
        _restore()


if __name__ == "__main__":
    main()
