#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Run from MaixVision on PC to inspect the box's live 20/40 YOLO output."""

import os

from maix import app, camera, display, image, nn, time


# This file is sent to /tmp/maixpy_run by MaixVision. Keep the model path
# absolute so the script uses the model already installed on the box.
MODEL_PATH = "/maixapp/apps/t11/model_317181.mud"
CAMERA_WIDTH = 640
CAMERA_HEIGHT = 480
CONF_TH = 0.65
IOU_TH = 0.45


def _flatten_objects(raw):
    if raw is None:
        return []
    if isinstance(raw, (list, tuple)):
        result = []
        for item in raw:
            if isinstance(item, (list, tuple)):
                result.extend(_flatten_objects(item))
            else:
                result.append(item)
        return result
    return [raw]


def main():
    if not os.path.isfile(MODEL_PATH):
        raise FileNotFoundError("model not found on box: " + MODEL_PATH)

    detector = nn.YOLOv5(model=MODEL_PATH, dual_buff=False)
    cam = camera.Camera(CAMERA_WIDTH, CAMERA_HEIGHT)
    disp = display.Display()

    labels = tuple(str(label) for label in detector.labels)
    print("[YOLO] model:", MODEL_PATH)
    print("[YOLO] labels:", labels)
    print("[YOLO] conf=%.2f iou=%.2f" % (CONF_TH, IOU_TH))

    fps = 0.0
    frame_count = 0
    last_log_ms = time.ticks_ms()

    while not app.need_exit():
        loop_start_ms = time.ticks_ms()
        img = cam.read()

        detect_start_ms = time.ticks_ms()
        raw = detector.detect(img, conf_th=CONF_TH, iou_th=IOU_TH)
        detect_ms = max(0, time.ticks_diff(time.ticks_ms(), detect_start_ms))
        objects = _flatten_objects(raw)

        candidates = []
        for obj in objects:
            class_id = int(obj.class_id)
            score = float(obj.score)
            label = labels[class_id] if 0 <= class_id < len(labels) else "unknown"
            color = image.COLOR_GREEN if label in ("20", "40") else image.COLOR_RED

            img.draw_rect(obj.x, obj.y, obj.w, obj.h, color=color)
            img.draw_string(
                obj.x,
                max(0, obj.y - 16),
                "%scm %.2f" % (label, score),
                color=color,
            )
            if label in ("20", "40"):
                candidates.append((score, label))

        loop_ms = max(1, time.ticks_diff(time.ticks_ms(), loop_start_ms))
        instant_fps = 1000.0 / float(loop_ms)
        fps = instant_fps if frame_count == 0 else fps * 0.9 + instant_fps * 0.1

        if candidates:
            best_score, best_label = max(candidates, key=lambda item: item[0])
            status = "TARGET %scm %.2f" % (best_label, best_score)
            status_color = image.COLOR_GREEN
        else:
            status = "TARGET UNKNOWN"
            status_color = image.COLOR_RED

        img.draw_string(5, 5, status, color=status_color)
        img.draw_string(
            5,
            25,
            "infer=%dms fps=%.1f boxes=%d" % (detect_ms, fps, len(objects)),
            color=image.COLOR_YELLOW,
        )
        disp.show(img)

        frame_count += 1
        now_ms = time.ticks_ms()
        if time.ticks_diff(now_ms, last_log_ms) >= 1000:
            print(
                "[YOLO] %s infer=%dms fps=%.1f boxes=%d"
                % (status, detect_ms, fps, len(objects))
            )
            last_log_ms = now_ms


if __name__ == "__main__":
    main()