import os
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


def set_requires_grad(nets, requires_grad=False):
        """Set requies_grad=Fasle for all the networks to avoid unnecessary computations
        Parameters:
            nets (network list)   -- a list of networks
            requires_grad (bool)  -- whether the networks require gradients or not
        """
        if not isinstance(nets, list):
            nets = [nets]
        for net in nets:
            if net is not None:
                for param in net.parameters():
                    param.requires_grad = requires_grad


if __name__ == "__main__": 
    # load model
    os.environ['CUDA_VISIBLE_DEVICES'] = '1' 
    checkpoint = '/fileser/gupc/code_update/grouplung/lungDetection3D/RibFractureDetector/rib_1019_3/checkpoints/145.pth'
    model = NoduleNet(config)
    model.eval()
    device = torch.device("cuda")
    model.to(device)
    checkpoint = torch.load(checkpoint)
    model.load_state_dict(checkpoint['state_dict'])
    print("model successfully loaded")


    # # 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.cuda().unsqueeze(0)
    # rpn_windows = make_rpn_windows(device, config)
    # print("data successfully generated")

    # load data
    input_data = torch.from_numpy(np.load('input_data.npy')).cuda()
    print(input_data.shape)
    print("data successfully loaded")


    # pth to torchscript
    model_path = 'NoduleNet_Rib.pt'
    model = model
    input_data = input_data
    traced_model = torch.jit.trace(model, (input_data))
    print(traced_model.graph)
    traced_model.save(model_path)
    


