import os
import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt_runtime = trt.Runtime(TRT_LOGGER)

def build_engine(max_batch_size=1, onnx_file_path="", engine_file_path="", fp16_mode=False, int8_mode=False, save_engine=False):
    """Takes an ONNX file and creates a TensorRT engine to run inference with"""
    explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    with trt.Builder(TRT_LOGGER) as builder, \
            builder.create_network(explicit_batch) as network, \
            trt.OnnxParser(network, TRT_LOGGER) as parser:

        builder.max_workspace_size = 1 << 30  # Your workspace size
        builder.max_batch_size = max_batch_size
        # pdb.set_trace()
        builder.fp16_mode = fp16_mode  # Default: False
        builder.int8_mode = int8_mode  # Default: False
        if int8_mode:
            # To be updated
            raise NotImplementedError

        # Parse model file
        if not os.path.exists(onnx_file_path):
            quit('ONNX file {} not found'.format(onnx_file_path))

        print('Loading ONNX file from path {}...'.format(onnx_file_path))
        with open(onnx_file_path, 'rb') as model:
            print('Beginning ONNX file parsing')
            parser.parse(model.read())

        print('Completed parsing of ONNX file')
        print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))


        last_layer = network.get_layer(network.num_layers - 1)
        network.mark_output(last_layer.get_output(0))

        engine = builder.build_cuda_engine(network)
        print("Completed creating Engine")

        if save_engine:
            with open(engine_file_path, "wb") as f:
                f.write(engine.serialize())
        return engine

engine = build_engine(max_batch_size=1, onnx_file_path="models/onnx/centerface_1088_1920.onnx", engine_file_path="models/tensorrt/centerface.trt", save_engine=True)
