#  Tencent is pleased to support the open source community by making GNES available.
#
#  Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from typing import List

import numpy as np

from ..base import BaseImageEncoder
from ...helper import batching


class BaseONNXImageEncoder(BaseImageEncoder):
    batch_size = 64

    def __init__(self, model_name: str,
                 model_dir: str,
                 use_cuda: bool = False,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.model_dir = model_dir
        self.model_name = model_name
        self._use_cuda = use_cuda

    def post_init(self):
        import onnxruntime as ort

        self.sess = ort.InferenceSession(self.model_dir + '/' + self.model_name)
        inputs_info = self.sess.get_inputs()

        if len(inputs_info) != 1:
            raise ValueError('Now only support encoder with one input')
        else:
            self.input_name = inputs_info[0].name
            self.input_shape = inputs_info[0].shape
            self.input_type = inputs_info[0].type
            self.batch_size = self.input_shape[0]

    @batching
    def encode(self, img: List['np.ndarray'], *args, **kwargs) -> np.ndarray:
        pad_batch = 0
        if len(img) != self.input_shape[0]:
            pad_batch = self.input_shape[0] - len(img)
            for _ in range(pad_batch):
                img.append(np.zeros_like(img[0]))

        img_ = np.array(img, dtype=np.float32).transpose(0, 3, 1, 2)
        if list(img_.shape) != self.input_shape:
            raise ValueError('Map size not match net, expect', self.input_shape, ',got', img_.shape)

        result_npy = self.sess.run(None, {self.input_name: img_})
        if pad_batch != 0:
            return result_npy[0][0:len(img)]
        else:
            return result_npy[0]
