import depthai as dai
import cv2
import numpy as np

# Criação do pipeline
pipeline = dai.Pipeline()

# Câmera RGB
cam_rgb = pipeline.create(dai.node.ColorCamera)
cam_rgb.setPreviewSize(1280, 720)  # Resolução HD
cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)
cam_rgb.setInterleaved(False)

# Criar XLinkOut para a imagem RGB
xout_rgb = pipeline.create(dai.node.XLinkOut)
xout_rgb.setStreamName("rgb")
cam_rgb.preview.link(xout_rgb.input)

# Câmera 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)

# Configuração do StereoDepth
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)

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

# Inicia o pipeline
with dai.Device(pipeline) as device:
    rgb_queue = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
    depth_queue = device.getOutputQueue(name="depth", maxSize=4, blocking=False)

    while True:
        # Captura RGB
        rgb_frame = rgb_queue.get().getCvFrame()

        # Captura Profundidade
        depth_frame = depth_queue.get().getFrame()
        depth_frame = cv2.normalize(depth_frame, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        depth_frame = cv2.applyColorMap(depth_frame, cv2.COLORMAP_JET)

        # Exibir os frames
        cv2.imshow("OAK-D Lite - RGB", rgb_frame)
        cv2.imshow("OAK-D Lite - Profundidade", depth_frame)

        if cv2.waitKey(1) == ord('q'):  # Pressione 'q' para sair
            break

cv2.destroyAllWindows()