import math
import os
import random
import sys
import traceback
import warnings

# import horovod.torch as hvd
import numpy as np
import pandas as pd
import torch
import torch.backends.cudnn as cudnn
from apex import amp
from tensorboardX import SummaryWriter
from torch.autograd import Variable
import torch.distributed
from torch.utils.data import DataLoader
from tqdm import tqdm
import torch.distributed as dist

from ..data import BboxReader
from ..data.collate import train_collate
from ..modeling.detector.centernet import CenterNet
from ..modeling.detector.nodulenet import NoduleNet
import copy
warnings.filterwarnings("ignore")
this_module = sys.modules[__name__]

# Set seed
SEED = 35202
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)


class DistributedSampler_new(torch.utils.data.distributed.DistributedSampler):
    def update_dataset(self, dataset):
        self.dataset = dataset
        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
        self.total_size = self.num_samples * self.num_replicas


# Horovod: average metrics from distributed training.
class Metric(object):
    def __init__(self, name):
        self.name = name
        self.sum = torch.tensor(0.)
        self.n = torch.tensor(0.)

    def update(self, val):
        # val = val.clone().detach()
        # output_tensors = [val.clone().to(val) for _ in range(torch.distributed.get_world_size())]
        # print('world_size: ', torch.distributed.get_world_size())
        if os.environ.get('RANK', -1) != -1:
            output_tensors = [val.clone() for _ in range(torch.distributed.get_world_size())]
            dist.all_gather(output_tensors, val)
            self.sum += torch.sum(torch.stack(output_tensors)).cpu().detach()
        else:
            self.sum += val.cpu().detach()
        self.n += 1

    @property
    def avg(self):
        return self.sum / self.n


def save_checkpoint(epoch, model, optimizer, model_out_dir):
    state_dict = model.state_dict()
    for key in state_dict.keys():
        state_dict[key] = state_dict[key].cpu()
    
    torch.save({
        'epoch': epoch,
        'out_dir': model_out_dir,
        'state_dict': state_dict,
        'optimizer': optimizer.state_dict()},
        model_out_dir)


