import os
import sys
import glob
import numpy as np

import tensorflow as tf
from keras import Model
from keras.regularizers import l2
from keras import backend as K
from keras.engine import Layer,InputSpec
from keras.layers.merge import concatenate
from keras.callbacks import TensorBoard,Callback
from keras.layers.advanced_activations import LeakyReLU
from keras.preprocessing.image import ImageDataGenerator
from keras.layers.normalization import BatchNormalization
from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
from keras import initializers, regularizers, constraints,optimizers
from keras.callbacks import ModelCheckpoint, LearningRateScheduler,TensorBoard
from keras.layers import Add,Input,Conv3D,Convolution3D,Dropout,UpSampling3D,Concatenate,MaxPooling3D,\
GlobalAveragePooling3D,Dense,GlobalMaxPooling3D,Lambda,Activation,Reshape,Permute, PReLU, Deconvolution3D,Conv3DTranspose

sys.path.append('../baseLayers/')
from unet_utils import *
from SEB import SEB_block
from OCR import SpatialGather,OCRModule
from CommonLayers import NormActi,ConvUnit


def UpsampleBlock(x,**kwargs):
    num_filters = kwargs.get('num_filters')
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',l2(1e-4))
    block_pre = kwargs.get('block_pre','')
    kernel_size = kwargs.get('kernel_size',3)
    rate = kwargs.get('rate',2)
    upsample_choice = kwargs.get('upsample_choice','upsample')
    norm_func = kwargs.get('norm_func')
    activation_func = kwargs.get('activation_func')
    atrous_rate = kwargs.get('atrous_rate',1)
    padding = kwargs.get('padding','same')
    if type(activation_func)==str:
        current_acti = activation_func
        activation_func = None
    else:
        current_acti = None
#     if rate>1:
#     print ('kernel of %s is %d'%(block_pre,kernel_size))
#     print ('current_acti is ',current_acti,norm_func)
    if upsample_choice == 'upsample':
        x = UpSampling3D(size=rate,name='%s_upsample'%(block_pre))(x)
        x = Convolution3D(nb_filter = num_filters,kernel_size=kernel_size, 
                             kernel_initializer=kernel_initializer, kernel_regularizer=kernel_regularizer,
                            border_mode=padding,activation=current_acti,name='%s_conv'%block_pre)(x)     
    elif upsample_choice == 'deconv':
        x = Deconvolution3D(filters = num_filters,kernel_size=kernel_size, strides=rate,
                     kernel_initializer=kernel_initializer, kernel_regularizer=kernel_regularizer,
                    padding=padding,activation=current_acti,name='%s_conv'%block_pre)(x) 
    else:
        x = Conv3DTranspose(filters = num_filters,kernel_size=kernel_size, strides=rate,
                     kernel_initializer=kernel_initializer, kernel_regularizer=kernel_regularizer,
                    padding=padding,activation=current_acti,name='%s_conv'%block_pre)(x)  

    x = NormActi(x,norm_func=norm_func,activation_func=activation_func,
                               block_prefix=block_pre,block_suffix=block_pre)
#     else:
#         x = ConvUnit(x,norm_func=norm_func,activation_func=activation_func,num_filters=num_filters,
#                     kernel_size=kernel_size,atrous_rate=atrous_rate,padding=padding,block_prefix=block_pre,
#                      kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
#                     layer_idx=1,conv_first=True)
    return x
    
    
    

