# -*- coding:utf-8 -*-
import os
import sys
import glob
import math
import random
import numpy as np
import pandas as pd
import scipy.ndimage as nd
from scipy.ndimage import zoom
# from keras.utils import to_categorical
from torch.utils.data import Dataset
from collections import Counter, defaultdict

from .ImagePreprocess import *
from .ReadData import load_data,load_single_data
from .data_aug_helper import data_augmentation,random_add_val
import json

class SegDataGenerator(Dataset):
    def __init__(self,data_paths, **kwargs):
        super(SegDataGenerator,self).__init__()
        
        self.aug = kwargs.get('aug',False)
        self.HU_min = kwargs.get('HU_min',-1024)
        self.HU_max = kwargs.get('HU_max',1024)
        self.nonfix_crop = kwargs.get('nonfix_crop',False)
        self.stride_ratio = kwargs.get('stride_ratio',8)
        self.densityTypeIndex = kwargs.get('densityTypeIndex',-1)
        self.input_size = kwargs.get('input_size')
        self.diameter_enlarge_ratio = kwargs.get('diameter_enlarge_ratio',1.5)
        self.mode = kwargs.get('mode','training')
        self.direction_z_ratio = kwargs.get('direction_z_ratio',0.5)
        self.diameter_idx = kwargs.get('diameter_idx',None)
        self.spacing_idx = kwargs.get('spacing_idx',None)
        self.clip_choice = kwargs.get('clip_choice',True)
        self.target_shape = kwargs.get('target_shape',72)
        self.min_patch_size = kwargs.get('min_patch_size')
        self.do_cls = kwargs.get('do_cls',False)
        self.zoom_shape = kwargs.get('zoom_shape',None)
        self.use_label_diameter = kwargs.get('use_label_diameter',False)
        self.calculate_shape_choice = kwargs.get('calculate_shape_choice',1)
        repeat_times = kwargs.get('repeat_times',4)

        # self.data, self.mask, self.info = self.__load_data([data_paths,mask_paths,info_paths])
        # print('data shape is ',self.data.shape)
        # print('mask shape is ',self.mask.shape)
        # print('info shape is ',self.info.shape)

        # self.indices = np.arange(self.data.shape[0])
        # if self.mode=='training':
        #     self.indices = np.repeat(self.indices,repeats=repeat_times)
        # print ('length of indices is',len(self.indices))
        # random.shuffle(self.indices)
        with open(data_paths, 'r+') as f:
            self.data_list = json.load(f)['data']
        
        
    def __load_single_data(self,data_paths,aux_task=False,target_shape=-1):
        # data = 0
        # for path in data_paths:
        #     temp_data = np.load(path)
        #     if len(temp_data.shape)==4 and target_shape>0:
        #         center = int(temp_data.shape[1]/2)
        #         margin_size = int(target_shape/2)
        #         temp_data = temp_data[:,center-margin_size:center+margin_size,
        #                             center-margin_size:center+margin_size,
        #                             center-margin_size:center+margin_size]
        #     if type(data) == int:
        #         data = np.asarray(temp_data)
        #     else:
        #         data = np.concatenate([data,temp_data],axis=0)
        data = np.load(data_paths)
        return data
    
    
    def __load_data(self,paths_list,aux_task=False,target_shape=-1): 
        result_list = [self.__load_single_data(data_paths,aux_task=aux_task,target_shape=target_shape) for data_paths in paths_list]
        return result_list
        
        
    def __len__(self):
        # return len(self.indices)
        return len(self.data_list)
        
    def __getitem__(self,index):
        # if self.mode == 'test':
        data_idx = index
        # else:
        #     data_idx = self.indices[index]
        return self.getSingleCase(data_idx)
        
    def getSingleCase(self,data_idx):
        HU_min,HU_max = self.HU_min,self.HU_max
        # data,mask = self.data[data_idx],self.mask[data_idx]
        data_path = self.data_list[data_idx]
        mask_path = data_path.replace('_data.npy', '_mask.npy')
        data = np.load(data_path)
        mask = np.load(mask_path)
        data = (data - HU_min)/float(HU_max - HU_min)
        return data[None, ...], mask[None, ...]
        # if self.nonfix_crop:
        #     '''
        #     是否使用info中记录的长径(infer的时候可以选择是)
        #     '''
        #     if self.use_label_diameter:
        #         diameter_mm = float(info[self.diameter_idx])
        #         spac_x = float(info[self.spacing_idx])
        #         diameter_px = diameter_mm/spac_x
        #     else:
        #         diameter_px = None
        #     if self.calculate_shape_choice==1:
        #         patch_shape = CalculatePatchSizeWithDiameter(mask,stride_ratio=self.stride_ratio,
        #                                        diameter_enlarge_ratio = self.diameter_enlarge_ratio,
        #                                        direction_z_ratio = self.direction_z_ratio,
        #                                        min_patch_size = self.min_patch_size,
        #                                        diameter_px=diameter_px)
        #     elif self.calculate_shape_choice==2:
        #         '''
        #         线上模型
        #          '''
        #         patch_shape = CalculatePatchSizeWithDiameterV2(mask,stride_ratio=self.stride_ratio,
        #                                        diameter_enlarge_ratio = self.diameter_enlarge_ratio,
        #                                        direction_z_ratio = self.direction_z_ratio,
        #                                        min_patch_size = self.min_patch_size,
        #                                        diameter_px=diameter_px)
        #     else:
        #         patch_shape = CalculatePatchSizeWithDiameterV3(mask,stride_ratio=self.stride_ratio,
        #                                        diameter_enlarge_ratio = self.diameter_enlarge_ratio,
        #                                        direction_z_ratio = self.direction_z_ratio,
        #                                        min_patch_size = self.min_patch_size,
        #                                        diameter_px=diameter_px)
        # else:
        #     patch_shape = self.input_size
        
        # ################## This pair of vals should also be output 
        # final_image,final_mask = Data_preprocess(data,mask,aug=self.aug,input_shape=patch_shape,shift_ratio=1.0)

        # non_zoom_image,non_zoom_mask = final_image,final_mask 

        # ################################## zoom if necessary
        # if self.zoom_shape is not None:
        #     if type(self.zoom_shape)==int:
        #         ratio = float(self.zoom_shape)/np.squeeze(patch_shape[2])
        #         current_zoom_shape = [int(val*ratio) for val in patch_shape]
        #     else:
        #         current_zoom_shape = eval(self.zoom_shape)
        #     ori_spacing = [float(val) for val in info[self.spacing_idx:self.spacing_idx+3]]

        #     final_image,final_mask,new_spacing = imagepari_target_shape_respacing(final_image,final_mask,ori_spacing,current_zoom_shape)
        #     final_image = final_image[np.newaxis,...,np.newaxis]
        #     final_mask = final_mask[np.newaxis,...,np.newaxis]
        # else:
        #     new_spacing = [1 for _ in range(3)]


        # if self.aug:
        #     if np.random.choice([0,1]):
        #         final_image = np.swapaxes(final_image,2,3)
        #         final_mask = np.swapaxes(final_mask,2,3)
        #     if np.random.choice([0,1]):
        #         theta = np.random.uniform(-3, 3)
        #         cval = HU_min if HU_min else np.amin(final_image)
        #         final_image = nd.interpolation.rotate(final_image,theta,axes=(2,3),reshape=False,cval=cval)
        #         final_mask = nd.interpolation.rotate(final_mask,theta,axes=(2,3),reshape=False)
        #         final_mask = final_mask>0

        #     final_image, final_mask = data_augmentation(np.squeeze(final_image), np.squeeze(final_mask))
        # if self.clip_choice:
        #     final_image = np.clip(final_image,HU_min,HU_max)
        # final_image = (final_image - HU_min)/float(HU_max - HU_min)
        
        # if self.aug:
        #     '''
        #     随机加入一点噪声值
        #     '''
        #     final_image = random_add_val(final_image)
        # final_image = np.squeeze(final_image)[np.newaxis,...]
        # final_mask = np.squeeze(final_mask)[np.newaxis,...]
        # mask_size = final_mask.shape[0]*final_mask.shape[1]*final_mask.shape[2]*final_mask.shape[3]
        # mask_region_size = np.sum(final_mask)  
        # final_mask = final_mask.astype(np.uint8)
        
        
        # if self.mode == 'test':
        #     return final_image,final_mask,non_zoom_image,non_zoom_mask,patch_shape,self._GenerateMargin(data_idx),new_spacing,info
        # else:
        #     return final_image,final_mask
    
    def CalculateDiameterPx(self,info):
        if info is None:
            return None
        if self.diameter_idx is not None and self.spacing_idx is not None and self.mode=='test':
            spacing = [float(val) for val in info[self.spacing_idx:self.spacing_idx+3]]
            if spacing[0]==spacing[1]:
                spac_x = spacing[0]
            else:
                spac_x = spacing[-1]
            diameter_mm = float(info[self.diameter_idx])
            diameter_px = int(round(diameter_mm/spac_x))
        else:
            diameter_px = None
        return diameter_px
    
    def _GenerateMargin(self,data_idx):
        data,mask = self.data[data_idx],self.mask[data_idx]
        info = self.info[data_idx]
        patch_shape = CalculatePatchSizeWithDiameter(mask,stride_ratio=self.stride_ratio,
                                           diameter_enlarge_ratio = self.diameter_enlarge_ratio,
                                           direction_z_ratio = self.direction_z_ratio,
                                           diameter_px = self.CalculateDiameterPx(info),
                                            aug=self.aug)
        data_shape = data.shape
        margins = []
        for idx in range(3):
            val = max(0,data_shape[idx]-patch_shape[idx])/2
            margins.append(val)
        return margins

    def _zoomPredBack(self,info,pred,new_spacing,patch_shape):
        ori_spacing = [float(val) for val in info[self.spacing_idx:self.spacing_idx+3]]
        _,current_new_pred,_ = imagepari_target_shape_respacing(np.squeeze(pred),np.squeeze(pred),new_spacing,[int(val) for val in patch_shape])

        return current_new_pred

    
    def DataStat(self,diameter_ths):
        diameter_index = self.diameter_idx
        densityTypeIndex = self.densityTypeIndex
        infos = self.info
        uids = set([x[0] for x in self.info])
        print ('Number of records in total is %d in %d series'%(len(self.info),len(uids)))
        density_map_dict = {
            1:'solid',
            2:'mix',
            3:'ggo',
            4:'calc',
            5:'solid',
            6:'clac',
            7:'p_patch'
        }

        keys = ['solid','mix','ggo','calc','p_patch']
        for diam_th_vals in diameter_ths:
            print ('diam_th_vals is',diam_th_vals)
            diameter_min,diameter_max = diam_th_vals
            density_stat_map = defaultdict(int)
            count = 0
            for record in infos:
             
                diameter = float(record[diameter_index])
