pipeline = dai.Pipeline()
camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setBoardSocket(dai.CameraBoardSocket.RGB)
camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
camRgb.setIspScale(2, 3)
camRgb.setFps(30)

monoleft = pipeline.create(dai.node.MonoCamera)
monoright = pipeline.create(dai.node.MonoCamera)
stereo = pipeline.create(dai.node.StereoDepth)
spatialLocationCalculator = pipeline.create(dai.node.SpatialLocationCalculator)

xoutDepth = pipeline.create(dai.node.XLinkOut) # 深度图输出
xoutRgb = pipeline.create(dai.node.XLinkOut) # 彩色图输出
spatial_config_in = pipeline.create(dai.node.XLinkIn) # 空间位置设置输入
depth_config_in = pipeline.create(dai.node.XLinkIn)
xoutSpatialData = pipeline.create(dai.node.XLinkOut) # 空间数据输出

xoutDepth.setStreamName("depth")
xoutRgb.setStreamName("color")
xoutSpatialData.setStreamName("spatialData")
spatial_config_in.setStreamName("spatialCalConfig")
depth_config_in.setStreamName("DepthConfig")

monoleft.setResolution(dai.MonoCameraProperties.SensorResolution.THE_720_P)
monoleft.setBoardSocket(dai.CameraBoardSocket.LEFT)
monoright.setResolution(dai.MonoCameraProperties.SensorResolution.THE_720_P)
monoright.setBoardSocket(dai.CameraBoardSocket.RIGHT)
stereo.setDepthAlign(dai.CameraBoardSocket.RGB)

stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_ACCURACY)
stereo.initialConfig.setLeftRightCheck(lr_check)
stereo.initialConfig.setExtendedDisparity(extended_disparity)
stereo.initialConfig.setSubpixel(subpixel)
stereo.initialConfig.setMedianFilter(dai.MedianFilter.KERNEL_3x3)

config = stereo.initialConfig.get()
config.postProcessing.speckleFilter.enable = False
config.postProcessing.speckleFilter.speckleRange = 50
config.postProcessing.temporalFilter.enable = True
config.postProcessing.spatialFilter.enable = True
config.postProcessing.spatialFilter.holeFillingRadius = 2
config.postProcessing.spatialFilter.numIterations = 1
config.postProcessing.thresholdFilter.minRange = 400
config.postProcessing.thresholdFilter.maxRange = 15000
config.postProcessing.decimationFilter.decimationFactor = 1

stereo.initialConfig.set(config)

config_spatial_location = dai.SpatialLocationCalculatorConfigData()
config_spatial_location.depthThresholds.lowerThreshold = 10
config_spatial_location.depthThresholds.upperThreshold = 10000
config_spatial_location.roi = dai.Rect(TOP_LEFT, BOTTOM_RIGHT)
config_spatial_location.calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MIN # 默认取平均

config_spatial_location1 = dai.SpatialLocationCalculatorConfigData()
config_spatial_location1.depthThresholds.lowerThreshold = 10
config_spatial_location1.depthThresholds.upperThreshold = 10000
config_spatial_location1.roi = dai.Rect(TOP_LEFT1, BOTTOM_RIGHT1)
config_spatial_location1.calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MIN # 默认取平均

config_spatial_location2 = dai.SpatialLocationCalculatorConfigData()
config_spatial_location2.depthThresholds.lowerThreshold = 10
config_spatial_location2.depthThresholds.upperThreshold = 10000
config_spatial_location2.roi = dai.Rect(TOP_LEFT2, BOTTOM_RIGHT2)
config_spatial_location2.calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MIN # 默认取平均

config_spatial_location3 = dai.SpatialLocationCalculatorConfigData()
config_spatial_location3.depthThresholds.lowerThreshold = 10
config_spatial_location3.depthThresholds.upperThreshold = 10000
config_spatial_location3.roi = dai.Rect(TOP_LEFT3, BOTTOM_RIGHT3)
config_spatial_location3.calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MIN # 默认取平均

config_spatial_location4 = dai.SpatialLocationCalculatorConfigData()
config_spatial_location4.depthThresholds.lowerThreshold = 10
config_spatial_location4.depthThresholds.upperThreshold = 10000
config_spatial_location4.roi = dai.Rect(TOP_LEFT4, BOTTOM_RIGHT4)
config_spatial_location4.calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MIN # 默认取平均

spatialLocationCalculator.inputConfig.setWaitForMessage(False)
spatialLocationCalculator.initialConfig.setROIs(
[config_spatial_location, config_spatial_location1, config_spatial_location2, config_spatial_location3,
config_spatial_location4])

# 链接
camRgb.isp.link(xoutRgb.input)
monoleft.out.link(stereo.left)
monoright.out.link(stereo.right)
spatialLocationCalculator.passthroughDepth.link(xoutDepth.input)
stereo.depth.link(spatialLocationCalculator.inputDepth)
spatial_config_in.out.link(spatialLocationCalculator.inputConfig)
spatialLocationCalculator.out.link(xoutSpatialData.input)
depth_config_in.out.link(stereo.inputConfig)

