# Copyright (c) 2023, NVIDIA CORPORATION.  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.

"""Helper functions for different DetectNet_v2 Models."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging
import os

import numpy as np

import pycuda.autoinit # noqa pylint: disable=unused-import
import pycuda.driver as cuda

try:
    import tensorrt as trt  # noqa pylint: disable=W0611 pylint: disable=W0611
    trt_available = True
    from nvidia_tao_tf1.core.export.data import TensorFile
except Exception as e:
    logger = logging.getLogger(__name__)
    logger.warning(
        "Failed to import TensorRT package, exporting TLT to a TensorRT engine "
        "will not be available."
    )
    trt_available = False

logger = logging.getLogger(__name__)


# Simple helper class to define host, device buffers.
class HostDeviceMem(object):
    """Clean data structure to handle host/device memory."""

    def __init__(self, host_mem, device_mem, name, npshape):
        """Initialize a HostDeviceMem data structure.

        Args:
            host_mem (cuda.pagelocked_empty): A cuda.pagelocked_empty memory buffer.
            device_mem (cuda.mem_alloc): Allocated memory pointer to the buffer in the GPU.
            name (str): Name of the binding blob in TensorRT.
            npshape (tuple): Shape of the input dimensions.

        Returns:
            HostDeviceMem instance.
        """
        self.host = host_mem
        self.device = device_mem
        self.numpy_shape = npshape
        self.name = name

    def __str__(self):
        """String containing pointers to the TRT Memory."""
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        """Return the canonical string representation of the object."""
        return self.__str__()


if trt_available:
    # Simple helper class for calibration.
    class Calibrator(trt.IInt8EntropyCalibrator2):
        """Calibrator class."""

        def __init__(self, data_filename, cache_filename,
                     n_batches, batch_size,
                     *args, **kwargs):
            """Init routine.

            This inherits from ``trt.IInt8EntropyCalibrator2`` to implement
            the calibration interface that TensorRT needs to calibrate the
            INT8 quantization factors.

            Args:
                data_filename (str): ``TensorFile`` data file to use.
                cache_filename (str): name of calibration file to read/write to.
                n_batches (int): number of batches for calibrate for.
                batch_size (int): batch size to use for calibration (this must be
                    smaller or equal to the batch size of the provided data).
            """
            super(Calibrator, self).__init__(*args, **kwargs)

            self._data_file = None
            if data_filename is not None and os.path.exists(data_filename):
                self._data_file = TensorFile(data_filename, 'r')
            else:
                logger.info("A valid tensorfile doesn't exist at {}. "
                            "The calibrator will attempt to read from a cache file "
                            "if provided.".format(self._data_file))
            self._cache_filename = cache_filename

            self._batch_size = batch_size
            self._n_batches = n_batches

            self._batch_count = 0
            self._data_mem = None

        def get_batch(self, names):
            """Return one batch.

            Args:
                names (list): list of memory bindings names.
            """
            if self._batch_count < self._n_batches:
                batch = np.array(self._data_file.read())
                if batch is not None:
                    # Adding a pylint error disable check due to a pylint issue
                    # TODO: <vpraveen> Remove this after pylint/issues/3139
                    batch_size = batch.shape[0]  # pylint: disable=E1136
                    if batch_size < self._batch_size:
                        raise ValueError("Data file batch size (%d) < request batch size (%d)" %
                                         (batch_size, self._batch_size))
                    batch = batch[:self._batch_size]

                    if self._data_mem is None:
                        self._data_mem = cuda.mem_alloc(batch.size * 4)  # 4 bytes per float32.

                    self._batch_count += 1

                    # Transfer input data to device.
                    cuda.memcpy_htod(self._data_mem, np.ascontiguousarray(batch, dtype=np.float32))
                    return [int(self._data_mem)]

            if self._data_mem is not None:
                self._data_mem.free()
            return None

        def get_batch_size(self):
            """Return batch size."""
            return self._batch_size

        def read_calibration_cache(self):
            """Read calibration from file."""
            if os.path.isfile(self._cache_filename):
                logger.warning("Calibration file exists at {}."
                               " Reading this cache.".format(self._cache_filename))
                with open(self._cache_filename, "rb") as cal_file:
                    return cal_file.read()
            return None

        def write_calibration_cache(self, cache):
            """Write calibration to file.

            Args:
                cache (memoryview): buffer to read calibration data from.
            """
            logger.info("Saving calibration cache (size %d) to %s",
                        len(cache), self._cache_filename)
            with open(self._cache_filename, 'wb') as f:
                f.write(cache)
