from maix import camera, display, image, nn, app, uart, pinmap, time, sys, err

# ---------------------- 1. Hardware & UART Config ----------------------
device_id = sys.device_id()
if device_id == "maixcam2":
    pin_function = {"A21": "UART4_TX", "A22": "UART4_RX"}
    device = "/dev/ttyS4"
else:
    pin_function = {"A16": "UART0_TX", "A17": "UART0_RX"}
    device = "/dev/ttyS0"

# Set Pin mapping
for pin, func in pin_function.items():
    err.check_raise(pinmap.set_pin_function(pin, func), f"Failed to set {pin} to {func}")

# Init UART with 115200 baud rate
serial_dev = uart.UART(device, 115200)

# ---------------------- 2. AI Model & Camera Init ----------------------
# Resolution is fixed at 320x224 for optimal inference speed
detector = nn.YOLOv5(model="/root/models/yolov5s.mud", dual_buff=True)
cam = camera.Camera(detector.input_width(), detector.input_height(), detector.input_format())
disp = display.Display()

print("Vision System Initialized. Resolution: 320x224")

# ---------------------- 3. Main Logic Loop ----------------------
while not app.need_exit():
    img = cam.read()
    objs = detector.detect(img, conf_th=0.5, iou_th=0.45)
    
    target_person = None
    max_area = 0

    # Find the largest 'person' in the frame to lock as the target
    for obj in objs:
        if detector.labels[obj.class_id] == "person":
            area = obj.w * obj.h
            if area > max_area:
                max_area = area
                target_person = obj

    # Process the locked target
    if target_person:
        obj = target_person
        center_x = int(obj.x + obj.w / 2)
        center_y = int(obj.y + obj.h / 2)
        
        # UI Feedback: Draw rectangle and center point
        img.draw_rect(obj.x, obj.y, obj.w, obj.h, color=image.COLOR_RED)
        img.draw_string(obj.x, obj.y, f"TARGET: {obj.score:.2f}", color=image.COLOR_RED)
        img.draw_string(center_x, center_y, "+", color=image.COLOR_GREEN)

        # Logic for proximity and movement
        # If the person is very close (large width/height), send STOP command
        if obj.w > 300 or obj.h > 200:
            serial_dev.write_str("stop\n")
            img.draw_string(10, 10, "STATUS: STOP (TOO CLOSE)", color=image.COLOR_RED)
        else:
            # Send coordinates to robot (Format matched with robot's parser)
            output_msg = f"x:{center_x},y:{center_y}\n"
            serial_dev.write_str(output_msg)
            img.draw_string(10, 10, f"STATUS: TRACKING X={center_x}", color=image.COLOR_GREEN)
    else:
        # Optional: You could send a "lost" signal here if needed
        pass

    disp.show(img)
    # Control loop frequency slightly to prevent UART flooding (approx 20Hz)
    time.sleep