import cv2
import depthai as dai
import numpy as np


FRAME_SIZE = (640, 400)
FPS = 30

MIN_DEPTH_MM = 300
MAX_DEPTH_MM = 5000


def colorize_depth(depth):
    valid = depth > 0

    clipped = np.clip(
        depth,
        MIN_DEPTH_MM,
        MAX_DEPTH_MM
    )

    depth8 = np.zeros(depth.shape, dtype=np.uint8)

    depth8[valid] = (
        (MAX_DEPTH_MM - clipped[valid])
        * 255.0
        / (MAX_DEPTH_MM - MIN_DEPTH_MM)
    ).astype(np.uint8)

    color = cv2.applyColorMap(
        depth8,
        cv2.COLORMAP_TURBO
    )

    color[~valid] = 0

    return color


with dai.Pipeline() as pipeline:

    # --------------------------------------------------
    # RGB camera
    # OAK-D: CAM_A
    # --------------------------------------------------
    color = pipeline.create(dai.node.Camera).build(
        dai.CameraBoardSocket.CAM_A,
        sensorFps=FPS
    )

    # --------------------------------------------------
    # Stereo cameras
    # OAK-D: CAM_B = left
    #        CAM_C = right
    # --------------------------------------------------
    left = pipeline.create(dai.node.Camera).build(
        dai.CameraBoardSocket.CAM_B,
        sensorFps=FPS
    )

    right = pipeline.create(dai.node.Camera).build(
        dai.CameraBoardSocket.CAM_C,
        sensorFps=FPS
    )

    # --------------------------------------------------
    # Stereo depth
    # --------------------------------------------------
    stereo = pipeline.create(dai.node.StereoDepth)

    stereo.setDefaultProfilePreset(
        dai.node.StereoDepth.PresetMode.DEFAULT
    )

    stereo.setRectifyEdgeFillColor(0)
    stereo.enableDistortionCorrection(True)

    # Left / Right camera -> StereoDepth
    left.requestOutput(FRAME_SIZE).link(
        stereo.left
    )

    right.requestOutput(FRAME_SIZE).link(
        stereo.right
    )

    # --------------------------------------------------
    # RGBD
    #
    # Important:
    # autocreate=True is NOT used.
    # --------------------------------------------------
    rgbd = pipeline.create(dai.node.RGBD).build(
        color,
        stereo,
        FRAME_SIZE,
        FPS
    )

    rgbd_queue = rgbd.rgbd.createOutputQueue(
        maxSize=4,
        blocking=False
    )

    pipeline.start()

    print("Pipeline started")
    print("Q / ESC : quit")

    while pipeline.isRunning():

        data = rgbd_queue.get()

        rgb_msg = data.getRGBFrame()
        depth_msg = data.getDepthFrame()

        if rgb_msg is None or depth_msg is None:
            continue

        rgb = rgb_msg.getCvFrame()

        # uint16, mm
        depth = depth_msg.getCvFrame()

        depth_color = colorize_depth(depth)

        # --------------------------------------------------
        # Center distance
        # --------------------------------------------------
        h, w = depth.shape

        cx = w // 2
        cy = h // 2

        d = int(depth[cy, cx])

        if d > 0:
            text = f"{d} mm"
        else:
            text = "invalid"

        cv2.drawMarker(
            rgb,
            (cx, cy),
            (0, 0, 255),
            cv2.MARKER_CROSS,
            20,
            2
        )

        cv2.putText(
            rgb,
            text,
            (20, 40),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 255, 0),
            2
        )

        cv2.imshow("RGB", rgb)
        cv2.imshow("Depth", depth_color)

        key = cv2.waitKey(1) & 0xff

        if key == ord("q") or key == 27:
            break


cv2.destroyAllWindows()