def NoduleSegDecoderBlock(x,encoder_tensors,**kwargs):
    block_idx = kwargs.get('block_idx')
    num_units = kwargs.get('num_units')
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',l2(1e-4))
    SEB_choice = kwargs.get('SEB_choice',False)
    kernel_size = kwargs.get('kernel_size',3)
    norm_func = kwargs.get('norm_func')
    activation_func = kwargs.get('activation_func')
    
    
    num_filters = kwargs.get('num_filters')
    padding = kwargs.get('padding','same')
    
    atrous_rate = kwargs.get('atrous_rate',1)
    stride = kwargs.get('upsample_rate',2)
    conv_first = kwargs.get('conv_first',True)
    merge_axis = kwargs.get('merge_axis',-1)
    upsample_choice = kwargs.get('upsample_choice','upsample')
    
    DIM1_AXIS,DIM2_AXIS,DIM3_AXIS = 1,2,3
    block_prefix = 'NoduleSegDecoder_Block%02d'%block_idx

    x1 = UpsampleBlock(x,num_filters=num_filters,kernel_size=kernel_size,rate=stride,atrous_rate=atrous_rate,
                 padding=padding,block_pre=block_prefix,kernel_initializer=kernel_initializer,
                 kernel_regularizer=kernel_regularizer,upsample_choice=upsample_choice,
                 norm_func=norm_func,activation_func=activation_func)
    
    target_depth,target_height,target_width = x1._keras_shape[DIM1_AXIS],x1._keras_shape[DIM2_AXIS],x1._keras_shape[DIM3_AXIS]
    
    depth_before_SEB,height_before_SEB,width_before_SEB =  x._keras_shape[DIM1_AXIS],x._keras_shape[DIM2_AXIS],x._keras_shape[DIM3_AXIS]
    
    if not SEB_choice:
        encoder_target_tensor = encoder_tensors[-(block_idx+1)]
    else:
        target_tensor = encoder_tensors[-(block_idx+1)]
        SEB_tensors = encoder_tensors[-(block_idx):][::-1]
        
        upsample_tensor_list = []
        for layer_idx,current_tensor in enumerate(SEB_tensors):
            current_depth,current_height,current_width = current_tensor._keras_shape[DIM1_AXIS],current_tensor._keras_shape[DIM2_AXIS],current_tensor._keras_shape[DIM3_AXIS]
            upsample_rate = [int(depth_before_SEB/current_depth),int(height_before_SEB/current_height),int(width_before_SEB/current_width)]
            upsample_tensor = UpSampling3D(size=upsample_rate,name='%s_SEB_preUpsample%02d'%(block_prefix,layer_idx+1))(current_tensor)
            upsample_tensor_list.append(upsample_tensor)
            
        target_num_filters = target_tensor._keras_shape[-1]
        
        encoder_target_tensor = SEB_block([upsample_tensor_list,target_tensor],norm_func=norm_func,activation_func=activation_func,
                                         num_filters=target_num_filters,kernel_size=kernel_size,atrous_rate=atrous_rate,padding=padding,
                                         block_prefix=block_prefix,kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                                         conv_first=conv_first,merge_axis=merge_axis)

    x = Concatenate(axis=merge_axis,name='%s_Concatenate'%(block_prefix))([encoder_target_tensor,x1])
    
    for layer_idx in range(1,num_units):
        x = ConvUnit(x,norm_func=norm_func,activation_func=activation_func,num_filters=num_filters,
                    kernel_size=kernel_size,atrous_rate=atrous_rate,padding=padding,block_prefix=block_prefix,
                     kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=layer_idx+1,conv_first=conv_first)
    return x


