# -*- 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 .DataAna import GenerateRatios
from .data_aug import random_crop
from .data_aug import data_augmentation
from collections import Counter
from .ImagePro import ShowEdges,HFFilter
from skimage import  exposure,data
from matplotlib import pyplot as plt
from collections import defaultdict
import json

def to_categorical(y, num_classes=None, dtype='float32'):
    """
        Converts a class vector (integers) to binary class matrix.

        E.g. for use with categorical_crossentropy.

        # Arguments
            y: class vector to be converted into a matrix
                (integers from 0 to num_classes).
            num_classes: total number of classes.
            dtype: The data type expected by the input, as a string
                (`float32`, `float64`, `int32`...)

        # Returns
            A binary matrix representation of the input. The classes axis
            is placed last.
    """
    y = np.array(y, dtype='int')
    input_shape = y.shape
    if input_shape and input_shape[-1] == 1 and len(input_shape) > 1:
        input_shape = tuple(input_shape[:-1])
    y = y.ravel()
    if not num_classes:
        num_classes = np.max(y) + 1
    n = y.shape[0]
    categorical = np.zeros((n, num_classes), dtype=dtype) 
    categorical[np.arange(n), y] = 1 
    output_shape = input_shape + (num_classes,)
    categorical = np.reshape(categorical, output_shape)
    return categorical

