import numpy as np
import triton_python_backend_utils as pb_utils  # type: ignore


class TritonPythonModel:
    def initialize(self, args):
        _ = args

    def execute(self, requests):
        responses = []
        for request in requests:
            bbox_2d = pb_utils.get_input_tensor_by_name(request, "bbox_2d").as_numpy().astype(np.float32, copy=False)
            function_2d = pb_utils.get_input_tensor_by_name(
                request, "function_embedding_2d"
            ).as_numpy().astype(np.float32, copy=False)
            vision_2d = pb_utils.get_input_tensor_by_name(
                request, "vision_embedding_2d"
            ).as_numpy().astype(np.float32, copy=False)
            text_2d = pb_utils.get_input_tensor_by_name(
                request, "text_embedding_2d"
            ).as_numpy().astype(np.float32, copy=False)

            n = int(bbox_2d.shape[0]) if bbox_2d.ndim >= 1 else 0
            padding_mask = np.ones((1, n), dtype=np.int32)

            responses.append(
                pb_utils.InferenceResponse(
                    output_tensors=[
                        pb_utils.Tensor("bbox", np.expand_dims(bbox_2d, axis=0)),
                        pb_utils.Tensor("function_embedding", np.expand_dims(function_2d, axis=0)),
                        pb_utils.Tensor("vision_embedding", np.expand_dims(vision_2d, axis=0)),
                        pb_utils.Tensor("text_embedding", np.expand_dims(text_2d, axis=0)),
                        pb_utils.Tensor("padding_mask", padding_mask),
                    ]
                )
            )
        return responses

    def finalize(self):
        pass