def NoduleSegDeepCombineBlock(input_tensors,**kwargs):
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',l2(1e-4))
    kernel_size = kwargs.get('kernel_size',3)
    norm_func = kwargs.get('norm_func')
    activation_func = kwargs.get('activation_func')
    
    num_filters = kwargs.get('num_filters')
    
    padding = kwargs.get('padding','same')
    atrous_rate = kwargs.get('atrous_rate',1)
    stride = kwargs.get('upsample_rate',2)
    conv_first = kwargs.get('conv_first',True)
    merge_axis = kwargs.get('merge_axis',-1)
    upsample_choice = kwargs.get('upsample_choice','upsample')
    
    print ('input_tensors is ',input_tensors)
    x1,x2,x3 = input_tensors
    block_prefix = 'NoduleSegDeepCombineBlock'

    uprate_1 = int(x2._keras_shape[1]/x1._keras_shape[1])
    uprate_2 = int(x3._keras_shape[1]/x2._keras_shape[1])
    x1 = ConvUnit(x1,norm_func=norm_func,activation_func=activation_func,num_filters=num_filters,
                    kernel_size=kernel_size,atrous_rate=atrous_rate,padding=padding,block_prefix=block_prefix+'_block01_',
                     kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=1,conv_first=conv_first)
    up_x1 = UpSampling3D(size=uprate_1,name='%s_block01_upsample'%block_prefix)(x1)

    
    x2 = ConvUnit(x2,norm_func=norm_func,activation_func=activation_func,num_filters=num_filters,
                    kernel_size=kernel_size,atrous_rate=atrous_rate,padding=padding,block_prefix=block_prefix+'_block02_',
                     kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=1,conv_first=conv_first)
    up_x2 = UpSampling3D(size=uprate_2,name='%s_block02_upsample'%block_prefix)(Add(name='%s_Add01'%block_prefix)([up_x1,x2]))
    
    
    x3 = ConvUnit(x3,norm_func=norm_func,activation_func=activation_func,num_filters=num_filters,
                    kernel_size=kernel_size,atrous_rate=atrous_rate,padding=padding,block_prefix=block_prefix+'_block03_',
                     kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=1,conv_first=conv_first)
    
    final_result = Add(name='%s_Add02'%block_prefix)([up_x2,x3])
    
    return final_result,[x1,x2,x3]
    
    
    
    
    
