import time
import sys

import numpy as np
import torch
import triton_python_backend_utils as pb_utils  # type: ignore
if "/" not in sys.path:
    sys.path.insert(0, "/")
from utils import non_max_suppression


class TritonPythonModel:
    def initialize(self, args):
        self.conf_thres = 0.31
        self.iou_thres = 0.55

    def execute(self, requests):
        logger = pb_utils.Logger

        responses = []
        st = time.time()
        for request in requests:
            output = pb_utils.get_input_tensor_by_name(request, "model_output").as_numpy()

            output_tensor = torch.tensor(output).reshape(1, 22, 18900)
            output_list = non_max_suppression(
                output_tensor,
                conf_thres=self.conf_thres,
                iou_thres=self.iou_thres,
            )
            detections = output_list[0]
            if len(detections) == 0:
                output_numpy = np.empty((1, 0, 4), dtype=np.float32)
            else:
                boxes = detections[:, :4].detach().cpu().numpy().astype(np.float32)
                output_numpy = np.expand_dims(boxes, axis=0)

            out_tensor_0 = pb_utils.Tensor("detections", output_numpy)

            responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0]))

        logger.log_info(f"od.widget.postprocess execute duration : {int((time.time() - st)*1000)} ms")
        return responses