import depthai as dai
import json
from pathlib import Path
import blobconverter
import cv2
import numpy as np

pipeline = dai.Pipeline()
# Define source and output
camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_12_MP)
camRgb.setInterleaved(False)
camRgb.setIspScale(1,5) # 4056x3040 -> 812x608
camRgb.setPreviewSize(812, 608)
camRgb.setIspNumFramesPool(2)
camRgb.setBoardSocket(dai.CameraBoardSocket.RGB)
# Slightly lower FPS to avoid lag, as ISP takes more resources at 12MP
camRgb.setFps(10)


manip = pipeline.create(dai.node.ImageManip)
manip.setMaxOutputFrameSize(812 * 608 * 3) # 300x300x3
rgbRr = dai.RotatedRect()
rgbRr.center.x, rgbRr.center.y = camRgb.getPreviewWidth() // 2, camRgb.getPreviewHeight() // 2
rgbRr.size.width, rgbRr.size.height = camRgb.getPreviewHeight(), camRgb.getPreviewWidth()
rgbRr.angle = 90
manip.initialConfig.setCropRotatedRect(rgbRr, False)    #manip.initialConfig.setCropRect(0, 0, 1, 1)
camRgb.preview.link(manip.inputImage)

resizeManip = pipeline.create(dai.node.ImageManip)
resizeManip.initialConfig.setResizeThumbnail(300, 300)
resizeManip.initialConfig.setFrameType(dai.RawImgFrame.Type.BGR888p)
manip.out.link(resizeManip.inputImage)

xoutIsp = pipeline.create(dai.node.XLinkOut)
xoutIsp.setStreamName("isp")
manip.out.link(xoutIsp.input)

#manip.initialConfig.setRotation(-90)
#manip.initialConfig.setRotationDegrees(90)


# NN to demonstrate how to run inference on full FOV frames
nn = pipeline.create(dai.node.MobileNetDetectionNetwork)
nn.setConfidenceThreshold(0.5)
nn.setBlobPath(blobconverter.from_zoo(name='mobilenet-ssd'))
resizeManip.out.link(nn.input)

xoutNn = pipeline.create(dai.node.XLinkOut)
xoutNn.setStreamName("nn")
nn.out.link(xoutNn.input)

xoutRgb = pipeline.create(dai.node.XLinkOut)
xoutRgb.setStreamName("pass")
nn.passthrough.link(xoutRgb.input)

with dai.Device(pipeline) as device:
    qIsp = device.getOutputQueue(name="isp", maxSize=4, blocking=False)
    qDet = device.getOutputQueue(name="nn", maxSize=4, blocking=False)
    qPass = device.getOutputQueue(name="pass", maxSize=4, blocking=False)

    detections = []
    frame = None

    def frameNorm(frame, bbox):
        normVals = np.full(len(bbox), frame.shape[0])
        normVals[::2] = frame.shape[1]
        return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int)

    def displayFrame(name, frame):
        color = (255, 0, 0)
        for detection in detections:
            bbox = frameNorm(frame, (detection.xmin, detection.ymin, detection.xmax, detection.ymax))
            cv2.putText(frame, str(detection.label), (bbox[0] + 10, bbox[1] + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
            cv2.putText(frame, f"{int(detection.confidence * 100)}%", (bbox[0] + 10, bbox[1] + 40), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
            cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
        # Show the frame
        cv2.imshow(name, frame)

    while True:
        if qDet.has():
            detections = qDet.get().detections
        if qPass.has():
            frame = qPass.get().getCvFrame()

        if frame is not None:
            displayFrame("pass", frame)

        if qIsp.has():
            isp_frame = qIsp.get().getCvFrame()
            print(isp_frame.shape)
            cv2.imshow("isp", isp_frame)
        if cv2.waitKey(1) == ord('q'):
            break