def NoduleSegDecoder_proxima(encoder_result,**kwargs):
    
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',l2(1e-4))
    kernel_size = kwargs.get('kernel_size',3)
    final_kernel_size = kwargs.get('final_kernel_size',3)
    norm_func = kwargs.get('norm_func')
    activation_func = kwargs.get('activation_func')
    padding = kwargs.get('padding','same')
    atrous_rate = kwargs.get('atrous_rate',1)
    conv_first = kwargs.get('conv_first',True)
    merge_axis = kwargs.get('merge_axis',-1)
    
    SEB_choice = kwargs.get('SEB_choice',True)
    ACF_choice = kwargs.get('ACF_choice',False)
    OCR_choice = kwargs.get('OCR_choice',False)
    deep_supervision = kwargs.get('deep_supervision',True)
    num_units = kwargs.get('num_units',[3,3,3,3])
    seg_num_class = kwargs.get('seg_num_class')
    aux_task = kwargs.get('aux_task',False)
    num_classes = kwargs.get('num_classes')
    classification_layers = kwargs.get('classification_layers')
    upsample_choice = kwargs.get('upsample_choice','upsample')
    divide_ratio = kwargs.get('divide_ratio',4)
    
    encoder_model = encoder_result[0]
    encoder_tensors = encoder_result[1]
    
    try:
        encoder_stride_list = encoder_result[2]
    except:
        encoder_stride_list = [None for _ in range(len(encoder_tensors))]


    input_tensor = encoder_model.input
    
    x = encoder_tensors[-1]
    decoder_tensors = []

    for block_idx in range(len(encoder_tensors)-2):
        if encoder_stride_list[-(block_idx+1)]:
            current_stride = encoder_stride_list[-(block_idx+1)]
        else:
            current_stride = int(encoder_tensors[-(block_idx+2)]._keras_shape[1]/x._keras_shape[1])
        num_filters = int(x._keras_shape[-1]/2)
        if block_idx == len(encoder_tensors)-3:
            SEB_choice = False
        x = NoduleSegDecoderBlock(x,encoder_tensors,block_idx=block_idx+1,num_units=num_units[block_idx],
                                 kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                                 SEB_choice=SEB_choice,kernel_size=kernel_size,norm_func=norm_func,activation_func=activation_func,
                                 num_filters=num_filters,padding=padding,atrous_rate=atrous_rate,upsample_rate=current_stride,conv_first=conv_first,
                                 merge_axis=merge_axis,upsample_choice=upsample_choice)
        decoder_tensors.append(x)
    
    start_idx = len(decoder_tensors)-3

    current_stride_list =   encoder_stride_list[-4:-1][::-1]

    combine_tensors = NoduleSegDeepCombineBlock(decoder_tensors[start_idx:start_idx+4],kernel_initializer=kernel_initializer,
                                               kernel_regularizer=kernel_regularizer,kernel_size=kernel_size,norm_func=norm_func,
                                               activation_func=activation_func,num_filters=int(decoder_tensors[-1]._keras_shape[-1]/divide_ratio),
                                                padding=padding,atrous_rate=atrous_rate,conv_first=conv_first,merge_axis=merge_axis,upsample_choice=upsample_choice,
                                                stride_list=current_stride_list)
    
    aux_conv_tensor = 0
    if seg_num_class==1:
        current_acti = 'sigmoid'
    else:
        current_acti = 'softmax'
    
    final_conv_func = Convolution3D(nb_filter = seg_num_class,kernel_size=final_kernel_size, 
                                 kernel_initializer=kernel_initializer, kernel_regularizer=kernel_regularizer,
                                border_mode='same',activation=current_acti,name='NoduleSegDecoder_conv3d_mins1')

    if not deep_supervision:
        block_prefix = 'NoduleSegDecoder_NotSupervision'
        conv_mins1 = final_conv_func(combine_tensors[0])
        aux_conv_tensor = conv_mins1
    
        ########### TODO:revise this part during non fix
        if ACF_choice:
            coarse_feature_map = conv_mins1

            conv_mins1 = []
            feature_map_class = CCB([coarse_feature_map,combine_tensors[0]],num_classes=seg_num_class)
            CAB_feature_map = CAB([coarse_feature_map,feature_map_class],num_classes=seg_num_class,activaion=current_acti)
            concatenate_feature_map = Concatenate()([CAB_feature_map,combine_tensors[0]])
            aux_conv_tensor = concatenate_feature_map
            final_result = Convolution3D(nb_filter = seg_num_class,kernel_size=final_kernel_size, 
                                 kernel_initializer=kernel_initializer, kernel_regularizer=kernel_regularizer,
                                border_mode='same',activation=current_acti,name='NoduleSegDecoder_conv_ACF')(concatenate_feature_map)
            
            conv_mins1 = [coarse_feature_map,final_result]
    else:
        conv_mins1 = []
        input_tensors = combine_tensors[1]
        rate = np.product(current_stride_list[:len(input_tensors)])
        for idx in range(len(input_tensors)):
            # rate = 2 ** (2-idx)
            rate /= current_stride_list[idx]
            rate = int(rate)
            print ('During deep supervision rate of stage %d is %d'%(idx,rate))

            # rate = int(input_tensor._keras_shape[1]/input_tensors[idx]._keras_shape[1])
            block_pre = 'NoduleSegDecoder_DeepSupervision_block%02d'%(idx+1)
            current_tensor = UpsampleBlock(input_tensors[idx],num_filters=seg_num_class,kernel_size=final_kernel_size,rate=rate,
                                  kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                                   block_pre=block_pre,upsample_choice=upsample_choice,activation_func=current_acti)
                
            conv_mins1.append(current_tensor)
            
        if OCR_choice:
            features = combine_tensors[1][-1]
            coarse_seg_output = final_conv_func(conv_mins1[-1])
            print ('feature shape and coarse_seg_output shape',features.shape,coarse_seg_output.shape)
            proxy = SpatialGather(features,coarse_seg_output)
            ocr_output = OCRModule(features,proxy,norm_func=norm_func,activation_func=activation_func,num_filters=features._keras_shape[-1]*4,
                                   kernel_size=1,kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer)
            ocr_output = Convolution3D(nb_filter = seg_num_class,kernel_size=final_kernel_size, 
                                 kernel_initializer=kernel_initializer, kernel_regularizer=kernel_regularizer,
                                border_mode='same',activation=current_acti,name='NoduleSegDecoder_conv3d_ocr_output')(ocr_output)
            
            conv_mins1.append(ocr_output)
 
        
        
    output = conv_mins1
    if aux_task:
        ### Do classification
        pool_result = classification_layers(name='aux_task_pool_layer')(combine_tensors[1][-1])
        ### Dense unit
        dense_result = Dense(num_classes*32)(pool_result)
        class_result = Dense(num_classes)(dense_result)
        output.append(class_result)
    model = Model(input_tensor,output)
    return model
    
    
        
    
    
    
    
    
    
    