class SurDataSet(Dataset):
    def __init__(self,data,**kwargs):
        super(SurDataSet,self).__init__()
        # self.data = data
        # self.info = info
        self.cls_index = kwargs.get('cls_index')
        self.cls_map_dict = kwargs.get('cls_map_dict')
        self.diameter_index = kwargs.get('diameter_index')
        self.cls_types = kwargs.get('cls_types',None)
        self.num_class = kwargs.get('num_class')
        self.HU_min = kwargs.get('HU_min')
        self.HU_max = kwargs.get('HU_max')
        self.aug = kwargs.get('aug',False)
        self.diameter_ratio = kwargs.get('diameter_ratio',1.5)
        self.input_shape = kwargs.get('input_shape')
        self.nonfix_crop = kwargs.get('nonfix_crop',False)
        self.diameter_max_th = kwargs.get('diameter_max_th',10000000)
        self.diameter_min_th = kwargs.get('diameter_min_th',0)
        self.smooth_val = kwargs.get('smooth_val',0)
        self.convert_dict = kwargs.get('convert_dict',{})
        self.rescale_intensity = kwargs.get('rescale_intensity',False)
        self.high_freq_filter = kwargs.get('high_freq_filter',False)
        self.mode = 'train'
        self.model_mode = kwargs.get('model_mode','basic')
        self.weight_list = []
        self.n_input_channel = kwargs.get('n_input_channel',1)
        self.regression_choice = kwargs.get('regression_choice',False)
        self.reg_map_dict = kwargs.get('reg_map_dict',None)
        self.cls_balance = kwargs.get('cls_balance',False)
             

        if len(self.cls_map_dict.keys())>0:
            self.cls_types = self.cls_map_dict.keys()
            self.num_class = len(self.cls_map_dict.values())
            print ('==================self.num_class is',self.num_class)
        
        with open(data, 'r+') as f:
            self.data_list = json.load(f)['data']
            
    def DataNorm(self,data):
        # self.data = np.clip(self.data,self.HU_min,self.HU_max)
        # self.data = (self.data - self.HU_min)/float(self.HU_max-self.HU_min)
        # self.data = (self.data-0.5)*2
        data = np.clip(data,self.HU_min,self.HU_max)
        data = (data - self.HU_min)/float(self.HU_max-self.HU_min)
        data = (data-0.5)*2
        return data
        
    def GetSampleProb(self):
        prob_val = 1.0/len(self.cls_types)
        self.prob_list = [prob_val for _ in range(len(self.info))]
        sums = sum(self.prob_list)
        self.prob_list = [val/sums for val in self.prob_list]
        

    def _GenerateRatios(self):
        values = []
        for val in self.info:
            if ',' in str(val[self.cls_index]):
                continue
            else:
                if float(val[self.cls_index]) not in self.cls_map_dict.keys():
                    continue
                values.append(self.cls_map_dict[float(val[self.cls_index])])
        counts = Counter(values)
        print ('counts is ',counts)
        counts = [counts[key]/float(len(self.info)) for key in sorted(counts.keys())]
        return counts
        
    
    def _CalculateDiameterRatios(self,idx,start_val,end_val,interval):
        total = len(info)
        diameters = [float(val[self.diameter_index]) for val in info]
        end_val = int(math.ceil(max(diameters)/float(interval)))*interval
        num_interval = round((end_val - start_val)/interval)
        diameter_dict = defaultdict(int)
        for interval_idx in range(int(num_interval)):
            min_val,max_val = interval_idx*interval,(interval_idx+1)*interval
            key = '%03d-%03d'%(min_val,max_val)
            indices = [data_idx for data_idx in range(len(info)) if float(info[data_idx][self.diameter_index])<max_val and float(info[data_idx][self.diameter_index])>=min_val]
            diameter_dict[key] = len(indices)


    def _GenerateTrainData(self,uids=None,uids_test=None,mode='val'):
        cact_index = self.cls_index
        diameter_index = self.diameter_index
        if self.cls_types is None:
            cls_types = list(set([val[cact_index] for val in self.info]))
        else:
            cls_types = self.cls_types
        num_task = 1
        indices = [idx for idx in range(len(self.info)) if float(self.info[idx][self.diameter_index])<=self.diameter_max_th and float(self.info[idx][self.diameter_index])>self.diameter_min_th]
        print ('Number of data filtered out is %d'%(self.info.shape[0]-len(indices)))
        if type(cact_index)==int:
            indices = [idx for idx in indices if ',' not in str(self.info[idx][self.cls_index]) and  float(self.info[idx][self.cls_index]) in cls_types]
        else:
            num_task = len(self.cls_index)
            for temp_index in self.cls_index:
                indices = [idx for idx in indices if ',' not in str(self.info[idx][temp_index]) and  float(self.info[idx][temp_index]) in cls_types]

        if uids is not None:
            indices = [idx for idx in indices if self.info[idx][0] in uids]
        if uids_test is not None:
            num_before = len(indices)
            indices = [idx for idx in indices if self.info[idx][0] not in  uids_test] 
            num_after = len(indices)
            print ('Number of filtered samples %d'%(num_before-num_after))
        self.data = np.take(self.data,indices,axis=0)
        self.info = np.take(self.info,indices,axis=0)
        self.label_list = []
        print ('==================== self.cls_index is ',self.cls_index)
        if type(self.cls_index)==int:
            if self.num_class<=2:
                self.label = np.array([to_categorical(self.cls_map_dict[float(self.info[idx][cact_index])],num_classes=self.num_class) for idx in range(len(self.info))])

                self.arglabel = np.argmax(self.label,axis=1)
            else:
                
                self.label = np.array([self.cls_map_dict[float(self.info[idx][cact_index])] for idx in range(len(self.info))])
                self.arglabel = self.label
        else:
            for inner_cls_index in self.cls_index:
                if self.num_class<=2:
                    self.label = np.array([to_categorical(self.cls_map_dict[float(self.info[idx][inner_cls_index])],num_classes=self.num_class) for idx in range(len(self.info))])
                else:
                    self.label = np.array([self.cls_map_dict[float(self.info[idx][inner_cls_index])] for idx in range(len(self.info))])
                self.label_list.append(self.label)
            self.label_list = np.array(self.label_list)
            self.label_list = np.transpose(self.label_list,(1,0,2))
            self.label = self.label_list
        self.DataNorm(self.data)
        if num_task>1 or (mode!='train'):
            self.indices = np.arange(self.data.shape[0])
        else:
            print ('='*60)
            print ('Balance data')
            print ('Number of samples before is %d'%self.data.shape[0])
            print ('self.label_list',np.array(self.label).shape)
            pos_indices = [idx for idx in range(self.data.shape[0]) if np.argmax(self.label[idx])==1]
            neg_indices = [idx for idx in range(self.data.shape[0]) if np.argmax(self.label[idx])==0]
            print ('Number of pos and neg',len(pos_indices),len(neg_indices))
            max_value = max(len(pos_indices),len(neg_indices))
            if len(pos_indices)!=max_value:
                pos_indices = np.random.choice(pos_indices,max_value)
            if len(neg_indices)!=max_value:
                neg_indices = np.random.choice(neg_indices,max_value)
            all_indices = list(pos_indices) + list(neg_indices)
            self.indices = all_indices
        # random.shuffle(self.indices)
        self.weight_list = [1.0/len(self.indices) for _ in range(len(self.indices))]

    def __getitem__(self,index):
        # if self.mode=='train':
        #     if not self.cls_balance:
        #         data_idx = self.indices[index]
        #     else:
        #         data_idx = np.random.choice(self.indices,size=1,p=self.prob_list)[0]
        # else:
        #     data_idx = index
        # if self.model_mode=='triple':
        #     base_image,base_label = self._get_single_case(data_idx)
        #     same_cls_idx = self._SelectSample(data_idx,True)
        #     diff_cls_idx = self._SelectSample(data_idx,False)
        #     same_cls_image,same_cls_label = self._get_single_case(same_cls_idx)
        #     diff_cls_image,diff_cls_label = self._get_single_case(diff_cls_idx)
        #     net_input_imgs = [base_image,same_cls_image,diff_cls_image]
        #     net_input_labels = [base_label,same_cls_label,diff_cls_label]
        #     net_input_cls_consis = [True,False]
        #     return net_input_imgs,net_input_labels,net_input_cls_consis
        # elif self.model_mode=='deux':
        #     base_image, base_label = self._get_single_case(data_idx)
        #     if np.random.rand()<0.5:
        #         same_flag = True
        #     else:
        #         same_flag = False

        #     second_idx = self._SelectSample(data_idx, same_flag)
        #     second_image,second_label = self._get_single_case(second_idx)
        #     net_input_imgs = [base_image,second_image]
        #     net_input_labels = [base_label,second_label]
        #     net_input_cls_consis = [same_flag]
        #     print('net_input_labels: {}'.format(net_input_labels))
        #     return net_input_imgs,net_input_labels,net_input_cls_consis
        # else:
        #     base_image,base_label = self._get_single_case(data_idx)
        #     net_input_imgs = [base_image]
        #     net_input_labels = [base_label]
        #     net_input_cls_consis = []
        #     return net_input_imgs,net_input_labels,net_input_cls_consis
        data_path = self.data_list[index]
        info_path = data_path.replace('_data.npy', '_info.npy')
        data = np.load(data_path)
        info = np.load(info_path, allow_pickle=True)
        label = info[-2]
        data = self.DataNorm(data)
        return [data], [label], []

    def _getRegLabel(self,label):
        return self.reg_map_dict[label]
            
    def _SelectSample(self,index,same_flag):

        cls_id = self.arglabel[index]

        same_cls_indices = [idx for idx in range(self.label.shape[0]) if self.arglabel[idx]==cls_id]
        diff_cls_indices = [idx for idx in range(self.label.shape[0]) if self.arglabel[idx]!=cls_id]
        random.shuffle(same_cls_indices)
        random.shuffle(diff_cls_indices)

        if same_flag:
            if len(same_cls_indices) == 1:
                return same_cls_indices[0]
            else:
                return same_cls_indices[0] if same_cls_indices[0] != index else same_cls_indices[1]
        else:
            if len(diff_cls_indices) == 0:
                return same_cls_indices[0]
            else:
                return diff_cls_indices[0]
           
        
        
    def _get_single_case(self,index):
        data_idx = index
        
        step_val = 16
        min_patch_size = 16
        
        image_single,label_single,info_single = self.data[data_idx],self.label[data_idx],self.info[data_idx]

        if self.aug:
            image_single = data_augmentation([image_single],prob=0.7)[0]
            if np.random.choice([0,1]):
                degree = random.choice([0, 1, 2])
                if degree != 0:
                    image_single = np.swapaxes(image_single,degree,0)
                else:
                    image_single = np.swapaxes(image_single,1,2)
                
        zero_count = 0
        if self.aug:
            delta_range=5
            deltas = np.random.randint(low=-delta_range,high=delta_range,size=3)
        else:
            deltas = [0 for _ in range(3)]
        if self.nonfix_crop:
            try:
                current_spacing = [float(val) for val in info_single[-7:-4]][::-1]
            except:
                current_spacing = [0.6 for _ in range(3)]
            diameter = float(info_single[self.diameter_index])
            crop_shape = [int(round(diameter*self.diameter_ratio/float(val)/float(step_val)))*step_val for val in current_spacing]
            crop_shape = [max(val,min_patch_size) for val in crop_shape]
            crop_shape = [min(val1,val2) for val1,val2 in zip(np.squeeze(image_single.shape),crop_shape)]