with dai.Device(pipeline) as device:
device.setIrLaserDotProjectorBrightness(1200)
device.setIrFloodLightBrightness(750)

depth_queue = device.getOutputQueue(name="depth", maxSize=4, blocking=False)
color_queue = device.getOutputQueue(name="color", maxSize=4, blocking=False)
spatial_data_queue = device.getOutputQueue(name="spatialData", maxSize=4, blocking=False)
spatial_config = device.getInputQueue("spatialCalConfig")
depth_config = device.getInputQueue("DepthConfig")
calibration = device.readCalibration()
w, h = camRgb.getIspSize()
print("w, h: ", (w, h))
# 获取内参矩阵和畸变系数
intrinsics = np.array(calibration.getCameraIntrinsics(dai.CameraBoardSocket.RGB, dai.Size2f(w, h)))
coefficients = np.array(calibration.getDistortionCoefficients(dai.CameraBoardSocket.RGB))
print("please put the aruco marker on the robot arm end")
while True:

    color = color_queue.get().getCvFrame()
    print("color.shape:", color.shape)
    depth = depth_queue.get().getFrame()
    print("depth.shape", depth.shape)
    spatialData = spatial_data_queue.get().getSpatialLocations()
    # depthFrame = cv2.normalize(depth, None, 255, 0, cv2.NORM_INF, cv2.CV_8UC1)
    # depthFrame = cv2.equalizeHist(depthFrame)
    # depthFrame = cv2.applyColorMap(depthFrame, cv2.COLORMAP_JET)
    depth_downsacled = depth[::4]
    min_depth = np.percentile(depth_downsacled[depth_downsacled != 0], 1)
    max_depth = np.percentile(depth_downsacled, 99)
    depthFrame = np.interp(depth, (min_depth, max_depth), (0, 255)).astype(np.uint8)
    depthFrame = cv2.applyColorMap(depthFrame, cv2.COLORMAP_JET)

    for depthData in spatialData:
        roi = depthData.config.roi
        roi = roi.denormalize(width=color.shape[1], height=color.shape[0])
        xmin = int(roi.topLeft().x)
        ymin = int(roi.topLeft().y)
        xmax = int(roi.bottomRight().x)
        ymax = int(roi.bottomRight().y)
        cv2.rectangle(depthFrame, (xmin, ymin), (xmax, ymax), (255, 255, 255), 1)
        cv2.rectangle(color, (xmin + 200, ymin), (xmax + 200, ymax), (255, 255, 255), 1)
        cv2.putText(color, f"X: {int(depthData.spatialCoordinates.x)}",
                    (xmin - 300, ymin + 20),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (51, 87, 255))
        cv2.putText(color, f"Y: {int(depthData.spatialCoordinates.y)}",
                    (xmin - 300, ymin + 45),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (51, 87, 255))
        cv2.putText(color, f"Z: {int(depthData.spatialCoordinates.z)}",
                    (xmin - 300, ymin + 70),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (51, 87, 255))

        cv2.putText(depthFrame, f"X: {int(depthData.spatialCoordinates.x)} mm",
                    (40, 40),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255))
        cv2.putText(depthFrame, f"Y: {int(depthData.spatialCoordinates.y)} mm",
                    (40, 75),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255))
        cv2.putText(depthFrame, f"Z: {int(depthData.spatialCoordinates.z)} mm",
                    (40, 110),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255))
        cv2.putText(depthFrame, f"LeftRightCheck: {lr_check}",
                    (40, 145),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255))
        cv2.putText(depthFrame, f"Extended: {extended_disparity}",
                    (40, 180),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255))
        cv2.putText(depthFrame, f"Subpixel: {subpixel}",
                    (40, 215),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255))
        cv2.putText(depthFrame, f"MedianFilter: {median_filter}",
                    (40, 250),
                    cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255))
        depthMin = depthData.depthMin
        depthMax = depthData.depthMax
        X = depthData.spatialCoordinates.x
        Y = depthData.spatialCoordinates.y
        Z = depthData.spatialCoordinates.z
        center = [X, Y, Z]
    corners, ids, _ = cv2.aruco.detectMarkers(color, dictionary,
                                              parameters=parameters)
    # 旋转矩阵和平移向量
    rotate_vec, trans_vec, _ = cv2.aruco.estimatePoseSingleMarkers(corners, MARKER_LENGTH,
                                                                   intrinsics, coefficients)
    if ids is not None:
        cv2.aruco.drawDetectedMarkers(color, corners)
        cv2.drawFrameAxes(color, intrinsics, coefficients, rotate_vec, trans_vec, 0.05)
        for i, corner in zip(ids, corners):
            # *todo: 机械臂没有python的api,所以现在用的是手动调整roi窗口至aruco marker中点获取x,y,z坐标*
            x = (corner[0][0][0] + corner[0][2][0]) / 2
            y = (corner[0][0][1] + corner[0][2][1]) / 2

            cv2.circle(color, (int(x), int(y)), 3, (0, 0, 0), -1)  # 像素坐标
            TOP_LEFT.x = (int(x) - 1) / color.shape[1]
            TOP_LEFT.y = (int(y) - 1) / color.shape[0]
            BOTTOM_RIGHT.x = (int(x) + 1) / color.shape[1]
            BOTTOM_RIGHT.y = (int(y) + 1) / color.shape[0]

            TOP_LEFT1.x = (int(corner[0][0][0]) - 1) / color.shape[1]
            TOP_LEFT1.y = (int(corner[0][0][1]) - 1) / color.shape[0]
            BOTTOM_RIGHT1.x = (int(corner[0][0][0]) + 1) / color.shape[1]
            BOTTOM_RIGHT1.y = (int(corner[0][0][1]) + 1) / color.shape[0]

            TOP_LEFT2.x = (int(corner[0][1][0]) - 1) / color.shape[1]
            TOP_LEFT2.y = (int(corner[0][1][1]) - 1) / color.shape[0]
            BOTTOM_RIGHT2.x = (int(corner[0][1][0]) + 1) / color.shape[1]
            BOTTOM_RIGHT2.y = (int(corner[0][1][1]) + 1) / color.shape[0]

            TOP_LEFT3.x = (int(corner[0][2][0]) - 1) / color.shape[1]
            TOP_LEFT3.y = (int(corner[0][2][1]) - 1) / color.shape[0]
            BOTTOM_RIGHT3.x = (int(corner[0][2][0]) + 1) / color.shape[1]
            BOTTOM_RIGHT3.y = (int(corner[0][2][1]) + 1) / color.shape[0]

            TOP_LEFT4.x = (int(corner[0][3][0]) - 1) / color.shape[1]
            TOP_LEFT4.y = (int(corner[0][3][1]) - 1) / color.shape[0]
            BOTTOM_RIGHT4.x = (int(corner[0][3][0]) + 1) / color.shape[1]
            BOTTOM_RIGHT4.y = (int(corner[0][3][1]) + 1) / color.shape[0]
            saptialCal_new_config = True
    cv2.imshow("depth", depthFrame)
    cv2.imshow("color", color)
    key = cv2.waitKey(1)

    if key == ord("q"):
        break
    elif key in [ord("0"), ord("3"), ord("5"), ord("7")]:
        median_filter = MedianFilter[chr(key)]
        depth_new_config = True
    elif key == ord("s"):
        subpixel = not subpixel
        if subpixel:
            extended_disparity = False
        depth_new_config = True
    elif key == ord("e"):
        extended_disparity = not extended_disparity
        if extended_disparity:
            subpixel = False
        depth_new_config = True
    elif key == ord("l"):
        lr_check = not lr_check
        depth_new_config = True

    if saptialCal_new_config:
        config = dai.SpatialLocationCalculatorConfigData()
        config.depthThresholds.lowerThreshold = 100
        config.depthThresholds.upperThreshold = 12000
        config.roi = dai.Rect(TOP_LEFT, BOTTOM_RIGHT)
        config.calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MIN

        config_1 = dai.config = dai.SpatialLocationCalculatorConfigData()
        config_1.depthThresholds.lowerThreshold = 100
        config_1.depthThresholds.upperThreshold = 12000
        config_1.roi = dai.Rect(TOP_LEFT1, BOTTOM_RIGHT1)

        config_2 = dai.config = dai.SpatialLocationCalculatorConfigData()
        config_2.depthThresholds.lowerThreshold = 100
        config_2.depthThresholds.upperThreshold = 12000
        config_2.roi = dai.Rect(TOP_LEFT2, BOTTOM_RIGHT2)

        config_3 = dai.config = dai.SpatialLocationCalculatorConfigData()
        config_3.depthThresholds.lowerThreshold = 100
        config_3.depthThresholds.upperThreshold = 12000
        config_3.roi = dai.Rect(TOP_LEFT3, BOTTOM_RIGHT3)

        config_4 = dai.config = dai.SpatialLocationCalculatorConfigData()
        config_4.depthThresholds.lowerThreshold = 100
        config_4.depthThresholds.upperThreshold = 12000
        config_4.roi = dai.Rect(TOP_LEFT4, BOTTOM_RIGHT4)

        cfg = dai.SpatialLocationCalculatorConfig()
        cfg.setROIs([config, config_1, config_2, config_3, config_4])
        spatial_config.send(cfg)
        saptialCal_new_config = False
    if depth_new_config:
        config = dai.StereoDepthConfig()
        config.setLeftRightCheck(lr_check)
        config.setExtendedDisparity(extended_disparity)
        config.setSubpixel(subpixel)
        config.setMedianFilter(median_filter)
        depth_config.send(config)
        depth_new_config = False