import sys
sys.path.append("../")

import time
from openvino.inference_engine import IECore
import numpy as np
import cv2
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator
import torch
import onnxruntime as ort
from timeit import default_timer as timer
from models.common import DetectMultiBackend
import pyvirtualcam

WEBCAM_PORT = 0
IMG_SIZE = 640
CONF_THRESH = 0.4
MODEL_XML = r"C:\Users\user\PycharmProjects\yolov5_v6.0\quantization\optimized\yolov5s_640x640_opt.xml"
MODEL_WEIGHTS = r"C:\Users\user\PycharmProjects\yolov5_v6.0\quantization\optimized\yolov5s_640x640_opt.bin"


# yolov5に入力するIMG_SIZExIMG_SIZEで背景が0パディングされている形式に変更
def convert_to_yolov5format_img(frame):
    height, width = frame.shape[:2]
    if width < height:
        size, limit = height, width
    else:
        limit, size = height, width
    start = int((size - limit) / 2)
    fin = int((size + limit) / 2)

    transformed = np.full((size, size, 3), 114, np.uint8)
    if size == height:
        transformed[:, start:fin] = frame
    else:
        transformed[start:fin, :] = frame
    transformed = cv2.resize(transformed, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_CUBIC)
    return transformed

class DispFps:
    def __init__(self):
        # 表示関連定義
        self.__width = 160
        self.__height = 40
        self.__font_size = 1.0
        self.__font_width = 1
        self.__font_style = cv2.FONT_HERSHEY_COMPLEX
        self.__font_color = (255, 255, 255)
        self.__background_color = (0, 0, 0)

        # フレーム数カウント用変数
        self.__frame_count = 0

        # FPS計算用変数
        self.__accum_time = 0
        self.__curr_fps = 0
        self.__prev_time = timer()
        self.__str = "FPS: "

    def __calc(self):
        # フレーム数更新
        self.__frame_count += 1

        # FPS更新
        self.__curr_time = timer()
        self.__exec_time = self.__curr_time - self.__prev_time
        self.__prev_time = self.__curr_time
        self.__accum_time = self.__accum_time + self.__exec_time
        self.__curr_fps = self.__curr_fps + 1
        if self.__accum_time > 1:
            self.__accum_time = self.__accum_time - 1
            self.__str = "FPS: " + str(self.__curr_fps)
            self.__curr_fps = 0

    def __disp(self, frame, str, x1, y1, x2, y2):
        cv2.rectangle(frame, (x1, y1), (x2, y2), self.__background_color, -1)
        cv2.putText(frame, str, (x1 + 5, y2 - 5), self.__font_style, self.__font_size, self.__font_color,
                    self.__font_width)

    def disp(self, frame):
        # 表示内容計算
        self.__calc()
        # フレーム数（左上に表示する）
        self.__disp(frame, str(self.__frame_count), 0, 0, x2=self.__width, y2=self.__height)
        # FPS(右上に表示する)
        screen_width = int(frame.shape[1])
        self.__disp(frame, self.__str, screen_width - self.__width, 0, screen_width, self.__height)


ie = IECore()
net = ie.read_network(model=MODEL_XML, weights=MODEL_WEIGHTS)

input_layer = next(iter(net.input_info))
print(f"input layout: {net.input_info[input_layer].layout}")
print(f"input precision: {net.input_info[input_layer].precision}")
print(f"input shape: {net.input_info[input_layer].tensor_desc.dims}")

# config = {"EXCLUSIVE_ASYNC_REQUESTS": "YES",
#           "CPU_THREADS_NUM": "8",
#           "CPU_THROUGHPUT_STREAMS": "CPU_THROUGHPUT_AUTO"}
config = {}

exec_net = ie.load_network(network=net, device_name="CPU", config=config)

# torch_net = DetectMultiBackend('C:\\Users\\user\\PycharmProjects\\yolov5_v6.0\\runs\\train\\exp10\\weights\\best.pt', "cuda:0", False)
# torch_net.warmup(imgsz=(1, 3, 640, 640), half=False)

# onnx_net = ort.InferenceSession(r"C:\Users\user\PycharmProjects\yolov5_v6.0\runs\train\exp10\weights\yolov5s_640x640_opt.onnx")

input_layer = next(iter(net.input_info))

cap = cv2.VideoCapture(0)
dispFps = DispFps()

ret, frame = cap.read()
with pyvirtualcam.Camera(width=640, height=640, fps=20.0) as cam:
    while True:
        t1 = time.time()
        ret, frame = cap.read()

        if not ret:
            break
        N, C, H, W = net.input_info[input_layer].tensor_desc.dims
        # OpenCV resize expects the destination size as (width, height)
        src = convert_to_yolov5format_img(frame)
        input_data = np.expand_dims(np.transpose(src, (2, 0, 1)), 0).astype(np.float32)
        input_data /= 255
        t2 = time.time()

        result = torch.tensor(exec_net.infer({input_layer: input_data})["output"])
        # result = torch_net(torch.tensor(input_data).to(torch.device("cuda:0")), augment=False, visualize=False)
        # result = torch.tensor(onnx_net.run(None, {"images": input_data})[0])
        t3 = time.time()
        # NMS
        pred = non_max_suppression(result, CONF_THRESH, 0.45, 0, False, max_det=10)
        for i, det in enumerate(pred):
            gn = torch.tensor(src.shape)[[1, 0, 1, 0]]
            annotator = Annotator(src, line_width=3, example=str(["white_line"]))
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(input_data.shape[2:], det[:, :4], src.shape).round()

                for *xyxy, conf, cls in reversed(det):
                    c = int(cls)
                    label = "white_line"
                    annotator.box_label(xyxy, label)

        t4 = time.time()
        # dispFps.disp(src)

        t5 = time.time()
        # cv2.imshow("detect", src)
        cam.send(cv2.cvtColor(src, cv2.COLOR_BGR2RGB))

        print(f"{(t2 - t1):0.3f}, {(t3 - t2):0.3f}, {(t4 - t3):0.3f}, {(t5 - t4):0.3f}")

        cam.sleep_until_next_frame()

        key = cv2.waitKey(1)
        if key == ord("q"):
            break