#             print ('crop_shape is',crop_shape)
        else:
            crop_shape = self.input_shape[:3]
            crop_shape = [min(val,80) for val in crop_shape]
        for delta_idx in range(3):
            max_val = max(0,(np.squeeze(image_single).shape[delta_idx]-crop_shape[delta_idx])/2-1)
            abs_val = abs(deltas[delta_idx])
            if abs_val ==0:
                continue
            if abs_val>max_val:
                deltas[delta_idx] = deltas[delta_idx]/abs_val * max_val
        image_single = random_crop(image_single,input_shape=crop_shape,deltas=deltas,aug=self.aug)
        center_slice_idx = int(image_single.shape[0]/2)
        if self.rescale_intensity:
            p2, p98 = np.percentile(image_single, (1, 99))
            image_single = exposure.rescale_intensity(image_single, in_range=(p2, p98))
        if self.high_freq_filter:
            image_single = np.array([HFFilter(single_slice) for single_slice in image_single])
                                        
        if self.n_input_channel==2:
            edges = np.squeeze(np.array([ShowEdges(img_slice,gaussian_choice=1) for img_slice in image_single]))
            edges = edges[np.newaxis,...]
        
        image_single = image_single[np.newaxis,...]
        if self.n_input_channel==2:
            image_single = np.concatenate([image_single,edges],axis=0)
        if self.num_class<=2:
            if type(self.cls_index)==int:
                label_single = np.argmax(label_single)
                label_single = self._labelSmooth(label_single)
            else:
                label_single = [np.argmax(val,axis=-1) for val in label_single]
                label_single = [self._labelSmooth(val) for val in label_single]
        image_single = image_single.copy()
        if self.regression_choice:
            label_single = [label_single,self._getRegLabel(label_single)]
        return image_single,label_single
            
        
    def __len__(self):
        # return len(self.indices)
        return len(self.data_list)


    def _labelSmooth(self,label):
        if label==0:
            label += self.smooth_val
        else:
            label -= self.smooth_val
        return label
    
    def _selectSamples(self,label_val,loss_list,num_samples,num_hard):
        arg_label = np.argmax(self.label,axis=1)
        base_indices = set([idx for idx in range(len(self.label)) if arg_label[idx]==label_val])
        loss_list = [val for val in loss_list if val[0] in base_indices]
        loss_list.sort(key=lambda x:x[1])
        indices0 = [val[0] for val in loss_list[:num_samples]]
        indices2 = [val[0] for val in loss_list[num_samples:]]
        indices1 = np.random.choice(indices2,num_hard,replace=False)
        return list(indices0)+list(indices1)
        

    def _updateIndices(self,loss_list,ratio_pos,num_hard=10):
        loss_list.sort(key=lambda x:x[1])
        ratios = self._GenerateRatios()
        num_samples = int(round(ratio_pos*len(self.data)))
        num_zero_samples = int(round(num_samples*ratios[0]))
        num_one_samples = num_samples-num_zero_samples
        num_zero_hard = max(1,int(round(num_hard*ratios[0])))
        num_one_hard = max(1,num_hard-num_zero_hard)
        zero_indices = self._selectSamples(0,loss_list,num_zero_samples,num_zero_hard)
        one_indices = self._selectSamples(1,loss_list,num_one_samples,num_one_hard)
        self.indices = list(zero_indices)+list(one_indices)
        random.shuffle(self.indices)
        self.GetSampleProb()
        print ('Number of indices remain is %d'%len(self.indices))
    

    def _updateWeights(self,diff_list,increase=False):
        beta = 1/(1+2.0**0.5)
        for idx in range(len(diff_list)):
            data_idx,diff = diff_list[idx]
            diff_val = diff.item()
            diff_val = diff_val**0.5
            if increase:
                self.weight_list[data_idx] = (beta**(-diff_val))
            else:
                self.weight_list[data_idx] = (beta**diff_val)
        
        sum_value = np.sum(self.weight_list)
        self.weight_list = [val/float(sum_value) for val in self.weight_list]
        new_sum_val = np.sum(self.weight_list)  

    def _SelectSamplesBasedOnWeights(self):
        ################# balance?
        new_sum_val = np.sum(self.weight_list)
        print ('new sum value is ',new_sum_val)
        self.indices = np.random.choice(np.arange(len(self.data)),p=self.weight_list,size=len(self.data))
        random.shuffle(self.indices)