#                 print ('diameter is',diameter)
                densityType = int(float(record[densityTypeIndex]))
                if densityType not in density_map_dict.keys():
                    continue
                if diameter>=diameter_min and diameter<diameter_max:
                    density_stat_map[density_map_dict[densityType]]+=1
                    count += 1
            print ('='*60)
            print ('%d records between [%d,%d] is '%(count,diameter_min,diameter_max))
            for key in keys:
                print ('%s:%d'%(key,density_stat_map[key]))

    def _transImage(self,img,label):
        # final_image,final_mask,new_spacing = imagepari_target_shape_respacing(final_image,final_mask,ori_spacing,current_zoom_shape)
        img = img.numpy()
        label = label.numpy()

        trans_list = {}
        ########### flip
        flip_list = []
        cnt = 3
        prob = 0.5
        while random.random() < prob and cnt > 0:
            degree = random.choice([0, 1, 2])
            img = np.flip(img, axis=degree)
            label = np.flip(label, axis=degree)
            cnt = cnt - 1
            flip_list.append(degree)
        trans_list['flip_list'] = flip_list

        ########## swap axis
        if np.random.choice([0,1]):
            img = np.swapaxes(img,2,3)
            label = np.swapaxes(label,2,3)
            trans_list['swap'] = True
        else:
            trans_list['swap'] = False

        ############## zoom
        '''
        todo: add this part
        '''

        copy_img = img.copy()
        copy_label = label.copy()
        return copy_img,copy_label,trans_list



    def _transWithFunc(self,imgs,func_list):
        outputs = []
        if func_list['swap']:
            for img in imgs:
                new_img = np.swapaxes(img,1,2)
                outputs.append(new_img)
        else:
            outputs = [val for val in imgs]
        for degree in func_list['flip_list']:
            for img_idx,img in enumerate(outputs):
                new_img = np.flip(img,axis=degree)
                outputs[img_idx] = new_img
        outputs = [val.copy() for val in outputs]
        # print ('length of outputs',outputs[0].shape)
        return outputs