# from tensorrtserver.api import *
import os
import numpy as np
import torch
import timeit
import numpy as np
from config import config
from net.nodule_net import NoduleNet
from dataset.bbox_reader import BboxReader
from net.layer.rpn_nms import make_rpn_windows
# import onnxruntime as ort


if __name__ == "__main__": 
    os.environ['CUDA_VISIBLE_DEVICES'] = '1' 
    
    # # generate data
    # data_dir = config['preprocessed_data_dir']
    # test_set_name = "hahaha"
    # dataset = BboxReader(data_dir, test_set_name, config, mode='train')
    # input_data = dataset[50][0]
    # input_data = input_data.unsqueeze(0)
    # device = torch.device("cuda")
    # rpn_windows = make_rpn_windows(device, config).cpu().numpy()
    # np.save('input_data.npy',input_data)
    # np.save('rpn_windows.npy',rpn_windows)
    # print("data successfully generated")


    # load data
    input_data = np.load('input_data.npy')
    rpn_windows = np.load('rpn_windows.npy')
    print("data successfully loaded")

    # pytorth infer
    input_data = torch.from_numpy(input_data).cuda()
    rpn_windows = torch.from_numpy(rpn_windows).cuda()
    checkpoint = '/fileser/gupc/NoduleNet/results/proxima_data_without_lungmask/model/009.pth'
    model = NoduleNet(config)
    model.eval()
    device = torch.device("cuda")
    model.to(device)
    checkpoint = torch.load(checkpoint)
    model.load_state_dict(checkpoint['state_dict'])
    logits,deltas = model(input_data,rpn_windows)
    print("-------------------pytorch infer-------------------")
    print("logits: {}".format(logits.detach().cpu().numpy()))
    print("deltas: {}".format(deltas.detach().cpu().numpy()))



    # torchscript infer
    # input_data = torch.from_numpy(input_data).cuda()
    # rpn_windows = torch.from_numpy(rpn_windows).cuda()
    model_path = 'NoduleNet.pt'
    model = torch.jit.load(model_path)
    logits,deltas = model(input_data,rpn_windows)
    print("-------------------torchscript infer-------------------")
    print("logits: {}".format(logits.detach().cpu().numpy()))
    print("deltas: {}".format(deltas.detach().cpu().numpy()))



    # # trtserver infer
    # url = "10.100.37.12:19012"
    # protocol = ProtocolType.from_str("http")
    # model_name = "nodulenet"
    # model_version = -1
    # batch_size = 1
    # ctx = InferContext(url, protocol, model_name, model_version,False)
    # result = ctx.run({ 'INPUT__0': (input_data,),
    #                    'INPUT__1': (rpn_windows,)},
    #                  { 'OUTPUT__0': InferContext.ResultFormat.RAW,
    #                    'OUTPUT__1': InferContext.ResultFormat.RAW,},
    #                    batch_size)
    # print("-------------------trtserver infer-------------------")
    # print("logits: {}".format(result['OUTPUT__0']))
    # print("deltas: {}".format(result['OUTPUT__1']))



    # # onnxruntime infer
    # model_path = 'NoduleNet.onnx'
    # sess = ort.InferenceSession(model_path)
    # logits,deltas = sess.run(None, {sess.get_inputs()[0].name: input_data,
    #                                 sess.get_inputs()[1].name: rpn_windows,})
    # print("-------------------onnxruntime infer-------------------")
    # print("logits: {}".format(logits))
    # print("deltas: {}".format(deltas))