from .tensorrt import *
import numpy as np
__version__ = "5.0.6.3"

# Provides Python's `with` syntax
# Logger does not have a destructor.
ILogger.__enter__ = lambda this: this
ILogger.__exit__ = lambda this, exc_type, exc_value, traceback : None

Builder.__enter__ = lambda this: this
Builder.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

ICudaEngine.__enter__ = lambda this: this
ICudaEngine.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

IExecutionContext.__enter__ = lambda this: this
IExecutionContext.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

Runtime.__enter__ = lambda this: this
Runtime.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

INetworkDefinition.__enter__ = lambda this: this
INetworkDefinition.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

UffParser.__enter__ = lambda this: this
UffParser.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

CaffeParser.__enter__ = lambda this: this
CaffeParser.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

OnnxParser.__enter__ = lambda this: this
OnnxParser.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

IHostMemory.__enter__ = lambda this: this
IHostMemory.__exit__ = lambda this, exc_type, exc_value, traceback : this.__del__()

# Computes the volume of an iterable.
def volume(iterable):
    '''
    Computes the volume of an iterable.

    :arg iterable: Any python iterable, including a :class:`Dims` object.

    :returns: The volume of the iterable (0 for empty iterables).
    '''
    vol = 1
    for elem in iterable:
        vol *= elem
    return vol if iterable else 0

# Converts a TensorRT datatype to the equivalent numpy type.
def nptype(trt_type):
    '''
    Returns the numpy-equivalent of a TensorRT :class:`DataType` .

    :arg trt_type: The TensorRT data type to convert.

    :returns: The equivalent numpy type.
    '''
    if trt_type == float32:
        return np.float32
    elif trt_type == float16:
        return np.float16
    elif trt_type == int8:
        return np.int8
    elif trt_type == int32:
        return np.int32
    raise TypeError("Could not resolve TensorRT datatype to an equivalent numpy datatype.")

# Add a numpy-like itemsize property to the datatype.
def _itemsize(trt_type):
    '''
    Returns the size in bytes of this :class:`DataType` .

    :arg trt_type: The TensorRT data type.

    :returns: The size of the type.
    '''
    if trt_type == float32:
        return 4
    elif trt_type == float16:
        return 2
    elif trt_type == int8:
        return 1
    elif trt_type == int32:
        return 4

DataType.itemsize = property(lambda this: _itemsize(this))

# TODO: Load the plugin library, so that plugin_creator_list is availble in the plugin registry.
# from ctypes import *
# try:
#     cdll.LoadLibrary('libnvinfer_plugin.so')
# except OSError:
#     print("Make sure libnvinfer_plugin.so is in your LD_LIBRARY_PATH to use the plugin registry properly.")