import os
import sys
import six
import json
import glob
import shutil
import logging
import datetime
import numpy as np
import scipy.ndimage as nd
from tqdm import tqdm
from matplotlib import pyplot as plt

import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader
from tensorboard_logger import Logger

sys.path.append('../../lungNoduleDensityClassification/nodulecls/utils')
from RWconfig import LoadJson,WriteJson
from OfflineEval import *


sys.path.append('utils')
sys.path.append('model')
from modelBuild import build_model
from baseLayers.BasicModules import weight_init
from SegmentationLoss import Tversky_loss,DiceLoss,SurfaceLoss,DiceLoss_TwoLabels
from SegGenCOVIDV2 import SegGenCOVID
from ImagePreprocess import *
from ResultEval import dice_coefficient_post
from respacing_func import *
from cls_loss import *




class Segmentation3D(object):
    def __init__(self,filename,mode,logger):
        self.filename = filename
        self.logger = logger
        self.mode = mode
        self.parseParams()
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        
        
    def __call__(self):

        self._makeTensorBoard()
        
        self.buildModel()
   
        self.compileModel()
        self.loadData()
        
        if self.mode == 'training':
            self.Train()
        else:
            self.MakeInference()
       
    
    def _makeTensorBoard(self):
        tensorboard_path = "%s/tensorboard_logs"%self.save_path
        self.tensorboard_logger = Logger(logdir=tensorboard_path,flush_secs=10)
        self.logger.info('tensorboard path is %s'%tensorboard_path)


    def parseParams(self):
        cfg = LoadJson(self.filename)
        self.cfg = cfg[self.mode]
        self.common_cfg = cfg['common']
        
        
        if self.mode=='training':
            current_DT = datetime.datetime.now()
            self.current_day = '%02d%02d'%(current_DT.month,current_DT.day)
            model_order = '%02d'%self.cfg['model_order']
            self.save_path = self.common_cfg['base_path']+'/'+self.current_day+'_'+self.cfg['model_name']+'/'+str(model_order)+'_'+self.cfg['prefix']
        else:
            self.pre_model_file = self.cfg['pre_model_file']
            self.save_path = '/'.join(self.cfg['pre_model_file'].split('/')[:-1])
        
        self.logger.debug('save path is %s'%self.save_path)
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)
        target_filename = os.path.join(self.save_path,self.filename.split('/')[-1])
        shutil.copy(self.filename,target_filename)
         
        
    
    def loadData(self):

        input_size = [self.cfg['num_slice'],self.cfg['voxel_size'],self.cfg['voxel_size']]
        generator_parameters = {'aug':self.cfg.get('aug',False),
                         'HU_min':self.common_cfg['normalizer'][0],'HU_max':self.common_cfg['normalizer'][1],
                          'input_size':input_size,'mode':self.mode,
                           'repeat_times':self.cfg.get('repeat_times',4),
                           'estimate_patch_per_img':self.common_cfg.get('estimate_patch_per_img',12),
                           'area_th_val':self.cfg.get('area_th_val',0),
                           'use_npy':self.common_cfg.get('use_npy',False),
                           'val_mode':self.cfg.get('val_mode',1),
                           'infer_stride':self.cfg.get('infer_stride',1),
                           'learn_feature':self.cfg.get('learn_feature',False),
                           'bbox_margin':self.cfg.get('bbox_margin',0),
                           'norm_func':self.cfg.get('norm_func','HU'),
                           'center_region_choice':self.cfg.get('center_region_choice',False),
                           'patch_pos_ratio_range':eval(self.cfg.get('patch_pos_ratio_range',"[0,1]")),
                           'cache_size':self.cfg.get('cache_size',100),
                            'logger':self.logger}

        train_data_files = self.common_cfg["train_file"]
        val_data_files = self.common_cfg["val_file"]
        test_data_files = self.common_cfg["test_file"]
        
        train_paths = self._parseDataJosn(train_data_files)
        val_paths =  self._parseDataJosn(val_data_files)
        test_paths = self._parseDataJosn(test_data_files)
        
        print ('train_paths is ',train_paths)
        print ('val_paths is',val_paths)
        print ('test_paths',test_paths)

        self.test_paths = test_paths

        if self.mode=='training':
            self.train_generator = SegGenCOVID(df_paths=train_paths,**generator_parameters)
            generator_parameters['aug'] = False
            generator_parameters['mode'] = 'val'
            print ('='*60)
            print ('length of train generator is',len(self.train_generator))
            
            self.val_generator =  SegGenCOVID(df_paths=val_paths,**generator_parameters)
            print ('length of val generator is',len(self.val_generator))
            
        else:
            self.test_generator = SegGenCOVID(df_paths=test_paths,**generator_parameters)
        
    
    def _parseDataJosn(self,filename):
        df_files = []
        target_filename = os.path.join(self.save_path,filename.split('/')[-1])
        filename = 'files/%s'%filename
        files = LoadJson(filename)
        shutil.copy(filename,target_filename)
        
        for record in files:
            df_file = record["df_path"]
            df_files.append(df_file)

        return df_files
    
    def _findLast(self,paths):
        length = len('CP_epoch')
        vals = [int(x.split('/')[-1][length:]) for x in paths]
        max_val = max(vals)
        idx = [idx for idx in range(len(vals)) if vals[idx]==max_val][0]
        return paths[idx]
    
    def loadModel(self,path=None):
        if path is not None:
            model_path = path
            pretrain = model_path
        else:
            pretrain = self.cfg['pre_model_file']
            self.pretrain = pretrain
            model_path = self.pretrain
        print ('path is %s'%path)
        print ('model path is %s'%model_path)
        if os.path.exists(model_path):
            model_path = model_path
        elif len(glob.glob(pretrain+'/*pth'))>0:
            paths = glob.glob(pretrain+'/*pth')
            model_path = self._findLast(paths)
        else:
            model_path = None
        print ('============================================================================== model_path is %s'%model_path)
        if model_path is not None:
            try:
                self.model.to(device=self.device)
                print ('===================================== put model to device')
                self.model.load_state_dict(torch.load(model_path,map_location=self.device))
            except:
                self.logger.error('Cannot load pretrain weight %s'%model_path)
    
    def Train(self):
        train_batch_size = self.cfg['train_batch_size'] 
        val_batch_size = self.cfg['val_batch_size'] 
        num_epoch = self.cfg['n_epochs']
        compare_choice = self.cfg.get('compare_choice',False)
        cls_choice = self.common_cfg.get('do_cls',False)
        flooding_val = self.cfg.get('flooding_val',0)
        loss_names = {
            0:'branch_1_dice_loss',
            1:'branch_2_dice_loss',
            2:'branch_3_dice_loss',
            3:'branch_1_surface_loss',
            4:'branch_2_surface_loss',
            5:'branch_3_surface_loss',
            6:'branch_1_trans_dice_loss',
            7:'branch_2_trans_dice_loss',
            8:'branch_3_trans_dice_loss',
            9:'branch_1_trans_surface_loss',
            10:'branch_2_trans_surface_loss',
            11:'branch_3_trans_surface_loss',
            12:'branch_1_sim_loss',
            13:'branch_2_sim_loss',
            14:'branch_3_sim_loss',
            15:'cls_loss'
        }
        
        
        '''
        dataloader
        '''
        train_Loader = DataLoader(self.train_generator,batch_size=self.cfg['train_batch_size'],shuffle=True)
        val_Loader = DataLoader(self.val_generator,batch_size=self.cfg['val_batch_size'],shuffle=True)
        


        n_train = len(self.train_generator)
        label_dtype = torch.long
        
        '''
        载入模型
        '''
        self.loadModel()
        
        for epoch in range(num_epoch):
            self.model.train()
            self.model.to(self.device)
            self.model.cuda()
            epoch_loss = 0
            self.train_losses = [[] for _ in range(len(loss_names.keys()))]
            self.val_losses = [[] for _ in range(len(loss_names.keys()))]
            with tqdm(total=int(len(self.train_generator)/self.cfg['train_batch_size']),desc=f'Epoch{epoch+1}/{num_epoch}',unit='img') as pbar:
                batch_count = 0
                for batch in zip(train_Loader):
                    batch_count +=1
                    if not cls_choice:
                        if not compare_choice:
                            loss = self.TrainSingle(batch)[0]
                        else:
                            loss = self.TrainSingleWithConstrain(batch)[0]
                    else:
                        loss = self.TrainSingleWithCls(batch,compare_choice,epoch,batch_count)
                    self.optimizer.zero_grad()
                    flood = abs(loss-flooding_val)+flooding_val
                    flood.backward()
                    epoch_loss += loss.item()

                    self.optimizer.step()
                    pbar.update()
            train_loss = (epoch_loss)/float(n_train)
            self.model.eval()
            val_loss = self.Val(val_Loader,epoch,num_epoch)
            self.model.train()
            self.scheduler.step(val_loss)
            
            self.logger.info("Train loss：{}".format(train_loss))
            self.logger.info("Validation loss：{}".format(val_loss))

            self.tensorboard_logger.log_value('train_loss', train_loss, epoch)
            self.tensorboard_logger.log_value('val_loss', val_loss, epoch)
            self._AddTensorData(self.train_losses,loss_names,'train',epoch)
            self._AddTensorData(self.val_losses,loss_names,'val',epoch)
            
            if self.cfg['save_cp']:
                model_save_path_current_epoch = self.save_path+f'/CP_epoch{epoch+1}_loss_%.2f.pth'%val_loss
                torch.save(self.model.state_dict(),model_save_path_current_epoch)
                self.logger.info(f'Checkpoint{epoch+1} saved to %s!'%model_save_path_current_epoch)
                
            
    
    def Val(self,val_Loader,epoch,num_epoch):
        compare_choice = self.cfg.get('compare_choice',False)
        cls_choice = self.common_cfg.get('do_cls',False)
        
        self.model.eval()
        total_loss = 0
        n_val = len(self.val_generator)
        with tqdm(total=int(len(self.val_generator)/self.cfg['val_batch_size']),desc=f'Epoch{epoch+1}/{num_epoch}',unit='img') as pbar: 
            batch_count = 0
            for batch in zip(val_Loader):
                batch_count += 1
                if not cls_choice:
                    if not compare_choice:
                        loss = self.TrainSingle(batch,mode='val')[0]
                    else:
                        loss = self.TrainSingleWithConstrain(batch,mode='val')[0]
                else:

                    loss = self.TrainSingleWithCls(batch,compare_choice,epoch,batch_count,mode='val')

                total_loss += loss.item()
                pbar.update()
        self.model.train()
        return total_loss/float(n_val)

    def _getBatchData(self,mode='training'):
        batch_size = self.cfg.get('train_batch_size')

        current_generator = self.train_generator if mode=='training' else self.val_generator
        indices = np.arange(len(current_generator))

        chosen_indices = np.random.choice(indices,size=batch_size,replace=False)
        imgs,masks,labels = [],[],[]
        for idx in chosen_indices:
            patch,mask,label = current_generator[idx]
            imgs.append(patch)
            masks.append(mask)
            labels.append(label)

        imgs = np.array(imgs)
        masks = np.array(masks)
        labels = np.array(labels)

        imgs = torch.from_numpy(imgs)
        masks = torch.from_numpy(masks)
        labels = torch.from_numpy(labels)

        batch = [imgs,masks,labels]
        return [batch]
    
    def _getBatchDataV2(self,mode='training'):
        batch_size = self.cfg.get('train_batch_size')

        current_generator = self.train_generator if mode=='training' else self.val_generator
        indices = np.arange(len(current_generator.cache))
        chosen_indices = np.random.choice(indices,size=batch_size,replace=False)
        
        samples = np.take(current_generator.cache,chosen_indices,axis=0)
        
        imgs = [sample[0] for sample in samples]
        masks = [sample[1] for sample in samples]
        labels =  [sample[2] for sample in samples]
        

        imgs = np.array(imgs)
        masks = np.array(masks)
        labels = np.array(labels)

        imgs = torch.from_numpy(imgs)
        masks = torch.from_numpy(masks)
        labels = torch.from_numpy(labels)

        batch = [imgs,masks,labels]
        return [batch]


    def TrainSingleWithCls(self,batch,compare_choice,epoch,batch_count,mode='training'):
        cache_size = self.cfg.get('cache_size',100)
        batch_size = self.cfg.get('train_batch_size')
        target_batch_count = int(cache_size/batch_size*0.8)
        
        '''
        执行分类任务 要对两个batch内的数据做
        1.正负样本分类
        2.相似性loss
        
        batchA:当前的batch
        batchB:从历史数据中抽取(维持一个cache)
        '''
        if epoch>0 or batch_count>=target_batch_count:
            batch_compare = self._getBatchDataV2(mode)
        else:
            batch_compare = self._getBatchData(mode)
        if compare_choice:
            loss0,cls_pred0 = self.TrainSingleWithConstrain(batch,mode)
            loss_compare,cls_pred_compare = self.TrainSingleWithConstrain(batch_compare,mode)
        else:
            loss0,cls_pred0 = self.TrainSingle(batch,mode)
            loss_compare,cls_pred_compare = self.TrainSingle(batch_compare,mode)
        cls_label0 = batch[0][-1]
        cls_label_compare = batch_compare[0][-1]
        cls_label0 = cls_label0.cuda().to(device=self.device,dtype=torch.float32)
        cls_label_compare = cls_label_compare.cuda().to(device=self.device,dtype=torch.float32)


        total_loss = loss0+loss_compare
        

        loss_input = [[cls_pred0,cls_pred_compare]]+[[cls_label0,cls_label_compare]]
        cls_loss_val = self.cls_criterion(loss_input)
        cls_loss_val = self.common_cfg.get('cls_loss_ratio',0.1)*cls_loss_val

        if mode=='training':
            self.train_losses[-1].append(cls_loss_val.item())
        else:
            self.val_losses[-1].append(cls_loss_val.item())
        total_loss += cls_loss_val
        
        return total_loss



    def TrainSingleWithConstrain(self,batch,mode='training'):
        loss_weights = eval(self.cfg.get('loss_weights','[1.0,1.0,1.0]'))
        calculate_surface_loss = self.common_cfg.get('surface_loss',False)
        surface_loss_ratio = self.common_cfg.get('e_value',0.1)
        compare_loss_ratio = self.common_cfg.get('comapre_loss_ratio',0.1)
        cls_choice = self.common_cfg.get('do_cls',False)

        surface_criterion = SurfaceLoss()
        compare_loss = torch.nn. SmoothL1Loss()

        data,label,cls_label = batch[0]

        '''
        将原图进行变换
        '''
        trans_imgs,trans_labels,trans_dicts = [],[],[]
        for inner_batch_idx in range(data.shape[0]):
            trans_img,trans_label,trans_dict = self.train_generator._transImage(data[inner_batch_idx],label[inner_batch_idx])


            trans_imgs.append(trans_img)
            trans_labels.append(trans_label)
            trans_dicts.append(trans_dict)

        '''
        计算surface loss
        '''
        if calculate_surface_loss:
            if label.shape[0]>1:
                distance_map_label = calc_dist_map_batch(np.squeeze(label.cpu().numpy()))
                distance_map_label_trans = calc_dist_map_batch(np.squeeze(trans_labels))
            else:
                distance_map_label = calc_dist_map(np.squeeze(label.cpu().numpy()))
                distance_map_label_trans = calc_dist_map(np.squeeze(trans_labels))
            if len(distance_map_label.shape)==4:
                distance_map_label = distance_map_label[:,np.newaxis,...]
                distance_map_label_trans = distance_map_label_trans[:,np.newaxis,...]
                
            distance_map_label = torch.from_numpy(distance_map_label).to(device=self.device,dtype=torch.float32)
            distance_map_label_trans = torch.from_numpy(distance_map_label_trans).to(device=self.device,dtype=torch.float32)

        '''
        数据格式转换
        '''
        data = data.cuda().to(device=self.device,dtype=torch.float32)
        label = label.cuda().to(device=self.device,dtype=torch.float32)

        trans_imgs = np.array(trans_imgs)
        trans_labels = np.array(trans_labels)
        trans_imgs = torch.tensor(trans_imgs).cuda().to(device=self.device,dtype=torch.float32)
        trans_labels = torch.tensor(trans_labels).cuda().to(device=self.device,dtype=torch.float32)

        if mode=='val':
            with torch.no_grad():
                pred = self.model(data)
                if cls_choice:
                    cls_pred = pred[-1]
                    pred = pred[:-1]
                else:
                    cls_pred = None
                pred_trans = self.model(trans_imgs)
        else:
            pred = self.model(data)
            if cls_choice:
                cls_pred = pred[-1]
                pred = pred[:-1]
            else:
                cls_pred = None
            pred_trans = self.model(trans_imgs)

        cpu_pred =  [val.data.cpu().numpy() for val in pred]
        cpu_pred_trans = []

        '''
        对pred结果进行变换
        '''
        for branch_idx in range(len(pred)):
            current_branch_result = []
            for inner_batch_idx in range(data.shape[0]):
                current_cpu_pred_trans = self.train_generator._transWithFunc(cpu_pred[branch_idx][inner_batch_idx],trans_dicts[inner_batch_idx])
                current_branch_result.append(np.array(current_cpu_pred_trans).astype(np.double))
            current_branch_result = np.array(current_branch_result)
            cpu_pred_trans.append(torch.tensor(current_branch_result).cuda().to(device=self.device,dtype=torch.float32))

        gpu_pred_trans = cpu_pred_trans

        for idx in range(len(pred)):
            loss_val = self.criterion(pred[idx],label)
            loss_val_trans = self.criterion(pred_trans[idx],trans_labels)

            if mode=='val':
                self.val_losses[idx].append(loss_val.item())
                self.val_losses[idx+len(pred)*2].append(loss_val.item())
            else:
                self.train_losses[idx].append(loss_val_trans.item())
                self.train_losses[idx+len(pred)*2].append(loss_val_trans.item())
                
            if idx==0:
                total_loss = loss_val*loss_weights[idx]+loss_val_trans*loss_weights[idx]
            else:
                total_loss += (loss_val*loss_weights[idx]+loss_val_trans*loss_weights[idx])

            if calculate_surface_loss:
                surface_loss = surface_criterion(pred[idx],distance_map_label) 
                surface_loss_trans  = surface_criterion(pred[idx],distance_map_label_trans) 

                total_loss += surface_loss_ratio*(surface_loss+surface_loss_trans)
                if mode=='val':
                    self.val_losses[idx+len(pred)].append(surface_loss.item()*surface_loss_ratio)
                    self.val_losses[idx+len(pred)*3].append(surface_loss_trans.item()*surface_loss_ratio)
                else:
                    self.train_losses[idx+len(pred)].append(surface_loss.item()*surface_loss_ratio)
                    self.train_losses[idx+len(pred)*3].append(surface_loss_trans.item()*surface_loss_ratio)



            trans_loss = compare_loss(gpu_pred_trans[idx],pred_trans[idx])
            if idx==len(pred)-1:
                total_loss += trans_loss * compare_loss_ratio
                if mode=='val':
                    self.val_losses[idx+len(pred)*4].append(trans_loss.item() * compare_loss_ratio)
                else:
                    self.train_losses[idx+len(pred)*4].append(trans_loss.item() * compare_loss_ratio) 

        total_loss /= (2*len(pred))

        return total_loss,cls_pred
        
                                        
    
    def TrainSingle(self,batch,mode='training'):
        loss_weights = eval(self.cfg.get('loss_weights','[1.0,1.0,1.0]'))
        val_mode = self.cfg.get('val_mode',1)
        cls_choice = self.common_cfg.get('do_cls',False)
        calculate_surface_loss = self.common_cfg.get('surface_loss',False)
        surface_loss_ratio = self.common_cfg.get('e_value',0.1)
    
        data,label,cls_label = batch[0][:3]
        surface_criterion = SurfaceLoss()
        
        if calculate_surface_loss:
            if label.shape[0]>1:
                distance_map_label = calc_dist_map_batch(np.squeeze(label.cpu().numpy()))
            else:
                distance_map_label = calc_dist_map(np.squeeze(label.cpu().numpy()))
            if len(distance_map_label.shape)==4:
                distance_map_label = distance_map_label[:,np.newaxis,...]
                
            distance_map_label = torch.from_numpy(distance_map_label).to(device=self.device,dtype=torch.float32)
        
        
        
        if mode=='val':
            with torch.no_grad():
                if val_mode==1:
                    data = data.cuda().to(device=self.device,dtype=torch.float32)
                    label = label.cuda().to(device=self.device,dtype=torch.float32)
                    pred = self.model(data)
                else:
                    label = torch.from_numpy(cp_label).cuda().to(device=self.device,dtype=torch.float32)
                    new_data = [single_data.unsqueeze(0).cuda().to(device=self.device,dtype=torch.float32) for single_data in data]
                    pred_list = [[] for _ in range(3)]
                    for single_data in new_data:
                        result = self.model(single_data)

                        for inner_idx in range(len(pred_list)):
                            pred_list[inner_idx].append(result[inner_idx])
                    pred = []
                    for inner_idx in range(len(pred_list)):
                        pred.append(torch.cat(pred_list[inner_idx],0))
                    

        else:
            pred = self.model(data)

        total_loss = 0
        
        if cls_choice:
            cls_pred = pred[-1]
            pred = pred[:-1]
        else:
            cls_pred = None

       
        
        for idx in range(len(pred)):
            loss_val = self.criterion(pred[idx],label)

            if mode=='val':
                self.val_losses[idx].append(loss_val.item())
            else:
                self.train_losses[idx].append(loss_val.item())

            if idx==0:
                total_loss = loss_val*loss_weights[idx]
            else:
                total_loss += loss_val*loss_weights[idx]

            if calculate_surface_loss:
                surface_loss = surface_criterion(pred[idx],distance_map_label) 
                total_loss += surface_loss_ratio*surface_loss

                if mode=='val':
                    self.val_losses[idx+len(pred)].append(surface_loss.item()*surface_loss_ratio)
                else:
                    self.train_losses[idx+len(pred)].append(surface_loss.item()*surface_loss_ratio)
            

        total_loss/=len(pred)

        return total_loss,cls_pred
        
    def compileModel(self):
        union_choice = self.cfg.get('union_choice',False)
        lr_schedule = self.cfg.get('lr_schedule','basic')
        self.learn_feature = self.cfg.get('learn_feature',False)
        self.optimizer = optim.Adam(self.model.parameters(),lr=self.common_cfg['initial_learning_rate'])

        if lr_schedule=='basic':
            self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,patience=self.common_cfg['learning_rate_patience'],
                                                factor=self.common_cfg['learning_rate_drop'],verbose=True)
        else:
            self.scheduler = optim.lr_scheduler.CyclicLR(self.optimizer,
                                    base_lr=self.common_cfg.get('base_lr',1e-5),
                                    max_lr=self.common_cfg.get('max_lr',1e-4),
                                    step_size_up=self.common_cfg.get('step_size_up',2000),
                                    cycle_momentum=False)


        if self.common_cfg.get('do_cls',False):
            self.cls_criterion = SeLossWithBCEFocalLoss()
        if union_choice:
            # self.criterion = DiceLoss_TwoLabels(alpha=self.common_cfg['T_loss_alpha'],beta=1-self.common_cfg['T_loss_alpha'])
            self.criterion = DiceLoss_TwoLabels(alpha=0.1,beta=1.0)
        else:
            self.criterion = Tversky_loss(alpha=self.common_cfg['T_loss_alpha'],beta=1-self.common_cfg['T_loss_alpha'])
        
    
    
    def buildModel(self):
        filename = self.common_cfg['model_parameters']
        target_model_cfg_name = os.path.join(self.save_path,filename.split('/')[-1])
        filename = 'files/%s'%filename
        model_cfg = LoadJson(filename)
        shutil.copy(filename,target_model_cfg_name)
        
        encoder_model_name = self.common_cfg['encoder']
        decoder_model_name = self.common_cfg['decoder']
        
        encoder_model_cfg = model_cfg[encoder_model_name]
        decoder_model_cfg = model_cfg[decoder_model_name]
        



        encoder_model_cfg['do_cls'] = False
        
        encoder_model = build_model(encoder_model_name,encoder_model_cfg)
        decoder_model_cfg['encoder_model'] = encoder_model
        decoder_model_cfg['do_cls'] =  self.common_cfg.get('do_cls',False)

        self.model = build_model(decoder_model_name,decoder_model_cfg)
        self.model.apply(weight_init)

    def MakeSinglePred(self,idx,model_order=0,model_path=None):

        infer_idx = self.cfg.get('infer_idx',-1)
        th = self.cfg.get('th',0.05)
        respacing_back = self.cfg.get('respacing_back',False)
        print ('='*60)
        
        '''
        读入数据
        '''
        output_img_patches,output_mask_patches,pad_mask,pad_vals,bounds,image,lungmask,mask_path = self.test_generator[idx]
        uid = mask_path.split('/')[-1][:-len('.nii.gz')]
        
        predictions = []
        
        '''
        逐patch进行infer
        '''
        for idx in range(len(output_img_patches)):
            data = torch.from_numpy(output_img_patches[idx][np.newaxis,np.newaxis,...]).to(device=self.device,dtype=torch.float32)
            pred = self.model(data)
            result_before_sigmoid = [val.data.cpu().numpy() for val in pred][infer_idx]
                      
            
            pred = [torch.sigmoid(val) for val in pred]
            pred = [val.data.cpu().numpy() for val in pred]
            pred = pred[infer_idx]
            predictions.append(pred)
        '''
        combine回原图大小
        '''     
        result,mask = self.test_generator._CombinePrediction(predictions,pad_mask,pad_vals)    
        result,mask = self.test_generator._padBackToOriSize(image,result,mask,pad_vals,bounds)
        
        '''
        二值化
        '''
        raw_result = np.array(result>th).astype(np.uint8)
        
        ''''
        后处理:去掉小的连通域(目前线上模型没有用到&dice变换不大)
        '''
        raw_result_post = self._PostProcess(raw_result)

        '''
        去掉肺外区域
        '''
        result = lungmask*result
        th_result = np.array(result>=th).astype(np.uint8)
        

        '''计算dice'''
        dice_val = dice_coefficient_post(mask,th_result)
        dice_raw = dice_coefficient_post(mask,raw_result)
        dice_post = dice_coefficient_post(mask,raw_result_post)
        self.post_dice.append(dice_post)
        
        
        '''
        如果必要,将mask respacing到原图大小
        '''
        if respacing_back:
            
            img_path = sorted(glob.glob('/fileser/CT_COVID/IMAGE/RAW/%s*'%uid))[0]
            lungmask_path =  sorted(glob.glob('/fileser/CT_COVID/LUNGMASK/RAW/%s*'%uid))[0]
            img_ori = sitk.GetArrayFromImage(sitk.ReadImage(img_path))
            lungmask_ori = sitk.GetArrayFromImage(sitk.ReadImage(lungmask_path))
            
            ori_spac_mask_itk = sitk.ReadImage(lungmask_path)
            ori_spac_mask = sitk.GetArrayFromImage(ori_spac_mask_itk)
            lungmask_ori = ori_spac_mask
            
            spac = ori_spac_mask_itk.GetSpacing()
            result_itk = sitk.GetImageFromArray(result)
            result_itk.SetSpacing((1,1,1))
            
            ori_spac_result_itk = NearestResample(result_itk,spac,img_ori.shape[::-1])
            ori_spac_result = sitk.GetArrayFromImage(ori_spac_result_itk)
            ori_spac_result = np.array(ori_spac_result>th).astype(np.uint8)
            ################################# 
            
            try:
                dice_val_ori = dice_coefficient_post(ori_spac_mask,ori_spac_result)
                
                '''
                计算原图dice并存储
                '''
                self.ori_spac_dice.append(dice_val_ori)

                base_save_path = self.cfg.get('base_save_path','')
                current_case_save_path = '%s/%s_dice%.2f.nii.gz'%(base_save_path,uid,dice_val_ori*100)
                itkWriter = sitk.ImageFileWriter()
                itkWriter.SetUseCompression(True)
                itkWriter.SetFileName(current_case_save_path)
                itkWriter.Execute(ori_spac_result_itk)

            except Exception as err:
                print ('err ',err)
                print ('no valid orisapcing result')
            

        else:                          
            '''
            存储infer结果
            '''
            self._WriteInferResult(uid,th_result,dice_val,model_order,model_path)

        self.no_lungmask_dice.append(dice_raw)
        print ('dice_val and dice_raw and dice_post is',dice_val,dice_raw,dice_post)
        
        '''
        badcase可视化
        '''
        if dice_val<self.cfg.get('iou_th',0):        
            for slice_idx in range(result.shape[0]):
                current_result_slice = np.array(result[slice_idx]>0.05).astype(np.uint8)
                if np.sum(current_result_slice)==0 and np.sum(mask[slice_idx])==0:
                    continue

                slice_dice_val = dice_coefficient_post(mask[slice_idx],current_result_slice)
                print ('slice_dice_val is',slice_dice_val)
                figsize=15
                _,axs = plt.subplots(2,2,figsize=(figsize,figsize))
                axs[1][0].imshow(current_result_slice,cmap='bone')
                axs[1][1].imshow(mask[slice_idx],cmap='bone')
                axs[0][0].imshow(image[slice_idx],cmap='bone')
                axs[0][1].imshow(lungmask[slice_idx],cmap='bone')
                plt.show()

        return image,mask,result

    def _WriteInferResult(self,uid,result,dice,model_order=0,model_path=None):
        base_save_path = self.cfg.get('base_save_path','')
        base_save_path = '%s/%02d/'%(base_save_path,model_order)
        if not os.path.exists(base_save_path):
            os.makedirs(base_save_path)
        if model_path is not None:
            model_name = model_path.split('/')[-1]
            target_model_path = '%s/%s'%(base_save_path,model_name)
            if not os.path.exists(target_model_path):
                shutil.copy(model_path,target_model_path)
        if os.path.exists(base_save_path):
            current_case_save_path = '%s/%s_dice%.2f.nii.gz'%(base_save_path,uid,dice*100)
            itkWriter = sitk.ImageFileWriter()
            itkWriter.SetUseCompression(True)
            itkWriter.SetFileName(current_case_save_path)
            itk_result = sitk.GetImageFromArray(result)
            itk_result.SetSpacing((1,1,1))
            
            itkWriter.Execute(itk_result)

    def _AddTensorData(self,losses,name_dict,mode,epoch):
        for idx in range(len(losses)):
            if len(losses[idx])>0:
                mean_val = np.mean(losses[idx])
                self.tensorboard_logger.log_value('%s_%s'%(mode,name_dict[idx]), mean_val, epoch)


    def _PostProcess(self,mask):
        region_th = 100
        from skimage.measure import label,regionprops
        final_mask = np.zeros_like(mask)
        label_mask = label(mask)
        for region in regionprops(label_mask):
            if region.area<region_th:
                continue
            final_mask += (label_mask==region.label)
        result = np.array(final_mask>0).astype(np.uint8)
        return np.array(final_mask>0).astype(np.uint8)

    def _featureMapVis(self,x,layer_idx):
        print ('#'*60)
        print ('Feature Map Vis')
        print ('#'*60)
        temp_layers = list(self.model.children())

        net = nn.Sequential(*list(self.model.children())[:layer_idx])
        net.eval()
        out = net(x)
        
        out_cpu = [val.data.cpu().numpy() for val in out]

        return out_cpu
    
    def _normFeatureMap(self,FeatureMap):
        FeatureMap = np.squeeze(np.array(FeatureMap))
        ChannelSums = [np.sum(val) for val in FeatureMap]
        totalSum = np.sum(ChannelSums)
        normSums = [val/float(totalSum) for val in ChannelSums]
        
        result = [val1*val2 for val1,val2 in zip(normSums,FeatureMap)]
        return result

    
    def MakeInference(self):
        th = 0.05
        pretrain = self.cfg['pre_model_file']
        if '*' in pretrain:
            cand_paths = sorted(glob.glob(pretrain+'*pth'))
        else:
            cand_paths = [pretrain]
        
        
        model_dice_vals = []
        for model_order,model_path in enumerate(cand_paths): 
            print ('#'*60)
            self.loadModel(model_path)
            save_result_chocie = self.cfg.get('save_result_choice',False)

            self.model.eval()
            indices = np.arange(len(self.test_generator))

            dice_vals = []
            image_list,mask_list,pred_list = [],[],[]
            self.ori_spac_dice = []
            self.no_lungmask_dice = []
            self.post_dice = []
            label_sums = []

            for idx in indices:
                image,mask,result = self.MakeSinglePred(idx,model_order,model_path)
                image_list.append(image)
                mask_list.append(mask)
                pred_list.append(result)
                result = np.array(result>=th).astype(np.uint8)
                mask = np.array(mask>0).astype(np.uint8)
                label_sums.append(np.sum(mask))
                
                print ('sum of mask is %d and sum of pred is %d'%(np.sum(mask),np.sum(result)))
                dice_val = dice_coefficient_post(mask,result)
                dice_vals.append(dice_val)

            self.logger.info("mean val of dice vals is %.4f"%(np.mean(dice_vals)))
            model_dice_vals.append(np.mean(dice_vals))
            if len(self.ori_spac_dice)>0:
                self.logger.info("mean val of orispacing dice vals is %.4f"%(np.mean(self.ori_spac_dice)))   
            self.logger.info("mean val of (without lung mask) dice vals is %.4f"%(np.mean(self.no_lungmask_dice)))
            self.logger.info("mean val of (postprocessing) dice vals is %.4f"%(np.mean(self.post_dice)))
            self.logger.info('model path is %s'%model_path)
        if len(model_dice_vals)>0:
            max_dice_idx = np.argmax(model_dice_vals)
            max_mean_dice = max(model_dice_vals)
            max_dice_model_path = cand_paths[max_dice_idx]
            self.logger.info('Max dice is %.4f and its model path is %s'%(max_mean_dice,max_dice_model_path))
            
        
        
    
        
            
    

  
            
            
