import cv2
import depthai as dai
import blobconverter
import time

# COCO labels do MobileNet-SSD
LABELS = [
    "background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
    "car", "cat", "chair", "cow", "diningtable", "dog", "horse",
    "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"
]

pipeline = dai.Pipeline()

# RGB
cam_rgb = pipeline.create(dai.node.ColorCamera)
cam_rgb.setPreviewSize(300, 300)
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
cam_rgb.setInterleaved(False)
cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
cam_rgb.setFps(30)

# Mono stereo
mono_left = pipeline.create(dai.node.MonoCamera)
mono_right = pipeline.create(dai.node.MonoCamera)

mono_left.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
mono_right.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)

mono_left.setBoardSocket(dai.CameraBoardSocket.CAM_B)
mono_right.setBoardSocket(dai.CameraBoardSocket.CAM_C)

# Depth
stereo = pipeline.create(dai.node.StereoDepth)
stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.DEFAULT)
stereo.setDepthAlign(dai.CameraBoardSocket.CAM_A)
stereo.setSubpixel(True)

mono_left.out.link(stereo.left)
mono_right.out.link(stereo.right)

# Spatial Detection Network - MobileNet SSD
detection = pipeline.create(dai.node.MobileNetSpatialDetectionNetwork)

detection.setBlobPath(blobconverter.from_zoo(
    name="mobilenet-ssd",
    shaves=3,
    version="2021.4"
))

detection.setConfidenceThreshold(0.5)
detection.input.setBlocking(False)
detection.setBoundingBoxScaleFactor(0.5)
detection.setDepthLowerThreshold(200)
detection.setDepthUpperThreshold(8000)

cam_rgb.preview.link(detection.input)
stereo.depth.link(detection.inputDepth)

# Tracker
tracker = pipeline.create(dai.node.ObjectTracker)

# Para obstáculo geral, rastreia tudo detectado.
# Para só pessoas: tracker.setDetectionLabelsToTrack([15])
tracker.setTrackerType(dai.TrackerType.ZERO_TERM_COLOR_HISTOGRAM)
tracker.setTrackerIdAssignmentPolicy(dai.TrackerIdAssignmentPolicy.SMALLEST_ID)

detection.passthrough.link(tracker.inputTrackerFrame)
detection.passthrough.link(tracker.inputDetectionFrame)
detection.out.link(tracker.inputDetections)

# Outputs
xout_rgb = pipeline.create(dai.node.XLinkOut)
xout_track = pipeline.create(dai.node.XLinkOut)

xout_rgb.setStreamName("rgb")
xout_track.setStreamName("tracklets")

tracker.passthroughTrackerFrame.link(xout_rgb.input)
tracker.out.link(xout_track.input)

with dai.Device(pipeline) as device:
    q_rgb = device.getOutputQueue("rgb", maxSize=4, blocking=False)
    q_track = device.getOutputQueue("tracklets", maxSize=4, blocking=False)

    last = time.time()
    fps = 0

    while True:
        frame = q_rgb.get().getCvFrame()
        tracklets = q_track.get().tracklets

        now = time.time()
        fps = 0.9 * fps + 0.1 * (1 / max(now - last, 1e-6))
        last = now

        h, w = frame.shape[:2]

        perigo = False

        for t in tracklets:
            roi = t.roi.denormalize(w, h)
            x1 = int(roi.topLeft().x)
            y1 = int(roi.topLeft().y)
            x2 = int(roi.bottomRight().x)
            y2 = int(roi.bottomRight().y)

            label = LABELS[t.label] if t.label < len(LABELS) else str(t.label)

            x_mm = t.spatialCoordinates.x
            y_mm = t.spatialCoordinates.y
            z_mm = t.spatialCoordinates.z

            dist_m = z_mm / 1000.0

            if dist_m < 1.5:
                perigo = True

            texto = f"ID {t.id} | {label} | {t.status.name} | Z={dist_m:.2f}m"

            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, texto, (x1, max(20, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1)

            cv2.circle(frame, ((x1 + x2) // 2, (y1 + y2) // 2), 4, (0, 255, 255), -1)

        cv2.putText(frame, f"FPS: {fps:.1f}", (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        if perigo:
            cv2.putText(frame, "PERIGO: objeto perto - PARAR", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)

        cv2.imshow("OAK-D Lite Spatial Object Tracker", frame)

        key = cv2.waitKey(1)
        if key == ord("q") or key == 27:
            break

cv2.destroyAllWindows()