import depthai as dai
import numpy as np
import cv2
import time
from scipy.ndimage import label

# 📏 Definições da resolução da câmera
WIDTH, HEIGHT = 640, 400
DEPTH_THRESHOLD = 300  # 🔹 Diferença mínima de profundidade para considerar obstáculo (mm)
CAMERA_HEIGHT_CM = 55  # 🔹 Altura da câmera em cm

# Criando o pipeline
pipeline = dai.Pipeline()

# 📷 Criar nó da câmera RGB
cam_rgb = pipeline.create(dai.node.ColorCamera)
cam_rgb.setPreviewSize(WIDTH, HEIGHT)
cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)
cam_rgb.setInterleaved(False)

xout_rgb = pipeline.create(dai.node.XLinkOut)
xout_rgb.setStreamName("rgb")
cam_rgb.preview.link(xout_rgb.input)

# 📏 Criar nó de profundidade
mono_left = pipeline.create(dai.node.MonoCamera)
mono_right = pipeline.create(dai.node.MonoCamera)
stereo = pipeline.create(dai.node.StereoDepth)

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

mono_left.setBoardSocket(dai.CameraBoardSocket.LEFT)
mono_right.setBoardSocket(dai.CameraBoardSocket.RIGHT)

stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
stereo.setLeftRightCheck(True)
stereo.setSubpixel(True)

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

xout_depth = pipeline.create(dai.node.XLinkOut)
xout_depth.setStreamName("depth")
stereo.depth.link(xout_depth.input)

# 🚀 Inicializar Dispositivo
with dai.Device(pipeline) as device:
    rgb_queue = device.getOutputQueue(name="rgb", maxSize=1, blocking=False)
    depth_queue = device.getOutputQueue(name="depth", maxSize=1, blocking=False)

    last_time = time.time()
    target_fps = 15  # 🔹 Limita o FPS para evitar travamentos

    while True:
        # 🔹 Controle de FPS
        current_time = time.time()
        if current_time - last_time < 1 / target_fps:
            time.sleep(0.01)
            continue
        last_time = current_time

        in_rgb = rgb_queue.tryGet()
        in_depth = depth_queue.tryGet()

        if in_rgb is None or in_depth is None:
            continue

        rgb_frame = in_rgb.getCvFrame()
        depth_frame = in_depth.getFrame()

        # 🔹 Normaliza profundidade para análise
        depth_visual = cv2.normalize(depth_frame, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        depth_visual = cv2.applyColorMap(depth_visual, cv2.COLORMAP_JET)

        # 🔹 Criar máscara para identificar variações significativas de profundidade
        median_depth = np.median(depth_frame)  # 🔹 Calcula a profundidade média do ambiente
        depth_mask = np.abs(depth_frame - median_depth) > DEPTH_THRESHOLD  # 🔹 Detecta regiões diferentes do fundo

        # 🔹 Agrupar pixels próximos em objetos distintos
        labeled_mask, num_features = label(depth_mask)

        detected_objects = []

        for obj_id in range(1, num_features + 1):
            y, x = np.where(labeled_mask == obj_id)
            if len(x) < 200:  # 🔹 Ignorar detecções muito pequenas (ruído)
                continue

            x_min, x_max = x.min(), x.max()
            y_min, y_max = y.min(), y.max()
            width_px = x_max - x_min
            height_px = y_max - y_min

            # 🔹 Calcula distância média do objeto
            object_depth_values = depth_frame[y, x]
            object_depth_values = object_depth_values[object_depth_values > 0]  # 🔹 Remove valores inválidos
            avg_depth = np.mean(object_depth_values) if len(object_depth_values) > 0 else 0

            # 🔹 Conversão aproximada de pixels para cm
            distance_cm = avg_depth / 10  # 🔹 Supondo 1 mm por unidade
            width_cm = (width_px / WIDTH) * (distance_cm / CAMERA_HEIGHT_CM) * 100
            height_cm = (height_px / HEIGHT) * (distance_cm / CAMERA_HEIGHT_CM) * 100

            # 🔹 Filtrar objetos irrelevantes
            if distance_cm > 500 or height_cm < 10:  # 🔹 Ignorar objetos muito longe ou pequenos
                continue

            # 🔹 Adicionar aos obstáculos detectados
            detected_objects.append({
                "x_min": x_min,
                "y_min": y_min,
                "width_px": width_px,
                "height_px": height_px,
                "width_cm": round(width_cm, 2),
                "height_cm": round(height_cm, 2),
                "distance_cm": round(distance_cm, 2),
            })

            # 🔹 Desenhar bounding box na imagem RGB
            cv2.rectangle(rgb_frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
            cv2.putText(rgb_frame, f"{distance_cm:.1f} cm", (x_min, y_min - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # 🚀 Exibir imagem processada
        cv2.imshow("Detecção de Obstáculos - OAK-D Lite", rgb_frame)

        if cv2.waitKey(1) == ord('q'):
            break

    cv2.destroyAllWindows()