import tensorrt as trt
import torch
import mmap


def trt_dtype_to_torch(dtype):
    if dtype == trt.DataType.FLOAT:
        return torch.float32
    elif dtype == trt.DataType.INT32:
        return torch.int32
    elif dtype == trt.DataType.BOOL:
        return torch.bool
    else:
        raise RuntimeError(f"Unsupported dtype {dtype}")


class TRTOutputAllocator(trt.IOutputAllocator):
    def __init__(self, execution_context, dtype, device):
        trt.IOutputAllocator.__init__(self)
        self.execution_context = execution_context
        self.outputs = self.execution_context.outputs
        self.dtype = dtype
        self.device = device

    def reallocate_output(self, tensor_name, memory, size, alignment):
        if tensor_name in self.outputs:
            x = self.outputs[tensor_name]
            if x.numel() * x.element_size() == size:
                return x.data_ptr()
            del self.outputs[tensor_name]
        x = torch.zeros(size // 4, dtype=self.dtype, device=self.device)
        x.requires_grad = False
        self.outputs[tensor_name] = x
        return x.data_ptr()

    def notify_shape(self, tensor_name, shape):
        x = self.outputs[tensor_name]
        if x.shape != shape:
            x_ = x.view(list(shape))
            assert x_.data_ptr() == x.data_ptr()
            self.outputs[tensor_name] = x_


class TRTExecutionContext:
    def __init__(self, runtime, engine, optimization_profile_index, stream):
        self.runtime = runtime
        self.engine = engine
        self.stream = stream
        self.optimization_profile_index = None
        self.context = engine.create_execution_context()
        self.set_optimization_profile_index(optimization_profile_index)
        self.tensor_descs = []
        self.outputs = {}
        self.torch_device = torch.device("cuda", self.stream.device.index)
        for i in range(engine.num_io_tensors):
            name = engine.get_tensor_name(i)
            assert engine.get_tensor_location(name) == trt.TensorLocation.DEVICE
            dtype = engine.get_tensor_dtype(name)
            assert dtype in (
                trt.DataType.FLOAT,
                trt.DataType.INT32,
                trt.DataType.BOOL,
            )
            assert (
                engine.get_tensor_format(name, optimization_profile_index)
                == trt.TensorFormat.LINEAR
            )
            mode = engine.get_tensor_mode(name)
            self.tensor_descs.append(
                (
                    name,
                    mode,
                    engine.get_tensor_shape(name),
                    dtype,
                )
            )
            if mode == trt.TensorIOMode.OUTPUT:
                self.context.set_output_allocator(
                    name,
                    TRTOutputAllocator(
                        self, trt_dtype_to_torch(dtype), self.torch_device
                    ),
                )

    def set_optimization_profile_index(self, optimization_profile_index):
        if self.optimization_profile_index == optimization_profile_index:
            return
        self.stream.synchronize()
        self.optimization_profile_index = optimization_profile_index
        self.context.set_optimization_profile_async(
            optimization_profile_index, self.stream.cuda_stream
        )
        self.stream.synchronize()

    def eval_async(self, **kwargs):
        for i, (name, mode, shape, dtype) in enumerate(self.tensor_descs):
            if mode == trt.TensorIOMode.INPUT:
                if name not in kwargs:
                    raise RuntimeError(f"Missing input {name}")
                arg = kwargs[name]
                assert arg.device.type == "cuda"
                assert arg.device.index == self.stream.device.index
                assert arg.dtype == trt_dtype_to_torch(dtype)
                assert len(arg.shape) == len(shape)
                assert all(b == -1 or a == b for a, b in zip(arg.shape, shape))
                assert arg.is_contiguous()
                self.context.set_input_shape(name, arg.shape)
                self.context.set_tensor_address(name, arg.data_ptr())
        self.context.execute_async_v3(self.stream.cuda_stream)

    def eval_await(self):
        self.stream.synchronize()
        return self.outputs

    def eval(self, **kwargs):
        self.eval_async(**kwargs)
        return self.eval_await()


class TRTModel:
    def __init__(self, runtime, path):
        self.runtime = runtime
        with open(path, "rb") as f, mmap.mmap(
            f.fileno(), 0, access=mmap.ACCESS_READ
        ) as region:
            self.engine = runtime.deserialize_cuda_engine(region)
        assert self.engine

    def create_execution_context(self, optimization_profile_index=0, stream=None):
        if stream is None:
            stream = torch.cuda.Stream()
        return TRTExecutionContext(
            self.runtime, self.engine, optimization_profile_index, stream
        )

    def get_profile_index_for_dim_constraint(self, input_name, dim_index, max_size):
        candidates = []  # (profile_index, size)
        for i in range(self.engine.num_optimization_profiles):
            _, opt_shape, max_shape = self.engine.get_tensor_profile_shape(
                input_name, i
            )
            if max_shape[dim_index] >= max_size:
                candidates.append((i, opt_shape[dim_index]))
        if not candidates:
            raise RuntimeError(
                f"No profile found for {input_name} with dim {dim_index} at least {max_size}"
            )
        candidates.sort(key=lambda x: x[1])
        return candidates[0][0]

    def get_profile_max_size(self, profile_index, input_name, dim_index):
        _, _, max_shape = self.engine.get_tensor_profile_shape(
            input_name, profile_index
        )
        return max_shape[dim_index]