# class InferDataSet(Dataset):
#     def __init__(self,data,**kwargs):
#         super(InferDataSet,self).__init__()
#         self.data = data
#         self.info = info
#         self.cls_index = kwargs.get('cls_index')
#         self.cls_map_dict = kwargs.get('cls_map_dict')
#         self.diameter_index = kwargs.get('diameter_index')
#         self.cls_types = kwargs.get('cls_types',None)
#         self.num_class = kwargs.get('num_class')
#         self.HU_min = kwargs.get('HU_min')
#         self.HU_max = kwargs.get('HU_max')
#         self.aug = kwargs.get('aug',False)
#         self.diameter_ratio = kwargs.get('diameter_ratio',1.5)
#         self.input_shape = kwargs.get('input_shape')
#         self.nonfix_crop = kwargs.get('nonfix_crop',False)
#         self.diameter_max_th = kwargs.get('diameter_max_th',10000000)
#         self.diameter_min_th = kwargs.get('diameter_min_th',0)
#         self.smooth_val = kwargs.get('smooth_val',0)
#         self.convert_dict = kwargs.get('convert_dict',{})
#         self.rescale_intensity = kwargs.get('rescale_intensity',False)
#         self.high_freq_filter = kwargs.get('high_freq_filter',False)
#         self.mode = 'train'
#         self.model_mode = kwargs.get('model_mode','basic')
#         self.weight_list = []
#         self.n_input_channel = kwargs.get('n_input_channel',1)
#         self.regression_choice = kwargs.get('regression_choice',False)
#         self.reg_map_dict = kwargs.get('reg_map_dict',None)
#         self.cls_balance = kwargs.get('cls_balance',False)
#         self.DataNorm()

