from backbone import Backbone
import torch
import argparse
import numpy as np
import time
import onnxruntime
import os


def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


def main(args):
    model = Backbone()
    if args.use_gpu:
        model.load_state_dict(torch.load(args.model_path), strict=True)
    else:
        model.load_state_dict(torch.load(args.model_path, map_location='cpu'), strict=True)

    device = 'cuda' if args.use_gpu else 'cpu'
    model = model.to(device)
    # print(model)

    model.eval()
    batch_size = 1
    x = torch.randn(batch_size, 3, args.input_size, args.input_size, requires_grad=True).to(device)
    torch_out = model(x)

    print("파이토치 모델 실행시간 측정")
    test_cnt = 5
    total_time = 0.
    for i in range(test_cnt):
        x = torch.randn(batch_size, 3, args.input_size, args.input_size, requires_grad=True).to(device)
        start = time.time()
        torch_out = model(x)
        total_time += time.time() - start
    pytorch_exec_time = total_time / test_cnt
    print("파이토치 모델 실행시간 측정 완료")

    os.makedirs(os.path.dirname(args.onnx_output_path), exist_ok=True)
    # 모델 변환
    torch.onnx.export(model,  # 실행될 모델
                      x,  # 모델 입력값 (튜플 또는 여러 입력값들도 가능)
                      args.onnx_output_path,  # 모델 저장 경로 (파일 또는 파일과 유사한 객체 모두 가능)
                      export_params=True,  # 모델 파일 안에 학습된 모델 가중치를 저장할지의 여부
                      # opset_version=11,          # 모델을 변환할 때 사용할 ONNX 버전
                      do_constant_folding=True,  # 최적하시 상수폴딩을 사용할지의 여부
                      input_names=['input'],  # 모델의 입력값을 가리키는 이름
                      output_names=['output'],  # 모델의 출력값을 가리키는 이름
                      # )
                      # dynamic_axes={'input': {2: 'height', 3: 'width'}}
                      )

    print("onnx 모델 컨버팅 완료")

    print("onnx 모델 로딩")
    ort_session = onnxruntime.InferenceSession(args.onnx_output_path)
    print("onnx 모델 로딩 완료")
    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
    ort_outs = ort_session.run(None, ort_inputs)

    print("onnx 모델 실행시간 측정")
    total_time = 0.
    for i in range(test_cnt):
        ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}

        start = time.time()
        ort_outs = ort_session.run(None, ort_inputs)
        total_time += time.time() - start
    onnx_exec_time = total_time / test_cnt
    print("파이토치 모델 실행시간 측정 완료")

    np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
    print("Exported model has been tested with ONNXRuntime, and the result looks good!")

    print("onnx 평균시간", onnx_exec_time, "파이토치 평균시간", pytorch_exec_time)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='for common image metric learning')

    parser.add_argument('--input_size', type=int, default=112)

    parser.add_argument('--model_path', default=None, type=str)
    parser.add_argument('--onnx_output_path', default=None, type=str)
    parser.add_argument('--use_gpu', default=False, action="store_true")

    main(parser.parse_args())
