import warnings
from collections import namedtuple
from functools import partial
from pathlib import Path
from typing import List, Optional, Union

import numpy as np
import onnxruntime

try:
    import tensorrt as trt
except Exception:
    trt = None
import torch

warnings.filterwarnings(action='ignore', category=DeprecationWarning)


class TRTWrapper(torch.nn.Module):
    dtype_mapping = {}

    def __init__(self, weight: Union[str, Path],
                 device: Optional[torch.device]):
        super().__init__()
        weight = Path(weight) if isinstance(weight, str) else weight
        assert weight.exists() and weight.suffix in ('.engine', '.plan')
        if isinstance(device, str):
            device = torch.device(device)
        elif isinstance(device, int):
            device = torch.device(f'cuda:{device}')
        self.weight = weight
        self.device = device
        self.stream = torch.cuda.Stream(device=device)
        self.__update_mapping()
        self.__init_engine()
        self.__init_bindings()

    def __update_mapping(self):
        self.dtype_mapping.update({
            trt.bool: torch.bool,
            trt.int8: torch.int8,
            trt.int32: torch.int32,
            trt.float16: torch.float16,
            trt.float32: torch.float32
        })

    def __init_engine(self):
        logger = trt.Logger(trt.Logger.ERROR)
        self.log = partial(logger.log, trt.Logger.ERROR)
        trt.init_libnvinfer_plugins(logger, namespace='')
        self.logger = logger
        with trt.Runtime(logger) as runtime:
            model = runtime.deserialize_cuda_engine(self.weight.read_bytes())

        context = model.create_execution_context()

        names = [model.get_binding_name(i) for i in range(model.num_bindings)]

        num_inputs, num_outputs = 0, 0

        for i in range(model.num_bindings):
            if model.binding_is_input(i):
                num_inputs += 1
            else:
                num_outputs += 1

        self.is_dynamic = -1 in model.get_binding_shape(0)

        self.model = model
        self.context = context
        self.input_names = names[:num_inputs]
        self.output_names = names[num_inputs:]
        self.num_inputs = num_inputs
        self.num_outputs = num_outputs
        self.num_bindings = num_inputs + num_outputs
        self.bindings: List[int] = [0] * self.num_bindings

    def __init_bindings(self):
        Binding = namedtuple('Binding', ('name', 'dtype', 'shape'))
        inputs_info = []
        outputs_info = []

        for i, name in enumerate(self.input_names):
            assert self.model.get_binding_name(i) == name
            dtype = self.dtype_mapping[self.model.get_binding_dtype(i)]
            shape = tuple(self.model.get_binding_shape(i))
            inputs_info.append(Binding(name, dtype, shape))

        for i, name in enumerate(self.output_names):
            i += self.num_inputs
            assert self.model.get_binding_name(i) == name
            dtype = self.dtype_mapping[self.model.get_binding_dtype(i)]
            shape = tuple(self.model.get_binding_shape(i))
            outputs_info.append(Binding(name, dtype, shape))
        self.inputs_info = inputs_info
        self.outputs_info = outputs_info
        if not self.is_dynamic:
            self.output_tensor = [
                torch.empty(o.shape, dtype=o.dtype, device=self.device)
                for o in outputs_info
            ]

    def forward(self, *inputs):

        assert len(inputs) == self.num_inputs

        contiguous_inputs: List[torch.Tensor] = [
            i.contiguous() for i in inputs
        ]

        for i in range(self.num_inputs):
            self.bindings[i] = contiguous_inputs[i].data_ptr()
            if self.is_dynamic:
                self.context.set_binding_shape(
                    i, tuple(contiguous_inputs[i].shape))

        # create output tensors
        outputs: List[torch.Tensor] = []

        for i in range(self.num_outputs):
            j = i + self.num_inputs
            if self.is_dynamic:
                shape = tuple(self.context.get_binding_shape(j))
                output = torch.empty(
                    size=shape,
                    dtype=self.output_dtypes[i],
                    device=self.device)

            else:
                output = self.output_tensor[i]
            outputs.append(output)
            self.bindings[j] = output.data_ptr()

        self.context.execute_async_v2(self.bindings, self.stream.cuda_stream)
        self.stream.synchronize()

        return tuple(outputs)


class ORTWrapper(torch.nn.Module):

    def __init__(self, weight: Union[str, Path],
                 device: Optional[torch.device]):
        super().__init__()
        weight = Path(weight) if isinstance(weight, str) else weight
        assert weight.exists() and weight.suffix == '.onnx'

        if isinstance(device, str):
            device = torch.device(device)
        elif isinstance(device, int):
            device = torch.device(f'cuda:{device}')
        self.weight = weight
        self.device = device
        self.__init_session()
        self.__init_bindings()

    def __init_session(self):
        providers = ['CPUExecutionProvider']
        if 'cuda' in self.device.type:
            providers.insert(0, 'CUDAExecutionProvider')

        session = onnxruntime.InferenceSession(
            str(self.weight), providers=providers)
        self.session = session

    def __init_bindings(self):
        Binding = namedtuple('Binding', ('name', 'dtype', 'shape'))
        inputs_info = []
        outputs_info = []
        self.is_dynamic = False
        for i, tensor in enumerate(self.session.get_inputs()):
            if any(not isinstance(i, int) for i in tensor.shape):
                self.is_dynamic = True
            inputs_info.append(
                Binding(tensor.name, tensor.type, tuple(tensor.shape)))

        for i, tensor in enumerate(self.session.get_outputs()):
            outputs_info.append(
                Binding(tensor.name, tensor.type, tuple(tensor.shape)))
        self.inputs_info = inputs_info
        self.outputs_info = outputs_info
        self.num_inputs = len(inputs_info)

    def forward(self, *inputs):

        assert len(inputs) == self.num_inputs

        contiguous_inputs: List[np.ndarray] = [
            i.contiguous().cpu().numpy() for i in inputs
        ]

        if not self.is_dynamic:
            # make sure input shape is right for static input shape
            for i in range(self.num_inputs):
                assert contiguous_inputs[i].shape == self.inputs_info[i].shape

        outputs = self.session.run([o.name for o in self.outputs_info], {
            j.name: contiguous_inputs[i]
            for i, j in enumerate(self.inputs_info)
        })

        return tuple(torch.from_numpy(o).to(self.device) for o in outputs)