class BaseDetection3D(object):
    def __init__(self, cfg, mode='train', log_fun=None):
        """3d detector initialization
        
        Arguments:
            DetectorBase {[object]} -- detector class providing basic detector functions
            config {[dict]} -- [config dict with hyper parameters]
        
        Raises:
            ValueError: only valid frames are implemented
        """
        # Init network
        if cfg.model['meta_architecture'] == 'NoduleNet':
            self.model = NoduleNet(cfg, mode).cuda()
        elif cfg.model['meta_architecture'] == 'CenterNet':
            self.model = CenterNet(cfg, mode).cuda()
        else:
            raise ValueError('Detector architecture %s is not implemented.' % cfg.model['meta_architecture'])

        self.cfg = cfg
        self.DataReader = BboxReader
        self.DistributedSampler = DistributedSampler_new

        if log_fun is None:
            self.log_fun = print
        else:
            self.log_fun = log_fun

        # 设置 torch.backends.cudnn.benchmark=True 将会让程序在开始时花费一点额外时间，
        # 为整个网络的每个卷积层搜索最适合它的卷积实现算法，进而实现网络的加速。
        # 适用场景是网络结构固定（不是动态变化的），网络的输入形状（包括 batch size，图片大小，输入的通道）是不变的，
        # 其实也就是一般情况下都比较适用。反之，如果卷积层的设置一直变化，将会导致程序不停地做优化，反而会耗费更多的时间。
        cudnn.benchmark = False

    def _create_dataloader(self):
        self.train_dataset = self.DataReader(self.cfg, mode='train', log_fun=self.log_fun)
        self.val_dataset = self.DataReader(self.cfg, mode='val', log_fun=self.log_fun)

        # Horovod: limit # of CPU threads to be used per worker.
        torch.set_num_threads(4)

        kwargs = {'num_workers': self.cfg.num_workers, 'pin_memory': True}
        # When supported, use 'forkserver' to spawn dataloader workers instead of 'fork' to prevent
        # issues with Infiniband implementations that are not fork-safe
        # if (kwargs.get('num_workers', 0) > 0 and hasattr(mp, '_supports_context') and
        #         mp._supports_context and 'forkserver' in mp.get_all_start_methods()):
        #     kwargs['multiprocessing_context'] = 'forkserver'

        # Partition dataset among workers using DistributedSampler
        self.train_sampler = self.DistributedSampler(self.train_dataset, num_replicas=self.cfg.world_size, rank=self.cfg.rank)
        self.val_sampler = self.DistributedSampler(self.val_dataset, num_replicas=self.cfg.world_size, rank=self.cfg.rank)
        self.train_loader = DataLoader(self.train_dataset, batch_size=self.cfg.batch_size,
                                       collate_fn=train_collate,
                                       sampler=self.train_sampler, **kwargs)
        self.val_loader = DataLoader(self.val_dataset, batch_size=self.cfg.batch_size,
                                     collate_fn=train_collate,
                                     sampler=self.val_sampler, **kwargs)

    def _create_optimizer(self):
        self.optimizer = getattr(torch.optim, self.cfg.training['solver']['optimizer'])
        # Horovod: lr = base_lr * hvd.size
        # special_lr_layers = ['feature_net.back3.1.conv1.conv_offset.weight', 
        #                      'feature_net.back3.1.conv1.conv_offset.bias', 
        #                      'feature_net.back3.1.conv1.conv_mask.weight',
        #                      'feature_net.back3.1.conv1.conv_mask.bias',
        #                      'feature_net.back3.1.conv2.conv_offset.weight', 
        #                      'feature_net.back3.1.conv2.conv_offset.bias',
        #                      'feature_net.back3.1.conv2.conv_mask.weight',
        #                      'feature_net.back3.1.conv2.conv_mask.bias']
        # special_params = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] in special_lr_layers, self.model.named_parameters()))))
        # base_params = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] not in special_lr_layers, self.model.named_parameters()))))
        # self.optimizer = self.optimizer([{"params": base_params},
        #                                  {"params": special_params, "lr": self.cfg.TRAINING.SOLVER.BASE_LR * hvd.size() * 0.05}], 
        #                                  lr=self.cfg.TRAINING.SOLVER.BASE_LR * hvd.size(),
        #                                  weight_decay=self.cfg.TRAINING.SOLVER.WRIGHT_DECAY, momentum=self.cfg.TRAINING.SOLVER.MOMENTUM)
        if self.cfg.training['solver']['optimizer'] == 'Adam':
            self.optimizer = self.optimizer(self.model.parameters(), lr=self.cfg.lr * self.cfg.world_size,
                                            weight_decay=self.cfg.training['solver']['wright_decay'])
        elif self.cfg.training['solver']['optimizer'] == 'SGD':
            self.optimizer = self.optimizer(self.model.parameters(), lr=self.cfg.lr * self.cfg.world_size,
                                            weight_decay=self.cfg.training['solver']['wright_decay'],
                                            momentum=self.cfg.training['solver']['momentum'])
        else:
            raise ValueError("Optimizer %s is not implemented." % self.cfg.training['solver']['optimizer'])
        # Horovod: (optional) compression algorithm.
        # compression = hvd.Compression.fp16 if self.cfg.training['fp16_allreduce'] else hvd.Compression.none
        # Horovod: Add Horovod Distributed Optimizer
        # self.optimizer = hvd.DistributedOptimizer(self.optimizer, named_parameters=self.model.named_parameters(),
        #                                           compression=compression)
        # self.optimizer._requires_update = set()

        # Horovod: load model only for rank 0
        self.resume_from_epoch = 0
        if self.cfg.pretrain_msg and self.cfg.rank == 0:
            # print(os.path.isfile(self.cfg.pretrain_msg))
            print('[Loading model from %s]' % self.cfg.pretrain_msg)
            # checkpoint = torch.load(self.cfg.pretrain_msg)
            # self.resume_from_epoch = checkpoint['epoch']
            # state = self.model.state_dict()
            # state.update(checkpoint['state_dict'])
            # self.model.load_state_dict(state)  # , strict=False
            checkpoint = torch.load(self.cfg.pretrain_msg)
            # print('load weight feature_net.preBlock.0.weight: ', checkpoint['feature_net.preBlock.0.weight'])
            ckpt_dict = checkpoint['state_dict']
            # print(ckpt_dict)
            model_dict = self.model.state_dict()
            # print('self.model feature_net.preBlock.0.weight: ', checkpoint['feature_net.preBlock.0.weight'])
            new_state_dict = {}
            for k, v in ckpt_dict.items():
                if 'rpn.' in k:
                    new_k = k.replace('rpn.', 'rpn_head.')
                    new_state_dict[new_k] = v
                elif 'rcnn_head.' in k:
                    new_k = k.replace('rcnn_head.', 'box_head.')
                    new_state_dict[new_k] = model_dict[new_k]
                else:
                    new_state_dict[k] = v
                print(k)
            self.model.load_state_dict(new_state_dict)
            self.model_init = copy.deepcopy(self.model)
            print('Load success')
            
            # self.optimizer.load_state_dict(checkpoint['optimizer'])
            # self.resume_from_epoch += 1

        # Horovod: broadcast resume_from_epoch from rank 0 (which will have
        # checkpoints) to other ranks.
        # self.resume_from_epoch = hvd.broadcast(torch.tensor(self.resume_from_epoch), root_rank=0,
        #                                        name='resume_from_epoch').item()

        # Horovod: Broadcast parameters from rank 0 to all other processes.
        # hvd.broadcast_parameters(self.model.state_dict(), root_rank=0)
        # hvd.broadcast_optimizer_state(self.optimizer, root_rank=0)

        if self.cfg.training['amp']:
            # Apex
            amp.register_float_function(torch, 'sigmoid')
            self.model, optimizer = amp.initialize(self.model, self.optimizer,
                                                   opt_level="O1")  # , keep_batchnorm_fp32=True
        # Horovod: print output only on first rank.
        if self.cfg.rank == 0:
            # self.model_out_dir = os.path.join(self.cfg.training['saver']['saver_dir'], 'checkpoints')
            tb_out_dir = os.path.join(self.cfg.training['saver']['saver_root'], 'tensorboards')
            # if not os.path.exists(self.model_out_dir):
                # os.makedirs(self.model_out_dir)
            if not os.path.exists(tb_out_dir):
                os.makedirs(tb_out_dir)

            self.log_fun('[start_epoch %d, out_dir %s]' % (self.resume_from_epoch, self.cfg.training['saver']['saver_root']))
            self.log_fun('[length of train loader %d, length of valid loader %d]' % (
                len(self.train_loader), len(self.val_loader)))

        # Write graph to tensorboard for visualization
        self.writer = SummaryWriter(tb_out_dir) if self.cfg.rank == 0 else None
        self.train_writer = SummaryWriter(os.path.join(tb_out_dir, 'train')) if self.cfg.rank == 0 else None
        self.val_writer = SummaryWriter(os.path.join(tb_out_dir, 'val')) if self.cfg.rank == 0 else None

    def do_train(self):
        self.model.loss.cuda()

        self._create_dataloader()
        self._create_optimizer()

        verbose = 1 if self.cfg.rank == 0 else 0

        for i in range(self.resume_from_epoch, self.cfg.epochs):
            if self.cfg.training['sheduler']['num_neg_shedule']:
                self.num_neg_shedule(i)

            if self.cfg.training['sheduler']['rand_crop_ratio_shedule']:
                self.rand_crop_ratio_shedule(i)

            if i >= self.cfg.training['sheduler']['switch_roi_epoch']:
                self.model.use_rcnn = True
                self.model.loss.use_rcnn = True

            if self.cfg.data['data_loader']['balanced_sampling']:
                self.train_loader.dataset.list_IDs = self.train_loader.dataset._get_list_ids()
                self.train_loader.sampler.update_dataset(self.train_loader.dataset)

            if i == self.cfg.training['sheduler']['switch_balanced_sampling_epoch']:
                self.cfg.data['data_loader']['balanced_sampling'] = True
                self.train_loader.dataset.balanced_sampling_flag = True
                self.train_loader.dataset.list_IDs = self.train_loader.dataset._get_list_ids()
                self.train_loader.sampler.update_dataset(self.train_loader.dataset)
                if self.cfg.rank == 0:
                    print('[length of train loader %d, length of valid loader %d]' % (
                        len(self.train_loader), len(self.val_loader)))

            self.train(i, verbose)

            self.validate(i, verbose)

            if self.cfg.rank == 0:
                save_checkpoint(i, self.cfg.training['saver']['saver_frequency'], self.model, self.optimizer,
                                self.model_out_dir)

        if self.writer:
            self.writer.close()
        if self.train_writer:
            self.train_writer.close()
        if self.val_writer:
            self.val_writer.close()

    def train(self, epoch, verbose):
        self.train_sampler.set_epoch(epoch)
        self.model.set_mode('train')
        rpn_cls_loss = Metric('rpn_cls_loss')
        rpn_reg_loss = Metric('rpn_reg_loss')
        rcnn_cls_loss = Metric('rcnn_cls_loss')
        rcnn_reg_loss = Metric('rcnn_reg_loss')
        total_loss = Metric('train_loss')
        rpn_stats = []
        rcnn_stats = []

        with tqdm(total=len(self.train_loader),
                  desc='Train Epoch     #{}'.format(epoch),
                  disable=not verbose) as t:
            for j, (input_data, truth_bboxes) in enumerate(self.train_loader):
                input_data = Variable(input_data).cuda()

                if self.cfg.training['sheduler']['lr_shedule']:
                    self.lr_shedule(epoch, j)

                self.optimizer.zero_grad()
                rpn_logits, rpn_deltas, rpn_labels, rpn_label_weights, rpn_targets, rpn_targets_weights, \
                rcnn_logits, rcnn_deltas, rcnn_labels, rcnn_targets = self.model(input_data, truth_bboxes)

                loss, [rpn_cls_loss_, rpn_reg_loss_, rcnn_cls_loss_, rcnn_reg_loss_], rpn_stat_, rcnn_stat_ = \
                    self.model.loss(rpn_logits, rpn_deltas, rpn_labels, rpn_label_weights,
                                    rpn_targets, rpn_targets_weights, rcnn_logits, rcnn_deltas, rcnn_labels,
                                    rcnn_targets)
                if self.cfg.training['amp'].AMP:
                    with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                        scaled_loss.backward()
                        self.optimizer.synchronize()
                    with self.optimizer.skip_synchronize():
                        # torch.nn.utils.clip_grad_norm_(amp.master_params(self.optimizer), max_norm = 1.0, norm_type=2)
                        self.optimizer.step()
                else:
                    loss.backward()
                    self.optimizer.step()

                rpn_cls_loss.update(rpn_cls_loss_)
                rpn_reg_loss.update(rpn_reg_loss_)
                rcnn_cls_loss.update(rcnn_cls_loss_)
                rcnn_reg_loss.update(rcnn_reg_loss_)
                total_loss.update(loss)
                rpn_stats.append(rpn_stat_)
                rcnn_stats.append(rcnn_stat_)

                t.set_postfix({'total_loss': total_loss.avg.item(),
                               'rpn_cls_loss': rpn_cls_loss.avg.item(),
                               'rpn_reg_loss': rpn_reg_loss.avg.item(),
                               'rcnn_cls_loss': rcnn_cls_loss.avg.item(),
                               'rcnn_reg_loss': rcnn_reg_loss.avg.item(), })
                t.update(1)

        if self.train_writer:
            rpn_stats = np.asarray(rpn_stats, np.float32)
            print('Train Epoch %d, loss %f' % (epoch, total_loss.avg.item()))
            print('rpn_cls %f, rpn_reg %f, rcnn_cls %f, rcnn_reg %f' % (
                rpn_cls_loss.avg.item(), rpn_reg_loss.avg.item(), rcnn_cls_loss.avg.item(), rcnn_reg_loss.avg.item()))
            print('rpn_stats: tpr %f, tnr %f, total pos %d, total neg %d, reg %.4f, %.4f, %.4f, %.4f, %.4f, %.4f' % (
                100.0 * np.sum(rpn_stats[:, 0]) / np.sum(rpn_stats[:, 1]),
                100.0 * np.sum(rpn_stats[:, 2]) / np.sum(rpn_stats[:, 3]),
                np.sum(rpn_stats[:, 1]),
                np.sum(rpn_stats[:, 3]),
                np.mean(rpn_stats[:, 4]),
                np.mean(rpn_stats[:, 5]),
                np.mean(rpn_stats[:, 6]),
                np.mean(rpn_stats[:, 7]),
                np.mean(rpn_stats[:, 8]),
                np.mean(rpn_stats[:, 9])))

            # Write to tensorboard
            self.train_writer.add_scalar('loss', total_loss.avg, epoch)
            self.train_writer.add_scalar('rpn_cls', rpn_cls_loss.avg, epoch)
            self.train_writer.add_scalar('rpn_reg', rpn_reg_loss.avg, epoch)
            self.train_writer.add_scalar('rcnn_cls', rcnn_cls_loss.avg, epoch)
            self.train_writer.add_scalar('rcnn_reg', rcnn_reg_loss.avg, epoch)

            self.train_writer.add_scalar('rpn_reg_z', np.mean(rpn_stats[:, 4]), epoch)
            self.train_writer.add_scalar('rpn_reg_y', np.mean(rpn_stats[:, 5]), epoch)
            self.train_writer.add_scalar('rpn_reg_x', np.mean(rpn_stats[:, 6]), epoch)
            self.train_writer.add_scalar('rpn_reg_d', np.mean(rpn_stats[:, 7]), epoch)
            self.train_writer.add_scalar('rpn_reg_h', np.mean(rpn_stats[:, 8]), epoch)
            self.train_writer.add_scalar('rpn_reg_w', np.mean(rpn_stats[:, 9]), epoch)

            if self.model.use_rcnn:
                rcnn_stats = np.asarray([stat[:-1] for stat in rcnn_stats], np.float32)

                print('rcnn_stats: reg %.4f, %.4f, %.4f, %.4f, %.4f, %.4f' % (
                    np.mean(rcnn_stats[:, 0]),
                    np.mean(rcnn_stats[:, 1]),
                    np.mean(rcnn_stats[:, 2]),
                    np.mean(rcnn_stats[:, 3]),
                    np.mean(rcnn_stats[:, 4]),
                    np.mean(rcnn_stats[:, 5])))
                self.train_writer.add_scalar('rcnn_reg_z', np.mean(rcnn_stats[:, 0]), epoch)
                self.train_writer.add_scalar('rcnn_reg_y', np.mean(rcnn_stats[:, 1]), epoch)
                self.train_writer.add_scalar('rcnn_reg_x', np.mean(rcnn_stats[:, 2]), epoch)
                self.train_writer.add_scalar('rcnn_reg_d', np.mean(rcnn_stats[:, 3]), epoch)
                self.train_writer.add_scalar('rcnn_reg_h', np.mean(rcnn_stats[:, 4]), epoch)
                self.train_writer.add_scalar('rcnn_reg_w', np.mean(rcnn_stats[:, 5]), epoch)

        torch.cuda.empty_cache()

    def validate(self, epoch, verbose):
        # self.val_sampler.set_epoch(epoch)
        self.model.set_mode('valid')
        rpn_cls_loss = Metric('rpn_cls_loss')
        rpn_reg_loss = Metric('rpn_reg_loss')
        rcnn_cls_loss = Metric('rcnn_cls_loss')
        rcnn_reg_loss = Metric('rcnn_reg_loss')
        total_loss = Metric('train_loss')
        rpn_stats = []
        rcnn_stats = []

        with tqdm(total=len(self.val_loader),
                  desc='Validate  Epoch     #{}'.format(epoch),
                  disable=not verbose) as t:
            for j, (input_data, truth_bboxes) in enumerate(self.val_loader):
                with torch.no_grad():
                    input_data = Variable(input_data).cuda()

                    rpn_logits, rpn_deltas, rpn_labels, rpn_label_weights, rpn_targets, rpn_targets_weights, \
                    rcnn_logits, rcnn_deltas, rcnn_labels, rcnn_targets = self.model(input_data, truth_bboxes)

                    loss, [rpn_cls_loss_, rpn_reg_loss_, rcnn_cls_loss_, rcnn_reg_loss_], rpn_stat_, rcnn_stat_ = \
                        self.model.loss(rpn_logits, rpn_deltas, rpn_labels, rpn_label_weights, rpn_targets,
                                        rpn_targets_weights,
                                        rcnn_logits, rcnn_deltas, rcnn_labels, rcnn_targets)

                    rpn_cls_loss.update(rpn_cls_loss_)
                    # if torch.is_nonzero(rpn_reg_loss_):
                    rpn_reg_loss.update(rpn_reg_loss_)
                    rcnn_cls_loss.update(rcnn_cls_loss_)
                    rcnn_reg_loss.update(rcnn_reg_loss_)
                    total_loss.update(loss)
                    rpn_stats.append(rpn_stat_)
                    rcnn_stats.append(rcnn_stat_)

        if self.val_writer:
            rpn_stats = np.asarray(rpn_stats, np.float32)
            print('Val Epoch %d, loss %f' % (epoch, total_loss.avg.item()))
            print('rpn_cls %f, rpn_reg %f, rcnn_cls %f, rcnn_reg %f' % (
                rpn_cls_loss.avg.item(), rpn_reg_loss.avg.item(), rcnn_cls_loss.avg.item(), rcnn_reg_loss.avg.item()))
            print('rpn_stats: tpr %f, tnr %f, total pos %d, total neg %d, reg %.4f, %.4f, %.4f, %.4f, %.4f, %.4f' % (
                100.0 * np.sum(rpn_stats[:, 0]) / np.sum(rpn_stats[:, 1]),
                100.0 * np.sum(rpn_stats[:, 2]) / np.sum(rpn_stats[:, 3]),
                np.sum(rpn_stats[:, 1]),
                np.sum(rpn_stats[:, 3]),
                np.mean(rpn_stats[:, 4]),
                np.mean(rpn_stats[:, 5]),
                np.mean(rpn_stats[:, 6]),
                np.mean(rpn_stats[:, 7]),
                np.mean(rpn_stats[:, 8]),
                np.mean(rpn_stats[:, 9])))

            # Write to tensorboard
            self.val_writer.add_scalar('loss', total_loss.avg, epoch)
            self.val_writer.add_scalar('rpn_cls', rpn_cls_loss.avg, epoch)
            self.val_writer.add_scalar('rpn_reg', rpn_reg_loss.avg, epoch)
            self.val_writer.add_scalar('rcnn_cls', rcnn_cls_loss.avg, epoch)
            self.val_writer.add_scalar('rcnn_reg', rcnn_reg_loss.avg, epoch)

            self.val_writer.add_scalar('rpn_reg_z', np.mean(rpn_stats[:, 4]), epoch)
            self.val_writer.add_scalar('rpn_reg_y', np.mean(rpn_stats[:, 5]), epoch)
            self.val_writer.add_scalar('rpn_reg_x', np.mean(rpn_stats[:, 6]), epoch)
            self.val_writer.add_scalar('rpn_reg_d', np.mean(rpn_stats[:, 7]), epoch)
            self.val_writer.add_scalar('rpn_reg_h', np.mean(rpn_stats[:, 8]), epoch)
            self.val_writer.add_scalar('rpn_reg_w', np.mean(rpn_stats[:, 9]), epoch)

            if self.model.use_rcnn:
                rcnn_stats = np.asarray([stat[:-1] for stat in rcnn_stats], np.float32)

                print('rcnn_stats: reg %.4f, %.4f, %.4f, %.4f, %.4f, %.4f' % (
                    np.mean(rcnn_stats[:, 0]),
                    np.mean(rcnn_stats[:, 1]),
                    np.mean(rcnn_stats[:, 2]),
                    np.mean(rcnn_stats[:, 3]),
                    np.mean(rcnn_stats[:, 4]),
                    np.mean(rcnn_stats[:, 5])))
                self.val_writer.add_scalar('rcnn_reg_z', np.mean(rcnn_stats[:, 0]), epoch)
                self.val_writer.add_scalar('rcnn_reg_y', np.mean(rcnn_stats[:, 1]), epoch)
                self.val_writer.add_scalar('rcnn_reg_x', np.mean(rcnn_stats[:, 2]), epoch)
                self.val_writer.add_scalar('rcnn_reg_d', np.mean(rcnn_stats[:, 3]), epoch)
                self.val_writer.add_scalar('rcnn_reg_h', np.mean(rcnn_stats[:, 4]), epoch)
                self.val_writer.add_scalar('rcnn_reg_w', np.mean(rcnn_stats[:, 5]), epoch)

        torch.cuda.empty_cache()

    def do_test(self):
        initial_checkpoint = self.cfg.testing['weight']
        save_dir = self.cfg.testing['saver_dir']
        if initial_checkpoint:
            print('[Loading model from %s]' % initial_checkpoint)
            checkpoint = torch.load(initial_checkpoint)
            self.model.load_state_dict(checkpoint['state_dict'])

            # state = self.model.state_dict()
            # state.update(checkpoint['state_dict'])
            # self.model.load_state_dict(state) #,strict=False
            epoch = checkpoint['epoch']
        else:
            print('No model weight file specified')
            return
        self.model.set_mode('test')
        self.model.use_rcnn = self.cfg.testing['use_rcnn']

        # model_path = self.cfg.DEPLOY.TORCHSCRIPT_SAVE_PATH
        # print('[Loading torchscript from %s]' % model_path)
        # torchscript_model = torch.jit.load(model_path)

        self.test_dataset = self.DataReader(self.cfg, mode='test')
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        res_dir = os.path.join(save_dir, str(epoch))
        if not os.path.exists(res_dir):
            os.makedirs(res_dir)

        print('Total # of eval data %d' % (len(self.test_dataset)))
        for i, (input, image) in enumerate(self.test_dataset):
            # if i == 10:
            #     break
            try:
                D, H, W = image.shape
                pid = self.test_dataset.sample_bboxes[i].get_field("filename")
                pid = pid.split('/')[-2].replace('.nii.gz', '')

                print('[%d] Predicting %s' % (i, pid), image.shape)

                with torch.no_grad():
                    input = input.cuda().unsqueeze(0)
                    if self.cfg.testing['use_rcnn']:
                        rpns, detections, ensembles = self.model.forward(input, [])
                    else:
                        rpns = self.model.forward(input, [])
                        # image = np.expand_dims(image, 0)
                        # image = np.expand_dims(image, 0)
                        # image = torch.from_numpy(image).cuda()
                        # rpns = torchscript_model(image.half())

                rpns = rpns.cpu().numpy()
                if len(rpns):
                    rpns = rpns[:, 1:]
                    np.save(os.path.join(res_dir, '%s_rpns.npy' % (pid)), rpns)

                if self.cfg.testing['use_rcnn']:
                    detections = detections.cpu().numpy()
                    ensembles = ensembles.cpu().numpy()

                    if len(detections):
                        detections = detections[:, 0:-1]
                        np.save(os.path.join(res_dir, '%s_rcnns.npy' % (pid)), detections)

                    if len(ensembles):
                        ensembles = ensembles[:, 1:]
                        np.save(os.path.join(res_dir, '%s_ensembles.npy' % (pid)), ensembles)

                # Clear gpu memory
                torch.cuda.empty_cache()

            except Exception as e:
                torch.cuda.empty_cache()
                traceback.print_exc()

        self.npy2csv('rpns', res_dir)

        if self.cfg.testing['use_rcnn']:
            self.npy2csv('rcnns', res_dir)
            self.npy2csv('ensembles', res_dir)

    def deploy(self):
        self.model.eval()
        initial_checkpoint = self.cfg.testing['weight']
        print('[Loading model from %s]' % initial_checkpoint)
        checkpoint = torch.load(initial_checkpoint)
        self.model.load_state_dict(checkpoint['state_dict'])
        set_requires_grad(self.model)

        self.test_dataset = self.DataReader(self.cfg, mode='test')
        input_data = self.test_dataset.__getitem__(0)[1]
        input_data = np.expand_dims(input_data, 0)
        input_data = np.expand_dims(input_data, 0)
        print("input_data shape", input_data.shape)
        input_data = torch.from_numpy(input_data[:, :, :228, :226, :234]).cuda()
        # print(input_data.shape)

        # pth to torchscript
        input_data = input_data.half()
        self.model = self.model.half()
        print('1')
        traced_model = torch.jit.trace(self.model, (input_data))
        print(traced_model.graph)

        model_path = self.cfg.deploy['torchscript_save_path']
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        traced_model.save(model_path)

    def torchscript_test(self):
        i = 0
        self.test_dataset = self.DataReader(self.cfg, mode='test')
        torch_input_data, torchscript_input = self.test_dataset.__getitem__(i)
        pid = self.test_dataset.sample_bboxes[i].get_field("filename")
        pid = pid.split('/')[-1].replace('.nii.gz', '')
        print("input_data shape", torchscript_input.shape)
        res_dir = self.cfg.deploy['torchscript_compare_res_dir']

        initial_checkpoint = self.cfg.testing['weight']
        print('[Loading model from %s]' % initial_checkpoint)
        checkpoint = torch.load(initial_checkpoint)
        self.model.load_state_dict(checkpoint['state_dict'])
        self.model.set_mode('test')
        self.model.use_rcnn = self.cfg.testing['use_rcnn']
        with torch.no_grad():
            input_data = torch_input_data.cuda().unsqueeze(0)
            torch_res = self.model.forward(input_data, [])
        torch_res = torch_res.cpu().numpy()
        torch_res = torch_res[:, 1:]
        np.save(os.path.join(res_dir, '%s_torch_res.npy' % (pid)), torch_res)
        self.npy2csv('torch_res', res_dir)

        input_data = np.expand_dims(torchscript_input, 0)
        input_data = np.expand_dims(input_data, 0)
        input_data = torch.from_numpy(input_data).cuda()
        model_path = self.cfg.deploy['torchscript_save_path']
        print('[Loading torchscript from %s]' % model_path)
        torchscript_model = torch.jit.load(model_path)
        torchscript_res = torchscript_model(input_data)
        torchscript_res = torchscript_res.cpu().numpy()
        torchscript_res = torchscript_res[:, 1:]
        np.save(os.path.join(res_dir, '%s_torchscript_res.npy' % (pid)), torchscript_res)
        self.npy2csv('torchscript_res', res_dir)

    def npy2csv(self, name, save_dir):
        bbox_border = self.cfg.data['data_process']['bbox_border']
        res = []
        for bbox in self.test_dataset.sample_bboxes:
            pid = bbox.get_field("filename").split('/')[-2].replace('.nii.gz', '')
            if os.path.exists(os.path.join(save_dir, '%s_' % (pid) + name + '.npy')):
                bboxs = np.load(os.path.join(save_dir, '%s_' % (pid) + name + '.npy'))
                bboxs[:, 4] = bboxs[:, 4] - bbox_border
                bboxs[:, 5] = bboxs[:, 5] - bbox_border
                bboxs[:, 6] = bboxs[:, 6] - bbox_border
                bboxs = bboxs[:, [3, 2, 1, 6, 5, 4, 0]]
                names = np.array([[pid]] * len(bboxs))
                res.append(np.concatenate([names, bboxs], axis=1))

        res = np.concatenate(res, axis=0)
        col_names = ['uid', 'x_px', 'y_px', 'z_px', 'diameter_x_px', 'diameter_y_px', 'diameter_z_px', 'probability']
        submission_path = os.path.join(save_dir, 'submission_' + name + '.csv')
        df = pd.DataFrame(res, columns=col_names)
        df.to_csv(submission_path, index=False)

    def lr_shedule(self, epoch, batch_idx):
        warmup = self.cfg.training['sheduler']['warmup']
        total = self.cfg.epochs
        if epoch < warmup:
            epoch += float(batch_idx + 1) / len(self.train_loader)
            lr_adj = 1. / self.cfg.world_size * (epoch * (self.cfg.world_size - 1) / warmup + 1)
        elif epoch <= total * 0.5:
            lr_adj = 1.
        elif epoch <= total * 0.8:
            lr_adj = 0.1
        else:
            lr_adj = 0.01
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.cfg.lr * self.cfg.world_size * lr_adj

    def num_neg_shedule(self, epoch):
        total = self.cfg.epochs
        if epoch <= total * 0.2:
            self.model.rpn.num_neg = 800
        elif epoch <= total * 0.5:
            self.model.rpn.num_neg = 1600
        elif epoch <= total * 0.8:
            self.model.rpn.num_neg = 2000
        else:
            self.model.rpn.num_neg = 3000

    def rand_crop_ratio_shedule(self, epoch):
        def _set_dataset_r_rand(fraction):
            self.train_loader.dataset.r_rand = fraction
            self.train_loader.sampler.update_dataset(self.train_loader.dataset)
            if self.cfg.world_rank == 0:
                print('[RAND_CROP_RATIO_SHEDULE: length of train loader %d, length of valid loader %d]' % (
                    len(self.train_loader), len(self.val_loader)))

        total = self.cfg.epochs
        if epoch <= total * 0.05:
            pass
        elif epoch <= total * 0.1:
            _set_dataset_r_rand(0.3)
        elif epoch <= total * 0.2:
            _set_dataset_r_rand(0.4)
        else:
            _set_dataset_r_rand(0.5)


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
            # for name, param in net.named_parameters():
            #     print(name,param.size())
            # raise