#     def __len__(self):
#         return self.data.shape[0]

#     def DataNorm(self):
#         self.data = np.clip(self.data,self.HU_min,self.HU_max)
#         self.data = (self.data - self.HU_min)/float(self.HU_max-self.HU_min)
#         self.data = (self.data-0.5)*2
        
#     def __getitem__(self,index):
#         data = self.data[index]
#         if self.info is not None:
#             info = self.info[index]
#             return data, info
#         else:
#             return data
class InferDataSet(Dataset):
    def __init__(self,data,**kwargs):
        super(InferDataSet,self).__init__()
        self.data = data
        self.cls_index = kwargs.get('cls_index')
        self.cls_map_dict = kwargs.get('cls_map_dict')
        self.diameter_index = kwargs.get('diameter_index')
        self.cls_types = kwargs.get('cls_types',None)
        self.num_class = kwargs.get('num_class')
        self.HU_min = kwargs.get('HU_min')
        self.HU_max = kwargs.get('HU_max')
        self.aug = kwargs.get('aug',False)
        self.diameter_ratio = kwargs.get('diameter_ratio',1.5)
        self.input_shape = kwargs.get('input_shape')
        self.nonfix_crop = kwargs.get('nonfix_crop',False)
        self.diameter_max_th = kwargs.get('diameter_max_th',10000000)
        self.diameter_min_th = kwargs.get('diameter_min_th',0)
        self.smooth_val = kwargs.get('smooth_val',0)
        self.convert_dict = kwargs.get('convert_dict',{})
        self.rescale_intensity = kwargs.get('rescale_intensity',False)
        self.high_freq_filter = kwargs.get('high_freq_filter',False)
        self.mode = 'train'
        self.model_mode = kwargs.get('model_mode','basic')
        self.weight_list = []
        self.n_input_channel = kwargs.get('n_input_channel',1)
        self.regression_choice = kwargs.get('regression_choice',False)
        self.reg_map_dict = kwargs.get('reg_map_dict',None)
        self.cls_balance = kwargs.get('cls_balance',False)
        with open(data, 'r+') as f:
            self.data_list = json.load(f)['data']
            
    def DataNorm(self,data):
        # self.data = np.clip(self.data,self.HU_min,self.HU_max)
        # self.data = (self.data - self.HU_min)/float(self.HU_max-self.HU_min)
        # self.data = (self.data-0.5)*2
        data = np.clip(data,self.HU_min,self.HU_max)
        data = (data - self.HU_min)/float(self.HU_max-self.HU_min)
        data = (data-0.5)*2
        return data

    def __len__(self):
        return len(self.data_list)
        
    def __getitem__(self,index):
        data_path = self.data_list[index]
        info_path = data_path.replace('_data.npy', '_info.npy')
        data = np.load(data_path)
        info = np.load(info_path, allow_pickle=True)
        data = self.DataNorm(data)
        return data, info

