# Adapted from: https://github.com/PRBonn/bonnetal/blob/master/train/common/trtCalibINT8.py
import numpy as np
import cv2
import os

import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda

TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE)


class EntropyCalibrator(trt.IInt8EntropyCalibrator2):
    """
    Prepare callibartion dataset for int8 callibration.

    :param model_path: Path to the model
    :type model_path: str
    :param calibration_files: callibaration image paths
    :type calibration_files: list
    :param batch_size: batch size
    :type batch_size: int
    :param h: input size height
    :type h: int
    :param w: input size width
    :type w: int
    :param means: means
    :type means: float
    :param stds: standard deviation
    :type stds: float

    """

    def __init__(self, model_path, calibration_files, batch_size, h, w, means, stds):
        """
        :param model_path: directory path where model will be loaded from
        :type model_path: str
        :param calibration_files: path to calibration image files
        :type calibration_files: str
        :param batch_size: batch_size
        :type batch_size: int
        :param h: height
        :type h: int
        :param w: width
        :type w: int
        :param means: means
        :type means: np.array
        :param stds: standard deviation
        :type stds: np.array
        """
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.batch_size = batch_size
        self.h = h
        self.w = w
        self.input_size = [h, w]
        self.means = np.array(means, dtype=np.float32)
        self.stds = np.array(stds, dtype=np.float32)
        assert isinstance(calibration_files, list)
        self.calib_image_paths = calibration_files
        assert os.path.exists(model_path)
        self.cache_file = os.path.join(model_path, "model.trt.int8calib")
        self.shape = [self.batch_size, 3] + self.input_size
        self.device_input = cuda.mem_alloc(
            trt.volume(self.shape) * trt.float32.itemsize
        )
        self.indices = np.arange(len(self.calib_image_paths))
        np.random.shuffle(self.indices)

        def load_batches():
            """
            loads batches
            """
            for i in range(
                0, len(self.calib_image_paths) - self.batch_size + 1, self.batch_size
            ):
                indxes = self.indices[i : i + self.batch_size]
                paths = [self.calib_image_paths[i] for i in indxes]
                files = self.read_batch_file(paths)
                yield files

        self.batches = load_batches()

    def read_batch_file(self, filenames):
        """
        reads batch files

        :param filenames: list of file names
        :type filenames: List
        :returns: contiguous array
        :returntype: np.ascontiguousarray
        """
        tensors = []
        for filename in filenames:
            assert os.path.exists(filename)
            bgr_img = cv2.imread(filename)
            assert bgr_img.data
            bgr_img = cv2.resize(
                bgr_img, (self.w, self.h), interpolation=cv2.INTER_LINEAR
            )
            rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB).astype(np.float32)
            rgb_tensor = (rgb_img / 255.0 - self.means) / self.stds
            rgb_tensor = np.transpose(rgb_tensor, (2, 0, 1))
            rgb_tensor = np.ascontiguousarray(rgb_tensor, dtype=np.float32)
            tensors.append(rgb_tensor)
        return np.ascontiguousarray(tensors, dtype=np.float32)

    def get_batch_size(self):
        """
        returns batch size

        :returns: batch size
        :returntype: int
        """
        return self.batch_size

    # def get_batch(self, bindings, names):
    # https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Int8/EntropyCalibrator2.html#tensorrt.IInt8EntropyCalibrator2.get_batch
    def get_batch(self, names):
        """
        returns a batch

        :param names:

        :returns: List[batch from self.batches iter]
        :returntype: List
        """
        try:
            data = next(self.batches)
            cuda.memcpy_htod(self.device_input, data)
            return [int(self.device_input)]
        except StopIteration:
            return None

    def read_calibration_cache(self):
        """
        reads calibration cache

        :returns: file content from self.cache_file
        :returntype: file.read()
        """
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()

    def write_calibration_cache(self, cache):
        """
        writes cache to self.cache_file

        :param cache: cache to write
        :type cache: str
        """
        with open(self.cache_file, "wb") as f:
            f.write(cache)
