#!/usr/bin/env python3

import cv2
import depthai as dai
import numpy as np

# Closer-in minimum depth, disparity range is doubled (from 95 to 190):
extendedDisparity = True
# Better accuracy for longer distance, fractional disparity 32-levels:
subpixel = False  # True
# Better handling for occlusions:
lrCheck = True

enableRectified = True

iscScale = (2, 3)  # (1, 3) / (2, 3)

# Create pipeline
pipeline = dai.Pipeline()

# Define sources and outputs
left_a = pipeline.create(dai.node.ColorCamera)
right_d = pipeline.create(dai.node.ColorCamera)

# Create stereo
stereo_ad = pipeline.create(dai.node.StereoDepth)
xoutDepth_ad = pipeline.create(dai.node.XLinkOut)
xoutDepth_ad.setStreamName("disparity_ad")

# Properties
left_a.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1200_P)
left_a.setCamera("CAMA-4L")
left_a.setIspScale(*iscScale)

right_d.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1200_P)
right_d.setCamera("CAMD-4L")
right_d.setIspScale(*iscScale)

stereo_ad.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
stereo_ad.initialConfig.setMedianFilter(dai.MedianFilter.KERNEL_7x7)
stereo_ad.setLeftRightCheck(lrCheck)
stereo_ad.setExtendedDisparity(extendedDisparity)
stereo_ad.setSubpixel(subpixel)

# Linking
left_a.isp.link(stereo_ad.left)
right_d.isp.link(stereo_ad.right)
if enableRectified:
    xoutRectR_ad = pipeline.create(dai.node.XLinkOut)
    xoutRectL_ad= pipeline.create(dai.node.XLinkOut)
    xoutRectR_ad.setStreamName("rectifiedRight_ad")
    xoutRectL_ad.setStreamName("rectifiedLeft_ad")
    stereo_ad.rectifiedLeft.link(xoutRectL_ad.input)
    stereo_ad.rectifiedRight.link(xoutRectR_ad.input)
stereo_ad.disparity.link(xoutDepth_ad.input)

# Define sources and outputs
left_b = pipeline.create(dai.node.ColorCamera)
right_c = pipeline.create(dai.node.ColorCamera)
# Create stereo
stereo_bc = pipeline.create(dai.node.StereoDepth)
xoutDepth_bc = pipeline.create(dai.node.XLinkOut)
xoutDepth_bc.setStreamName("disparity_bc")

# Properties
left_b.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1200_P)
left_b.setCamera("CAMB-2L")
left_b.setIspScale(*iscScale)

right_c.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1200_P)
right_c.setCamera("CAMC-2L")
right_c.setIspScale(*iscScale)

stereo_bc.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
stereo_bc.initialConfig.setMedianFilter(dai.MedianFilter.KERNEL_7x7)
stereo_bc.setLeftRightCheck(lrCheck)
stereo_bc.setExtendedDisparity(extendedDisparity)
stereo_bc.setSubpixel(subpixel)

# Linking
left_b.isp.link(stereo_bc.left)
right_c.isp.link(stereo_bc.right)
if enableRectified:
    xoutRectR_bc = pipeline.create(dai.node.XLinkOut)
    xoutRectL_bc = pipeline.create(dai.node.XLinkOut)
    xoutRectR_bc.setStreamName("rectifiedRight_bc")
    xoutRectL_bc.setStreamName("rectifiedLeft_bc")
    stereo_bc.rectifiedLeft.link(xoutRectL_bc.input)
    stereo_bc.rectifiedRight.link(xoutRectR_bc.input)
stereo_bc.disparity.link(xoutDepth_bc.input)

maxDisp = stereo_ad.initialConfig.getMaxDisparity()

# Connect to device and start pipeline
with dai.Device(pipeline) as device:
    while not device.isClosed():
        queueNames = device.getQueueEvents()
        for q in queueNames:
            message = device.getOutputQueue(q).get()
            # Display arrived frames
            if type(message) == dai.ImgFrame:
                frame = message.getCvFrame()
                if 'disparity' in q:
                    disp = (frame * (255.0 / maxDisp)).astype(np.uint8)
                    disp = cv2.applyColorMap(disp, cv2.COLORMAP_JET)
                    cv2.imshow(q, disp)
                else:
                    cv2.imshow(q, frame)
        if cv2.waitKey(1) == ord('q'):
            break