class TestDataSet(Dataset):
    def __init__(self,data,**kwargs):
        super(TestDataSet,self).__init__()
        self.data = data
        self.cls_index = kwargs.get('cls_index')
        self.cls_map_dict = kwargs.get('cls_map_dict')
        self.diameter_index = kwargs.get('diameter_index')
        self.cls_types = kwargs.get('cls_types',None)
        self.num_class = kwargs.get('num_class')
        self.HU_min = kwargs.get('HU_min')
        self.HU_max = kwargs.get('HU_max')
        self.aug = kwargs.get('aug',False)
        self.diameter_ratio = kwargs.get('diameter_ratio',1.5)
        self.input_shape = kwargs.get('input_shape')
        self.nonfix_crop = kwargs.get('nonfix_crop',False)
        self.diameter_max_th = kwargs.get('diameter_max_th',10000000)
        self.diameter_min_th = kwargs.get('diameter_min_th',0)
        self.smooth_val = kwargs.get('smooth_val',0)
        self.convert_dict = kwargs.get('convert_dict',{})
        self.rescale_intensity = kwargs.get('rescale_intensity',False)
        self.high_freq_filter = kwargs.get('high_freq_filter',False)
        self.mode = 'train'
        self.model_mode = kwargs.get('model_mode','basic')
        self.weight_list = []
        self.n_input_channel = kwargs.get('n_input_channel',1)
        self.regression_choice = kwargs.get('regression_choice',False)
        self.reg_map_dict = kwargs.get('reg_map_dict',None)
        self.cls_balance = kwargs.get('cls_balance',False)
        print('data cfg:')
        print(kwargs)
        with open(data, 'r+') as f:
            self.data_list = json.load(f)['data']
            
    def DataNorm(self,data):
        # self.data = np.clip(self.data,self.HU_min,self.HU_max)
        # self.data = (self.data - self.HU_min)/float(self.HU_max-self.HU_min)
        # self.data = (self.data-0.5)*2
        data = np.clip(data,self.HU_min,self.HU_max)
        data = (data - self.HU_min)/float(self.HU_max-self.HU_min)
        data = (data-0.5)*2
        return data

    def __len__(self):
        return len(self.data_list)
        
    def __getitem__(self,index):
        data_path = self.data_list[index]
        info_path = data_path.replace('_data.npy', '_info.npy')
        data = np.load(data_path)
        info = np.load(info_path, allow_pickle=True)
        label = info[-2]
        is_limit = info[5]
        data = self.DataNorm(data)
        return data, label, is_limit