import sys
import torch
import tensorrt as trt
from torch2trt import torch2trt


def get_model(model_name, n_classes=None):
    if model_name == "yolov4":
        from models.detection import Yolov4
        model = Yolov4(n_classes=n_classes, inference=True)
    elif model_name == "yolov4-tiny":
        from models.detection import Yolov4Tiny
        model = Yolov4Tiny(n_classes=n_classes, inference=True)
    else:
        raise ValueError('Not implemented yet')

    return model


def transform_to_trt(model_name, weight_file, trt_output_name, max_batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W):
    model = get_model(model_name, n_classes)
    model.cuda()

    pretrained_dict = torch.load(weight_file, map_location=torch.device('cuda'))
    model.load_state_dict(pretrained_dict)

    input_names = ["input"]
    output_names = ['boxes', 'confs']

    x = torch.randn((max_batch_size, 3, IN_IMAGE_H, IN_IMAGE_W)).cuda()

    trt_model = torch2trt(model, [x], input_names=input_names, output_names=output_names, max_batch_size=max_batch_size, fp16_mode=True)

    y = model(x)
    y_trt = trt_model(x)

    # check the output against PyTorch
    print("difference between pytorch and tensorrt:" , torch.max(torch.abs(y - y_trt)))

    with open(trt_output_name, 'wb') as f:
        f.write(trt_model.engine.serialize())

if __name__ == '__main__':
    if len(sys.argv) == 8:
        model_name = sys.argv[1]
        weight_file = sys.argv[2]
        trt_output_name = sys.argv[3]
        max_batch_size = int(sys.argv[4])
        n_classes = int(sys.argv[5])
        IN_IMAGE_H = int(sys.argv[6])
        IN_IMAGE_W = int(sys.argv[7])

        print("Converting to TensorRT...")
        transform_to_trt(model_name, weight_file, trt_output_name, max_batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W)
    else:
        print('Please run this way:\n')
        print('  python pytorch2trt.py <model_name> <weight_file> <trt_output_name> <max_batch_size> <n_classes> <IN_IMAGE_H> <IN_IMAGE_W>')


import sys
import torch
import tensorrt as trt
from torch2trt import torch2trt


def get_model(model_name, n_classes=None):
    if model_name == "yolov4":
        from models.detection import Yolov4
        model = Yolov4(n_classes=n_classes, inference=True)
    elif model_name == "yolov4-tiny":
        from models.detection import Yolov4Tiny
        model = Yolov4Tiny(n_classes=n_classes, inference=True)
    else:
        raise ValueError('Not implemented yet')

    return model


def transform_to_trt(model_name, weight_file, trt_output_name, max_batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W):
    model = get_model(model_name, n_classes)
    model.cuda()

    pretrained_dict = torch.load(weight_file, map_location=torch.device('cuda'))
    model.load_state_dict(pretrained_dict)

    input_names = ["input"]
    output_names = ['boxes', 'confs']

    x = torch.randn((max_batch_size, 3, IN_IMAGE_H, IN_IMAGE_W)).cuda()

    trt_model = torch2trt(model, [x], input_names=input_names, output_names=output_names, max_batch_size=max_batch_size, fp16_mode=True)

    y = model(x)
    y_trt = trt_model(x)

    # check the output against PyTorch
    print("difference between pytorch and tensorrt:" , torch.max(torch.abs(y - y_trt)))

    with open(trt_output_name, 'wb') as f:
        f.write(trt_model.engine.serialize())

if __name__ == '__main__':
    if len(sys.argv) == 8:
        model_name = sys.argv[1]
        weight_file = sys.argv[2]
        trt_output_name = sys.argv[3]
        max_batch_size = int(sys.argv[4])
        n_classes = int(sys.argv[5])
        IN_IMAGE_H = int(sys.argv[6])
        IN_IMAGE_W = int(sys.argv[7])

        print("Converting to TensorRT...")
        transform_to_trt(model_name, weight_file, trt_output_name, max_batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W)
    else:
        print('Please run this way:\n')
        print('  python pytorch2trt.py <model_name> <weight_file> <trt_output_name> <max_batch_size> <n_classes> <IN_IMAGE_H> <IN_IMAGE_W>')


# check the output against PyTorch