#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
摄像头实时 YOLOv5 简易测试脚本。

特点：
- 完全独立脚本，直接 python test/test_yolo_camera_simple.py 运行，不需要传参。
- 不 import config，不依赖项目模块。
- 直接调用 maix.nn.YOLOv5(model=..., dual_buff=False)。
- camera.read() 得到的 Maix image 直接送 det.detect()。
- 在画面上画检测框、类别、置信度，并显示到屏幕。

运行环境：MaixCAM / MaixPy。
"""

import os


CAMERA_WIDTH = 640
CAMERA_HEIGHT = 480
# 默认与主项目 config.TRIANGLE_YOLO_MODEL_PATH 一致（勿用 /root/yolo26_int8.mud，那是占位路径）
_MODEL_DEFAULT = "/maixapp/apps/t11/model_270139.mud"
try:
    import config as _cfg

    MODEL_PATH = getattr(_cfg, "TRIANGLE_YOLO_MODEL_PATH", _MODEL_DEFAULT) or _MODEL_DEFAULT
except Exception:
    MODEL_PATH = _MODEL_DEFAULT
CONF_TH = 0.7
IOU_TH = 0.45
# native: Maix detect 返回框已映射到 camera.read() 图像坐标；letterbox: 需要从网络输入坐标反算
COORD_MODE = "native"
# 只用于 DRAW_ONLY_CLASS_IDS=True 时过滤显示；默认画所有框
CLASS_IDS = (0,)
DRAW_ONLY_CLASS_IDS = False  # True=只画 CLASS_IDS 里的类别；False=画所有 YOLO 返回框


def _det_obj_class_id(o):
    for key in ("class_id", "cls", "label", "category", "cat_id", "id"):
        if hasattr(o, key):
            v = getattr(o, key)
            if v is None:
                continue
            try:
                return int(float(v))
            except (TypeError, ValueError):
                continue
    return None


def _det_obj_from_seq(t):
    if not isinstance(t, (list, tuple)) or len(t) < 6:
        return None

    class Box:
        pass

    b = Box()
    b.x = float(t[0])
    b.y = float(t[1])
    b.w = float(t[2])
    b.h = float(t[3])
    b.score = float(t[4])
    b.class_id = int(float(t[5]))
    return b


def _normalize_objs(objs):
    out = []
    for o in objs or []:
        if isinstance(o, (list, tuple)):
            m = _det_obj_from_seq(o)
            if m is not None:
                out.append(m)
        else:
            out.append(o)
    return out


def _letterbox_net_to_src_xyxy(x, y, w, h, src_w, src_h, net_w, net_h):
    scale = min(net_w / float(src_w), net_h / float(src_h))
    new_w = src_w * scale
    new_h = src_h * scale
    pad_x = (net_w - new_w) * 0.5
    pad_y = (net_h - new_h) * 0.5
    x0 = (x - pad_x) / scale
    y0 = (y - pad_y) / scale
    x1 = (x + w - pad_x) / scale
    y1 = (y + h - pad_y) / scale
    return x0, y0, x1, y1


def _det_to_src_xyxy(o, coord_mode, src_w, src_h, net_w, net_h):
    x = float(getattr(o, "x", 0.0))
    y = float(getattr(o, "y", 0.0))
    w = float(getattr(o, "w", 0.0))
    h = float(getattr(o, "h", 0.0))
    if coord_mode in ("native", "source", "camera", "full"):
        return x, y, x + w, y + h
    return _letterbox_net_to_src_xyxy(x, y, w, h, src_w, src_h, net_w, net_h)


def _clip_xywh(x0, y0, x1, y1, src_w, src_h):
    x0 = max(0, min(int(round(x0)), src_w - 1))
    y0 = max(0, min(int(round(y0)), src_h - 1))
    x1 = max(x0 + 1, min(int(round(x1)), src_w))
    y1 = max(y0 + 1, min(int(round(y1)), src_h))
    return x0, y0, x1 - x0, y1 - y0


def _label(det, cid):
    labels = getattr(det, "labels", None)
    if labels is None:
        return str(cid)
    try:
        return str(labels[int(cid)])
    except Exception:
        return str(cid)


def main():
    from maix import camera, display, nn, time, image

    if not MODEL_PATH or not os.path.isfile(MODEL_PATH):
        print("[ERR] 模型文件不存在:", MODEL_PATH)
        return

    print("[INFO] 初始化 YOLO 模型:", MODEL_PATH)
    det = nn.YOLOv26(model=MODEL_PATH, dual_buff=False)
    net_w = int(det.input_width())
    net_h = int(det.input_height())
    print(
        "[INFO] net_in=%dx%d conf=%.2f iou=%.2f coord=%s class_ids=%s"
        % (net_w, net_h, CONF_TH, IOU_TH, COORD_MODE, str(CLASS_IDS))
    )

    print("[INFO] 初始化摄像头: %dx%d" % (CAMERA_WIDTH, CAMERA_HEIGHT))
    cam = camera.Camera(CAMERA_WIDTH, CAMERA_HEIGHT)
    disp = display.Display()

    color_cycle = []
    for name in ("RED", "GREEN", "BLUE", "ORANGE", "YELLOW", "CYAN", "MAGENTA"):
        c = getattr(image, "COLOR_" + name, None)
        if c is not None:
            color_cycle.append(c)
    if not color_cycle:
        color_cycle = [getattr(image, "COLOR_RED", 0)]

    frame_idx = 0
    last_log_ms = time.ticks_ms()
    fps_count = 0

    while True:
        frame = cam.read()
        src_w = frame.width()
        src_h = frame.height()

        t0 = time.ticks_ms()
        raw = det.detect(frame, conf_th=CONF_TH, iou_th=IOU_TH)
        detect_ms = time.ticks_ms() - t0
        objs = _normalize_objs(raw if raw is not None else [])

        draw_count = 0
        for i, o in enumerate(objs):
            cid = _det_obj_class_id(o)
            if cid is None:
                cid = -1
            if DRAW_ONLY_CLASS_IDS and cid not in CLASS_IDS:
                continue
            try:
                score = float(getattr(o, "score", 0.0))
            except Exception:
                score = 0.0

            x0, y0, x1, y1 = _det_to_src_xyxy(o, COORD_MODE, src_w, src_h, net_w, net_h)
            ix, iy, iw, ih = _clip_xywh(x0, y0, x1, y1, src_w, src_h)
            col = color_cycle[cid % len(color_cycle)] if cid >= 0 else color_cycle[0]
            frame.draw_rect(ix, iy, iw, ih, color=col)
            frame.draw_string(ix, max(0, iy - 16), "%s %.2f" % (_label(det, cid), score), color=col)
            draw_count += 1

        frame.draw_string(4, 4, "YOLO boxes:%d draw:%d %dms" % (len(objs), draw_count, detect_ms), color=color_cycle[0])
        disp.show(frame)

        frame_idx += 1
        fps_count += 1
        now = time.ticks_ms()
        if now - last_log_ms >= 1000:
            print(
                "[INFO] frame=%d fps=%d raw_boxes=%d draw_boxes=%d detect_ms=%d"
                % (frame_idx, fps_count, len(objs), draw_count, detect_ms)
            )
            fps_count = 0
            last_log_ms = now


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("[INFO] exit")
    except Exception as e:
        print("[ERR]", e)
        try:
            import traceback

            traceback.print_exc()
        except Exception:
            pass