#!/usr/bin/env python3

from pathlib import Path
import sys
import cv2
import depthai as dai
import numpy as np
from datetime import datetime
import time
import json
import random
from pathlib import Path
import blobconverter
# Get argument first
nnPath = blobconverter.from_zoo(name="mobilenet-ssd")

use_nn = True
preview_size = (300, 300)
sensor_resolution = dai.ColorCameraProperties.SensorResolution.THE_4_K
codec = dai.VideoEncoderProperties.Profile.H265_MAIN
file_extension = 'h265'
fps = 30

pipeline = dai.Pipeline()

camRgb = pipeline.create(dai.node.ColorCamera)
videoEncoder = pipeline.create(dai.node.VideoEncoder)
nn = None

videoOut = pipeline.create(dai.node.XLinkOut)
videoOut.setStreamName("h265")

# Properties
camRgb.setBoardSocket(dai.CameraBoardSocket.RGB)
camRgb.setResolution(sensor_resolution)
camRgb.setPreviewSize(preview_size)
camRgb.setInterleaved(False)

videoEncoder.setDefaultProfilePreset(fps, codec)
# Linking
camRgb.video.link(videoEncoder.input)
videoEncoder.bitstream.link(videoOut.input)

if use_nn:
  manip = pipeline.createImageManip()
  manip.setResize(300,300)
  manip.inputImage.setBlocking(False)
  manip.inputImage.setQueueSize(2)
  camRgb.preview.link(manip.inputImage)

  nn = pipeline.create(dai.node.MobileNetDetectionNetwork)
  nn.setConfidenceThreshold(0.5)
  nn.setBlobPath(nnPath)
  nn.setNumInferenceThreads(2)
  # nn.input.setBlocking(False)
  # nn.input.setBlocking(False)
  manip.out.link(nn.input)

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


# Connect to device and start pipeline
print(datetime.now().strftime('%H:%M.%S.%f: Starting device'))
with dai.Device(pipeline) as device:

    print(datetime.now().strftime('%H:%M.%S.%f: Device started'))


    # Queues
    queue_size = 8
    qDet = None
    if use_nn:
      qDet = device.getOutputQueue("nn", queue_size)
    qRgbEnc = device.getOutputQueue('h265', maxSize=30, blocking=True)

    frameCount = 0
    frameStart = None

    while True:
        inDet = None
        if qDet:
            inDet = qDet.tryGet()

        while qRgbEnc.has():
            encFrame = qRgbEnc.get()
            # Ordinarily we might write this to a file
            frameCount += 1

        if not frameStart:
            frameStart = datetime.now()
            frameCount = 0
        if (datetime.now() - frameStart).total_seconds() > 10:
            print("Saw %s frames in 10 seconds" % frameCount)
            frameStart = datetime.now()
            frameCount = 0

        time.sleep(.01)