import os
import sys
import keras
import tensorflow as tf
from keras import Model
from keras import backend as K
from keras.regularizers import l2
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,Dot,\
GlobalAveragePooling3D,Dense,GlobalMaxPooling3D,Lambda,Activation,Reshape,Permute, PReLU, Deconvolution3D,\
Multiply,GlobalAveragePooling3D,Dense,Softmax


sys.path.append('../baseLayers/')
from CommonLayers import NormActi,ConvUnit


def SpatialGather(x,probs,upsample_ratio=None):
    '''
    x represents features and probs represents coarse segmentation prob map
    '''
    
    batch_size,depth,height,width,num_classes = probs._keras_shape
    if upsample_ratio:
        upsample_ratio = 2
    upsample_ratio = int(depth/x._keras_shape[1])
    if upsample_ratio>1:
        x = UpSampling3D(size=upsample_ratio,name='SpatialGather')(x)

    num_features = x._keras_shape[-1]
    
    reshape_x = Reshape(target_shape=(depth*height*width,num_features))(x)
#     if num_classes == 1:
#         probs = Concatenate()([1-probs,probs])
#         num_classes = 2
    reshape_probs = Reshape(target_shape=(depth*height*width,num_classes))(probs)
    
    result = Dot(axes=1)([reshape_probs,reshape_x])
    result = Reshape(target_shape=(1,1,result._keras_shape[1],result._keras_shape[2]))(result)
    return result
    
def OCRModule(x,proxy,**kwargs):
    
    norm_func = kwargs.get('norm_func')
    activation_func = kwargs.get('activation_func')
    num_filters = kwargs.get('num_filters')
    kernel_size = kwargs.get('kernel_size',1)
    padding = kwargs.get('padding','same')

    atrous_rate = kwargs.get('atrous_rate',1)
    conv_first = kwargs.get('conv_first',True)
    dropout_rate= kwargs.get('dropout_rate',0)
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',l2(1e-4))
    
    block_prefix = 'OCRModule'
    
    input_features = x._keras_shape[-1]
    query = x
    for layer_idx in range(2):
        query = ConvUnit(query,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+'_query',
                     kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=layer_idx+1,conv_first=conv_first)
    
    key,value = proxy,proxy
    for layer_idx in range(2):
        key = ConvUnit(key,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+'_key',
                     kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=layer_idx+1,conv_first=conv_first)
    value = ConvUnit(value,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+'_value',
                     kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=1,conv_first=conv_first)
    
    batch,query_depth,query_height,query_width,query_filters = query._keras_shape
    batch,key_depth,key_height,key_width,key_filters = key._keras_shape
    
    query = Reshape(target_shape=(query_depth*query_height*query_width,query_filters))(query)
    
    key = Reshape(target_shape=(key_depth*key_height*key_width,key_filters))(key)
    value = Reshape(target_shape=(key_depth*key_height*key_width,key_filters))(value)
    value = Permute(dims=[2,1])(value)

    sim_map = Dot(name='%s_SimMapDot'%block_prefix,axes=2)([query,key])
    sim_map = Softmax(axis=-1,name='%s_SimMapSoftmax'%block_prefix)(sim_map)
    
    context = Dot(axes=2)([sim_map,value])

    context = Reshape(target_shape=(query_depth,query_height,query_width,query_filters))(context)
    
    combine_feature = Concatenate()([x,context])
    
    final_result = ConvUnit(combine_feature,norm_func=norm_func,activation_func=activation_func,num_filters=input_features,
                    kernel_size=kernel_size,atrous_rate=atrous_rate,padding=padding,block_prefix=block_prefix+'_FinalCombine',
                     kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=1,conv_first=conv_first)
    return final_result