from typing import Tuple, Dict,List
import tensorrt as trt
import numpy as np
import torch
import time
import os
from collections import OrderedDict, namedtuple

class TRT_Base():
    def __init__(self, 
                 input_shape: Tuple[int, int, int], 
                 model_path: str,
                 device: str='0'):
        """ Tensor RT base class for inference.

        Args:
            input_shape (Tuple[int, int, int]): image size (3, H, W)
            model_path (str): path to the model.trt
            device (str, optional): CUDA device. Defaults to '0'.
        """
        self.input_shape = input_shape
        self.model_path = model_path
        self.device = self.select_device(device)
        self.init_model()

    def select_device(self, device: str)->torch.device:
        """ Select device to be used for inference.
        Args:
            param device: 'cpu' or '0' or '0,1,2,3'
        Return: 
            torch.device
        """
        cpu = device.lower() == "cpu"
        if cpu:
            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
            return torch.device("cpu")
        else:
            assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested'
            torch.cuda.set_device(int(device))
            return torch.device("cuda")

    def init_model(self):
        """ Initialize TensorRT engine and context."""
        logger = trt.Logger(trt.Logger.INFO)
        trt.init_libnvinfer_plugins(logger, namespace="")
        with open(self.model_path, 'rb') as f, trt.Runtime(logger) as runtime:
            engine = runtime.deserialize_cuda_engine(f.read())
        context = engine.create_execution_context()
        self.model = {
            "engine": engine,
            "context": context
        }
        bindings, binding_addrs = self.get_bindings(input_shape=self.input_shape)
        input_names = [binding_name for binding_name in binding_addrs.keys() if (self.model["engine"].binding_is_input(binding_name))]
        
        for _ in range(10):
            for name in input_names:
                binding_addrs[name] = int(torch.randn(bindings[name].shape).to(self.device).data_ptr())
            context.execute_v2(list(binding_addrs.values()))
        self.model.update({
            'binding_addrs': binding_addrs,
            'bindings': bindings,
            'rt_shapes': self.input_shape
        })

    def get_bindings(self, input_shape: Tuple[int, int, int]):
        """ Get bindings and binding addresses for TensorRT engine.
        Args: 
            input_shape (Tuple[int, int, int]): image size (3, H, W)
        """
        self.model["context"].set_binding_shape(0, input_shape)
        bindings = OrderedDict()
        Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
        for index in range(self.model["engine"].num_bindings):
            name = self.model["engine"].get_binding_name(index)
            dtype = trt.nptype(self.model["engine"].get_binding_dtype(index))
            shape = tuple(self.model["context"].get_binding_shape(index))
            data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(self.device)
            bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
        binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
        return bindings, binding_addrs

    def change_runtime_dimension(self, input_shape: Tuple[int, int, int]):
        """ Support inference with Dynamic shape.

        Args:
            input_shape (Tuple[int, int, int]): image size (3, H, W)
        """
        if (input_shape == self.model["rt_shapes"]): return
        bindings, binding_addrs = self.get_bindings(input_shape)
        self.model['binding_addrs'] = binding_addrs
        self.model['bindings'] = bindings
        self.model['rt_shapes'] = input_shape