import os
import sys
import random
import numpy as np
import pandas as pd
import scipy.ndimage as nd
from ReadData import load_data,load_single_data
from data_aug import data_augmentation,random_add
from tensorflow.keras.utils import to_categorical
from ImagePreprocess import Data_preprocess,CalculatePatchSizeWithDiameter

# densityTypeConvertMap = {
#     1:1,
#     2:2,
#     3:3,
#     4:4,
#     5:1,
#     6:4
    
    
# }

def UNet_3D_Generator(data_paths,mask_paths,info_paths,**kwargs):
    '''
    Define generator for Unet
      @param:data_paths,mask_paths,info_paths. Pahts of data/mask/info, they should have same shape
      @param:clip_choice. Whether to do clip on intensity or not
      @param:HU_min,HU_max. If clip_choice is True then this range will be used to do clip
      @param:aug choice
      @param:batch_size, default value is 32
    '''
    clip_choice = kwargs.get('clip_choice',True)
    HU_min,HU_max = kwargs.get('HU_min'),kwargs.get('HU_max')
    aug = kwargs.get('aug',True)
    nonfix_crop = kwargs.get('nonfix_crop',False)
    stride_ratio = kwargs.get('stride_ratio',8)
    densityTypeIndex = kwargs.get('densityTypeIndex',-1)
    batch_size = kwargs.get('batch_size',4)
    input_size = kwargs.get('input_size')
    diameter_enlarge_ratio = kwargs.get('diameter_enlarge_ratio',1.5)
    direction_z_ratio = kwargs.get('direction_z_ratio',1.0)
    min_patch_size = stride_ratio




    config_dict = kwargs.get('config_dict')
    deep_supervision = config_dict['deep_supervision']
    aux_task = config_dict['aux_task']
    num_classes = config_dict['num_classes']
    densityTypeConvertMap = config_dict['densityTypeConvertMap']
 
    print ("densityTypeIndex  is",densityTypeIndex)
    

    batch_image,batch_label,batch_densityType = 0,0,0
    result_list = load_data([data_paths,mask_paths,info_paths],aux_task=aux_task)
    data,mask,info = result_list

    print ('After loading data',data.shape,mask.shape,info.shape)
    test_idx = 0
    # while test_idx<2:
    #     test_idx +=1

    while True:
        sample_indices = np.arange(data.shape[0])
        random.shuffle(sample_indices)
        
        for inner_idx in range(0,len(sample_indices)):
            current_sample_idx = sample_indices[inner_idx]
            current_image,current_mask = data[current_sample_idx],mask[current_sample_idx]
            current_info = info[current_sample_idx]
            if aux_task:
                densityType = int(float(current_info[densityTypeIndex]))

                if densityType==5:
                    densityType = 1
                elif densityType == 6:
                    densityType = 4
                if densityType>=7:
                    continue
                
                densityType -=1
 
                current_densityType = to_categorical(densityType,num_classes=4)                
                current_densityType = np.reshape(np.array(current_densityType),(1,4))

            else:
                current_densityType = 0
            # print ('shape of current_densityType',current_densityType.shape)
            #### Crop patch 
            if nonfix_crop:
                patch_shape = CalculatePatchSizeWithDiameter(current_mask,stride_ratio=stride_ratio,diameter_enlarge_ratio=diameter_enlarge_ratio,direction_z_ratio=direction_z_ratio)
            else:
                patch_shape = input_size
            final_image,final_mask = Data_preprocess(current_image,current_mask,aug=aug,input_shape=patch_shape,shift_ratio=1.0)
            if 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(-5, 5)
                    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 clip_choice:
                final_image = np.clip(final_image,HU_min,HU_max)
            
            final_image = (final_image - HU_min)/float(HU_max - HU_min)
            
            if aug:
                final_image = random_add(final_image)
        
            final_image = np.squeeze(final_image)[np.newaxis,...,np.newaxis]
            final_mask = np.squeeze(final_mask)[np.newaxis,...,np.newaxis]
            if type(batch_image) == int:
                batch_image = np.asarray(final_image)
                batch_mask = np.asarray(final_mask)
                batch_densityType = np.array([current_densityType])
            else:
                batch_image = np.concatenate([batch_image,final_image],axis=0)
                batch_mask = np.concatenate([batch_mask,final_mask],axis=0)
                batch_densityType = np.concatenate([batch_densityType,[current_densityType]],axis=0)
                
            if batch_mask.shape[0] == batch_size:
                # batch_densityType = np.squeeze(batch_densityType)
                # print ('batch_densityType shape',batch_densityType.shape)
                
                if deep_supervision:
                    batch_mask = [batch_mask for _ in range(3)]
                if aux_task:
                    batch_densityType = np.reshape(batch_densityType,(batch_size,4))
                    if deep_supervision:
                        batch_mask.append(batch_densityType)
                    else:
                        batch_mask = [batch_mask,batch_densityType]
                yield batch_image,batch_mask
                batch_image,batch_label,batch_densityType = 0,0,0


def DataStat(info_paths,diameter_ths,**kwargs):

    infos = load_single_data(info_paths,aux_task=True)
    densityTypeIndex = kwargs.get('densityTypeIndex',-4)
    from collections import defaultdict
    diameter_index = kwargs.get('diameter_index',4)
    
    uids = set([x[0] for x in infos])
    print ('Number of records in total is %d in %d series'%(len(infos),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:
        diameter_min,diameter_max = diam_th_vals
        density_stat_map = defaultdict(int)
        count = 0
        for record in infos:
            diameter = float(record[diameter_index])
            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 record between [%d,%d] is '%(count,diameter_min,diameter_max))
        for key in keys:
            print ('%s:%d'%(key,density_stat_map[key]))