import depthai as dai
import time
from datetime import timedelta
import math

FPS = 30

# Start defining a pipeline
pipeline = dai.Pipeline()

# Define a source - color camera
camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_4_K)
camRgb.setVideoNumFramesPool(10)
camRgb.setFps(FPS)

left = pipeline.create(dai.node.MonoCamera)
left.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
left.setBoardSocket(dai.CameraBoardSocket.LEFT)
left.setFps(FPS)

right = pipeline.create(dai.node.MonoCamera)
right.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
right.setBoardSocket(dai.CameraBoardSocket.RIGHT)
right.setFps(FPS)

stereo = pipeline.create(dai.node.StereoDepth)
stereo.initialConfig.setMedianFilter(dai.MedianFilter.KERNEL_7x7)
stereo.setLeftRightCheck(True)
stereo.setExtendedDisparity(False)
stereo.setSubpixel(False)
stereo.setDepthAlign(dai.CameraBoardSocket.RGB)
left.out.link(stereo.left)
right.out.link(stereo.right)

script = pipeline.create(dai.node.Script)
script.setScript(R"""
FPS=30
import time
from datetime import timedelta
import math

MS_THRESHOL = math.ceil(500 / FPS)

def check_sync(queues, timestamp):
    matching_frames = []
    for name, list in queues.items():
        for i, msg in enumerate(list):
            time_diff = abs(msg.getTimestamp() - timestamp)
            if time_diff <= timedelta(milliseconds=MS_THRESHOL):
                matching_frames.append(i)
                break

    if len(matching_frames) == len(queues):
        i = 0
        for name, list in queues.items():
            queues[name] = queues[name][matching_frames[i]:]
            i += 1
        return True
    else:
        return False

names = ['disp', 'rgb']
frames = dict()
for name in names:
    frames[name] = []

while True:
    for name in names:
        f = node.io[name + "_in"].tryGet()
        if f is not None:
            frames[name].append(f)
            node.warn(f"Added {name} to queue, seq {str(f.getSequenceNum())}, list size {str(len(frames[name]))}")

            if check_sync(frames, f.getTimestamp()):
                node.info("Synced frame!")
                for name, list in frames.items():
                    syncedF = list.pop(0)
                    node.info(f"{name}, ts: {str(syncedF.getTimestamp())}, seq {str(syncedF.getSequenceNum())}")
                    node.io[name + '_out'].send(syncedF)

    time.sleep(0.001)
""")
stereo.disparity.link(script.inputs["disp_in"])
camRgb.video.link(script.inputs["rgb_in"])

scriptOut = ["disp", "rgb"]
for name in scriptOut:
    xout = pipeline.create(dai.node.XLinkOut)
    xout.setStreamName(name)
    script.outputs[name + "_out"].link(xout.input)

# Connect to device and start pipeline
with dai.Device(pipeline) as device:
    device.setLogOutputLevel(dai.LogLevel.INFO)
    device.setLogLevel(dai.LogLevel.INFO)

    streams = ["rgb", "disp"]
    queue_map = {}
    for name in streams:
        queue_map[name] = device.getOutputQueue(name)

    while True:
        for name in streams:
            img = queue_map[name].get()
            print(f"Stream {name}, timestamp: {img.getTimestamp().total_seconds()}, sequence number: {img.getSequenceNum()}")