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 pltF
import imageio
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader
# from tensorboard_logger import Logger
from k8s_utils import CMetricsWriter
# sys.path.append('../../lungNoduleDensityClassification/nodulecls/utils')
# from RWconfig import LoadJson,WriteJson
# from OfflineEval import *


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

def LoadJson(path):
    with open(path, 'r+') as f:
        return json.load(f)

def Dice(label, pred):
    x = label.sum()
    y = pred.sum()
    dice = 2*(label*pred).sum()/(x+y+1e-5)
    return dice

def norm(image, hu_min=-1000.0, hu_max=600.0):
    image = (np.clip(image.astype(np.float32), hu_min, hu_max) - hu_min) / float(hu_max - hu_min)
    return image * 255.

class Segmentation3D(object):
    def __init__(self,cfg,logger):
        # self.filename = filename
        self.cfg_g = cfg
        self.logger = logger
        self.mode = cfg.mode
        self.parseParams(cfg)
        if self.mode == 'training':
            self.device = torch.device("cuda:%s"%(cfg.rank) if torch.cuda.is_available() else "cpu")
            self.metrics_writer = CMetricsWriter(
                filename=self.common_cfg['train_metrics'],
                headernames=['epoch', 'lr', 'loss', 'performanceAccuracy']
            )
            self.cfg = cfg.training
        else:
            self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
            self.cfg = cfg.testing
        
    def __call__(self):
        if self.mode == 'infer':
            self.Infer()

            # self._makeTensorBoard()
        else:
            self.buildModel()
    
            self.compileModel()
            self.loadData()
            
            if self.mode == 'training':
                self.model = self.model.cuda()
                # torch.nn.parallel.DistributedDataParallel(self.model, 
                #                                           broadcast_buffers=False, 
                #                                           find_unused_parameters=True)
                self.Train()
            elif self.mode == 'testing':
                self.Test()
        
       
    # 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):
        # cfg = LoadJson(self.filename)
        mode = 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.save_path = os.path.join(os.path.pardir(self.pre_model_file), 'run')
        self.save_path = self.cfg_g.common['save_path']
        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 = [64, 64, 64]
        generator_parameters = {'aug':self.cfg.get('aug',False),
                         'HU_min':self.common_cfg['normalizer'][0],'HU_max':self.common_cfg['normalizer'][1],
                         'nonfix_crop':self.cfg['nonfix_crop'],'stride_ratio':self.common_cfg['stride_ratio'],
                         'densityTypeIndex':self.common_cfg['densityTypeIndex'],
                         'diameter_enlarge_ratio':self.common_cfg['diameter_enlarge_ratio'],
                         'diameter_idx':self.common_cfg['diameter_idx'],'spacing_idx':self.common_cfg['spacing_idx'],
                          'input_size':input_size,'mode':self.mode,
                          'min_patch_size':self.common_cfg.get('min_patch_size',16),
                           'do_cls':self.common_cfg.get('do_cls',False),
                           'repeat_times':self.cfg.get('repeat_times',4),
                           'zoom_shape':self.common_cfg.get('zoom_shape',None),
                           'calculate_shape_choice':self.common_cfg.get('calculate_shape_choice',1),
                           'use_label_diameter':self.cfg.get('use_label_diameter',False)}

        # train_data_files = self.common_cfg["train_file"]
        # val_data_files = self.common_cfg["val_file"]
        # test_data_files = self.common_cfg["test_file"]
        if self.mode == 'training':
            # train_paths = [self.common_cfg['train_data_path'], self.common_cfg['train_mask_path'], self.common_cfg['train_info_path']]
            # val_paths =  [self.common_cfg['val_data_path'], self.common_cfg['val_mask_path'], self.common_cfg['val_info_path']]
            # test_paths = [self.common_cfg['val_data_path'], self.common_cfg['val_mask_path'], self.common_cfg['val_info_path']]
            train_paths = self.common_cfg['train_data_path']
            val_paths = self.common_cfg['val_data_path']

        elif self.mode == "testing":
            # val_paths =  [self.common_cfg['val_data_path'], self.common_cfg['val_mask_path'], self.common_cfg['val_info_path']]
            val_paths =  self.common_cfg['val_data_path']

        # self.test_paths = test_paths
        
        if self.mode=='training':
            self.train_generator = SegDataGenerator(data_paths=train_paths,**generator_parameters)
            generator_parameters['aug'] = False
            generator_parameters['mode'] = 'val'
            
            self.val_generator =  SegDataGenerator(data_paths=val_paths,**generator_parameters)
            
        else:
            self.val_generator = SegDataGenerator(data_paths=val_paths,**generator_parameters)
        
    
    # def _parseDataJosn(self,filename):
    #     data_files = []
    #     mask_files = []
    #     info_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:
    #         data_file = record["data_path"]
    #         mask_file = record["mask_path"]
    #         info_file = record.get("info_path",None)
    #         data_files.append(data_file)
    #         mask_files.append(mask_file)
    #         if info_file is not None:
    #             info_files.append(info_file)
    #     if len(info_files)==0:
    #         info_files = None
    #     return data_files,mask_files,info_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
        else:
            pretrain = self.cfg_g.pretrain_msg
            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
        if model_path is not None:
            try:
                print('load model weights')
                self.model.to(device=self.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_g.batch_size
        val_batch_size = self.cfg_g.batch_size
        
        print ('train_batch_size is ',train_batch_size,val_batch_size)
        
        train_Loader = DataLoader(self.train_generator,batch_size=train_batch_size,shuffle=True)
        val_Loader = DataLoader(self.val_generator,batch_size=val_batch_size,shuffle=False,drop_last=False)
        num_epoch = self.cfg_g.epochs
        compare_choice = self.cfg.get('compare_choice',False)
        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'
        }

        n_train = int(len(self.train_generator)/train_batch_size)
        label_dtype = torch.long
        
        self.loadModel()
        best_val_loss = 100000
        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=n_train,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 compare_choice:
                    loss, _ = self.TrainSingle(batch)
                    # else:
                    #     loss = self.TrainSingleWithConstrain(batch)
                    self.optimizer.zero_grad()
                    loss.backward()
                    epoch_loss += loss.item()
                    self.optimizer.step()
                    pbar.update()

            train_loss = (epoch_loss)/float(n_train)
            self.model.eval()
            val_loss, dice = 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 torch.distributed.get_rank() == 0:
                self.metrics_writer.append_one_line([epoch, self.optimizer.param_groups[0]['lr'], train_loss, dice])
                with open(self.common_cfg['eval_pjson'], 'w+') as file:
                    json.dump(
                        {
                            "loss": val_loss,
                            "dice": dice
                        }, file, indent=4)
                        
                # 更新eval中的performance.md
                with open(self.common_cfg['eval_pmd'], 'w+') as file:
                    file.write('# overall performance \n')
                    file.write('| loss | dice |\n')
                    file.write('| ---- | ---- | \n')
                    file.write(fr'| {val_loss} | {dice} |\n')


                # 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)

                if val_loss <= best_val_loss:
                    with open(self.common_cfg['train_result'], 'w+') as file:
                        json.dump(
                            {
                                "successFlag":"TRAINING", 
                                "bestModelEpoch": epoch
                            }, file, indent=4)
                    save_path = self.common_cfg['save_path']+'/model.pth'
                    torch.save(self.model.state_dict(), save_path)
                    print(f'model save to {save_path}')
                    best_val_loss = val_loss

        if torch.distributed.get_rank() == 0:
            json_info = json.load(open(self.common_cfg['train_result'], 'r'))
            json_info['successFlag'] = 'SUCCESS'
            with open(self.common_cfg['train_result'], 'w') as file:
                json.dump(json_info, file, indent=4)    
            print('train done')    
            
    def Test(self):
        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'
        }
        self.val_losses = [[] for _ in range(len(loss_names.keys()))]
        
        val_batch_size = 1
        val_Loader = DataLoader(self.val_generator,batch_size=val_batch_size,shuffle=False)
        self.loadModel()
        self.model.to(self.device)
        # self.model.cuda()
        val_loss, dice = self.Val(val_Loader,1,1)
        
        with open(self.common_cfg['eval_pjson'], 'w+') as file:
            json.dump(
                {
                    "loss": val_loss,
                    "dice": dice,
                }, file, indent=4)

        with open(self.common_cfg['eval_pmd'], 'w+') as file:
            file.write('# overall performance \n')
            file.write('| loss | dice | \n')
            file.write('| ---- | ---- | \n')
            file.write(fr'| {val_loss} | {dice} |\n')

        with open(self.common_cfg['eval_result'], 'w+') as file:
            json.dump(
                    {
                        "successFlag": "SUCCESS",
                    },file, indent=4)


    def Val(self,val_Loader,epoch,num_epoch):
        compare_choice = self.cfg.get('compare_choice',False)
        val_batch_size = self.cfg_g.batch_size

        self.model.eval()
        total_loss = 0
        total_dice = 0
        n_val = int(len(self.val_generator)/val_batch_size)
        print('n_val:', n_val)
        with tqdm(total=n_val,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 compare_choice:
                loss, dice = self.TrainSingle(batch,mode='val')
                # else:
                #     loss = self.TrainSingleWithConstrain(batch,mode='val')
                
                total_loss += loss.item()
                total_dice += dice.item()
                pbar.update()
        self.model.train()
        return total_loss/float(n_val), total_dice/float(n_val)

    # 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 Infer(self):
        import cv2
        from skimage import measure, morphology
        from scipy.spatial import distance
        import SimpleITK as sitk
        data = self.common_cfg['val_data_path']
        info = self.common_cfg['val_info_path']
        data = np.load(data)
        info = np.load(info)
        print(f'data shape: {data.shape}')
        self.buildModel()
        self.loadModel()
        num_imgs = data.shape[0]
        self.model = self.model.to(self.device)
        self.model.eval()
        all_images = []
        target_path = self.cfg_g.save_img_dir
        save_json = os.path.dirname(target_path)+'/diameter.json'
        diameter_dict = {}
        for i in range(num_imgs):
            print(f'img: {i}')
            spacing_x, spacing_y, spacing_z = info[i]
            input_data = torch.from_numpy(data[i][None, None, ...]).float().to(self.device)
            pred = self.model(input_data)
            print(f'input shape: {input_data.shape}, output shapeL {pred[0].shape}')
            pred = torch.sigmoid(pred[0][0, 0, ...]).detach().cpu().numpy()
            mask = np.where(pred>0.9, 1, 0)
            # print(1)
            labeled_mask, num_features = measure.label(mask, connectivity=1, return_num=True)
            # print(2)
            regions = measure.regionprops(labeled_mask)
            # print(3)
            max_region = max(regions, key=lambda r:r.area)
            # print(4)
            coords =max_region.coords
            # print(coords.shape)
            if coords.shape[0]>10000:
                coords = coords[:10000, :]
            diameter = max(distance.pdist(coords)*spacing_x)
            # print(5)
            # diameter = 11
            print(f'max diameter: {diameter}')
            diameter_dict[i] = diameter
            imgs = input_data[0,0,...].cpu().numpy()
            imgs = norm(imgs)
            z = imgs.shape[0]
            for j in range(z):
                img = imgs[j]
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
                mask_i = mask[j][...,None]
                zero_mask = np.zeros_like(mask_i)
                mask_i = np.concatenate((zero_mask, zero_mask, mask_i*255), -1).astype(np.uint8)
                img = 0.5*img + 0.5*mask_i
            #     all_images.append(img)
            # for idx, img in enumerate(all_images):
                save_path = os.path.join(target_path, str(i)+'_'+str(diameter))
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                # print(img)
                cv2.imwrite(os.path.join(save_path, '{:03d}.png'.format(j)), img)
            
        print(f'save diameter json')
        with open(save_json, 'w+') as f:
            json.dump(diameter_dict, f)
        output_gif_path = self.cfg_g.save_gif_dir
        if not os.path.exists(output_gif_path):
            os.makedirs(output_gif_path)

        for file in os.listdir(target_path):
            files = []
            for f in os.listdir(os.path.join(target_path, file)):
                files.append(f)
            files.sort(key=lambda x: x[:-4])
            frames = []
            for i in tqdm(range(len(files))):
                im = imageio.imread(os.path.join(target_path, file) + '/' + files[i])
                frames.append(im)
            fps = 24.0
            imageio.mimsave(os.path.join(output_gif_path, file+'.gif'), frames, 'GIF', duration=1/fps)

    # 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.MSELoss() 

    #     data,label = batch[0]
    #     cls_label = 1


    #     '''
    #     trans img and 
    #     '''

    #     trans_imgs,trans_labels,trans_dicts = [],[],[]
    #     for inner_batch_idx in range(np.squeeze(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,...]
    #         if len(distance_map_label_trans.shape)==4:
    #             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 = []

    #     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])
    #         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))

    #     if cls_choice:
    #         loss_input = [cls_pred]+[cls_label]+[label_diff]
    #         cls_loss_val = self.cls_criterion(loss_input)
    #         cls_loss_val = self.common_cfg.get('cls_loss_ratio',0.1)*cls_loss_val
    #         total_loss += cls_loss_val

    #         if mode=='val':
    #             self.val_losses[-1].append(cls_loss_val.item())
    #         else:
    #             self.train_losses[-1].append(cls_loss_val.item()) 

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

        data,label = batch[0]
        cls_label = 1
        surface_criterion = SurfaceLoss()
        print(f'input data shape: {data.shape}, mask shape: {label.shape}')
        # 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)
        
        data = data.cuda().to(device=self.device,dtype=torch.float32)
        label = label.cuda().to(device=self.device,dtype=torch.float32)
        
        
        if mode=='val':
            with torch.no_grad():
                pred = self.model(data)
        else:
            pred = self.model(data)
        print(11111, pred[0].shape)
        if cls_choice:
            cls_pred = pred[-1]
            pred = pred[:-1]

        # pred = torch.sigmoid(pred[0][0, 0, ...]).detach().cpu().numpy()
        # mask = np.where(pred>0.8, 1, 0)
        total_dice = 0
        total_loss = 0
        for idx in range(len(pred)):
            loss_val = self.criterion(pred[idx],label)
            
            if idx==0:
                total_loss = loss_val*loss_weights[idx]
            else:
                total_loss += loss_val*loss_weights[idx]

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

            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_dice += Dice(label, pred[idx])
        total_loss /= (len(pred))
        total_dice /= (len(pred))

        if cls_choice:
            loss_input = [cls_pred]+[cls_label]+[label_diff]
            cls_loss_val = self.cls_criterion(loss_input)
            cls_loss_val = self.common_cfg.get('cls_loss_ratio',0.1)*cls_loss_val
            total_loss += cls_loss_val

            if mode=='val':
                self.val_losses[-1].append(cls_loss_val.item())
            else:
                self.train_losses[-1].append(cls_loss_val.item()) 

        return total_loss, total_dice
        
    def compileModel(self):
        lr_schedule = 'basic'
        self.optimizer = optim.Adam(self.model.parameters(),
            lr=self.cfg_g.lr)

        
        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.criterion =  SeLossWithBCEFocalLoss()
        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('TorchNoduleSeg/files/'+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 = 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)
#         summary(self.model.cuda(),(1,32,32,32))
        

    def AnayInferResult(self,dice_list,info_list,model_path=None,file_order=0,diameter_ths = [[0,4],[4,10],[10,30],[30,100000],[4,1000000]]):
        if model_path is None:
            model_path = self.pre_model_file
        
        inference_result_dict = {
            'model_weights':model_path,
            'test_data':self.common_cfg["test_file"],
            'result':{}
        }
        inference_result_dict['result']['overall_infer_dice'] = [np.mean(dice_list),np.std(dice_list)]
        print ('length of info list and dice list is',len(info_list),len(dice_list))
        print ('length of info is',len(info_list[0]))
        
        print ('info_list[0]',info_list[0],info_list[0][self.common_cfg['diameter_idx']],type(info_list[0][self.common_cfg['diameter_idx']]))
        
        self.logger.info('overall_infer_dice is %s'%(str([np.mean(dice_list),np.std(dice_list)])))
        for diameter_th in diameter_ths:
            diameter_min,diameter_max = diameter_th
            indices = [idx for idx in range(len(dice_list)) if float(info_list[idx][self.common_cfg['diameter_idx']])>=diameter_min and float(info_list[idx][self.common_cfg['diameter_idx']])<diameter_max]
            dice_vals = [dice_list[sample_idx] for sample_idx in indices]
#             diff_vals = [diam_diff_list[sample_idx] for sample_idx in indices]
            
            
            inference_result_dict['result']['Diameter_[%d,%d]_infer_dice'%(diameter_min,diameter_max)] = [np.mean(dice_vals),np.std(dice_vals)]
            self.logger.info('Diameter_[%d,%d]_infer_dice is %s'%(diameter_min,diameter_max,str([np.mean(dice_vals),np.std(dice_vals)])))
            

        WriteJson(os.path.join(os.path.dirname(self.pre_model_file),'inference_result_%02d.json'%file_order),inference_result_dict)
        print ('After writing json file to %s'%os.path.join(os.path.dirname(self.pre_model_file),'inference_result.json'))

    def MakeSinglePred(self,input_list):
        
        th = 0.05

        current_data,current_mask,non_zoom_image,non_zoom_mask,patch_shape,margins,new_spacing,info = input_list


        data = torch.from_numpy(current_data[np.newaxis,...]).to(device=self.device,dtype=torch.float32)
        pred = self.model(data)
        pred = [torch.sigmoid(val) for val in pred]
        pred = [val.data.cpu().numpy() for val in pred]
    

        th_pred = [val>th for val in pred]
        th_pred = [np.array(val).astype(np.uint8) for val in th_pred]


        if self.common_cfg.get('zoom_shape',None) is None:
            return data,pred,current_mask,margins
        else:
            zoom_back_preds = [self.test_generator._zoomPredBack(info,single_mask,new_spacing,patch_shape) for single_mask in th_pred]
            return non_zoom_image,zoom_back_preds,non_zoom_mask,margins

    
    def MakeInference(self):
        
        ################# load predtrain weights
        pretrain = self.cfg['pre_model_file']
        if '*' in pretrain:
            cand_paths = sorted(glob.glob(pretrain+'*pth'))
        else:
            cand_paths = [pretrain]
        model_dice_lists = [] 
        model_order = 0
        for model_path in cand_paths: 
            model_order+=1
            self.loadModel(model_path)
            infer_idx = self.cfg.get('infer_idx',-1)
            save_result_choice = 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 = [],[],[]
            info_list = []
            margin_list = []

            for idx in indices:
                current_data,pred,current_mask,margins = self.MakeSinglePred(self.test_generator[idx])
                image_list.append(current_data)
                mask_list.append(current_mask)
                pred_list.append(pred)
                info_list.append(self.test_generator.info[idx])
                margin_list.append(margins)

                dice_val = dice_coefficient_post(current_mask,pred[infer_idx])
                dice_vals.append(dice_val)

            ########################### Analy
            self.AnayInferResult(dice_vals,info_list,model_path,model_order)
            model_dice_lists.append(np.mean(dice_vals))

            self.image_list = image_list
            self.mask_list = mask_list
            self.pred_list = pred_list
            self.dice_vals = dice_vals
            self.info_list = info_list
            if save_result_choice:
                self.savePredResult(pred_list,margin_list,infer_idx)
        
        if len(model_dice_lists)>0:
            max_idx = np.argmax(model_dice_lists)
            max_dice = max(model_dice_lists)
            max_dice_model_path = cand_paths[max_idx]
            self.logger.info('model_path %s has max avg dice val %.4f'%(max_dice_model_path,max_dice))
        
        if self.cfg.get('badcase_vis',False):
            result_list = [image_list,mask_list,pred_list,dice_vals,info_list]
            diameter_ths = eval(self.cfg['diameter_ths'])
            dice_range = eval(self.cfg['dice_range'])
            self.BadCaseVis(result_list,diameter_ths,dice_range)
        
            
    def savePredResult(self,result_list,margin_list,infer_idx):
        final_result_list = []
        for idx in range(len(result_list)):
            pad_result = self.padImageBack(result_list[idx][infer_idx],margin_list[idx])
            final_result_list.append(pad_result)
        save_path = os.path.join(os.path.dirname(self.pre_model_file),'prediction.npy')
        np.save(save_path,final_result_list)
        logger.info('After saving pred result to %s'%save_path)
        
        
   
    def padImageBack(self,pred,margin):
        current_shape = pred.shape
        target_shape = [val+2*val_m for val,val_m in zip(current_shape,margin)]
        
        target_mask = np.zeros(target_shape)
        target_mask[margin[0]:margin[0]+current_shape[0],
                margin[1]:margin[1]+current_shape[1],
                margin[2]:margin[2]+current_shape[2]] = pred
        return target_mask
    
    
    def BadCaseVis(self,result_list=None,diameter_ths=[0,1000000],dice_range=[0,1]):
        if result_list is None:
            result_list = [self.image_list,self.mask_list,self.pred_list,self.dice_vals,self.info_list]

        diameter_min,diameter_max = diameter_ths
        data,mask,pred,dices,infos = result_list
        indices = [idx for idx in range(len(infos)) if float(infos[idx][self.common_cfg['diameter_idx']])>=diameter_min and float(infos[idx][self.common_cfg['diameter_idx']])<diameter_max]
        indices = [idx for idx in indices if dices[idx]<dice_range[1] and dices[idx]>=dice_range[0]]
        
        for idx in indices:
            print ('='*60)
            current_data,current_mask,current_pred = np.squeeze(data[idx]),np.squeeze(mask[idx]),np.squeeze(pred[idx])
            center_slice_idx = int(current_data.shape[0]/2)
            _,axs =plt.subplots(1,3)
            axs[0].imshow(current_data[center_slice_idx],cmap='bone')
            axs[1].imshow(current_mask[center_slice_idx],cmap='bone')
            axs[2].imshow(current_pred[center_slice_idx],cmap='bone')
            plt.show()
        
    
    def _DataStat(self,diameter_ths=[[0,1000000000],[0,4],[4,10],[10,20],[20,30],[30,1000000],[40,100000],[4,100000000000000000]]):
        if self.mode=='training':
            self.logger.info('Nodule distribution on diameter in train dataset is')
            self.train_generator.DataStat(diameter_ths)
            self.logger.info('Nodule distribution on diameter in val dataset is')
            self.val_generator.DataStat(diameter_ths)
        else:
            self.test_generator.DataStat(diameter_ths)
            
