"""
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
"""

import cv2
import numpy as np
import openvino
from openvino.runtime import Core


class OvInfer:
    def __init__(self, model_path, device_name="AUTO"):
        self.resized_image = None
        self.ratio = None
        self.resize_image = None
        self.ori_image = None
        self.device = device_name
        self.model_path = model_path
        self.core = Core()
        self.available_device = self.core.available_devices
        self.compile_model = self.core.compile_model(self.model_path, device_name)
        self.target_size = [
            self.compile_model.inputs[0].get_partial_shape()[2].get_length(),
            self.compile_model.inputs[0].get_partial_shape()[3].get_length(),
        ]
        self.query_num = self.compile_model.outputs[0].get_partial_shape()[1].get_length()

    def infer(self, inputs: dict):
        infer_request = self.compile_model.create_infer_request()
        for input_name, input_data in inputs.items():
            input_tensor = openvino.Tensor(input_data)
            infer_request.set_tensor(input_name, input_tensor)
        infer_request.infer()
        outputs = {
            "labels": infer_request.get_tensor("labels").data,
            "boxes": infer_request.get_tensor("boxes").data,
            "scores": infer_request.get_tensor("scores").data,
        }
        return outputs

    def process_image(self, ori_image, keep_ratio: bool):
        self.ori_image = ori_image
        h, w = ori_image.shape[:2]
        if keep_ratio:
            r = min(self.target_size[0] / h, self.target_size[1] / w)
            self.ratio = r
            new_w = int(w * r)
            new_h = int(h * r)
            temp_image = cv2.resize(ori_image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
            resized_image = np.full(
                (self.target_size[0], self.target_size[1], 3), 114, dtype=temp_image.dtype
            )
            resized_image[:new_h, :new_w, :] = temp_image
            self.resized_image = resized_image
        else:
            self.resized_image = cv2.resize(
                ori_image, self.target_size, interpolation=cv2.INTER_LINEAR
            )
        blob_image = cv2.dnn.blobFromImage(self.resized_image, 1.0 / 255.0)
        orig_size = np.array([self.resized_image.shape[0], self.resized_image.shape[1]], dtype=np.int64).reshape(
            1, 2
        )

        inputs = {
            "images": blob_image,
            "orig_target_sizes": orig_size,
        }
        return inputs

    def get_available_device(self):
        return self.available_device

    def draw_and_save_image(self, infer_result, image_path, score_threshold=0.6):
        draw_image = self.ori_image
        scores = infer_result["scores"]
        labels = infer_result["labels"]
        boxes = infer_result["boxes"]
        for i in range(self.query_num):
            if scores[0, i] > score_threshold:
                cx = boxes[0, i, 0] / self.ratio
                cy = boxes[0, i, 1] / self.ratio
                bx = boxes[0, i, 2] / self.ratio
                by = boxes[0, i, 3] / self.ratio
                cv2.rectangle(
                    draw_image, (int(cx), int(cy), int(bx - cx), int(by - cy)), (255, 0, 0), 1
                )
        cv2.imwrite(image_path, draw_image)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("-image", "--image", type=str, required=True)
    parser.add_argument("-ov_model", "--ov_model", type=str, required=True)
    args = parser.parse_args()
    img = cv2.imread(args.image)
    mOvInfer = OvInfer(args.ov_model)
    inputs = mOvInfer.process_image(img, True)
    outputs = mOvInfer.infer(inputs)
    mOvInfer.draw_and_save_image(outputs, "openvino_result.jpg")