from keras.layers import GlobalAveragePooling2D, GlobalMaxPooling2D, Reshape, Dense, multiply, Permute, Concatenate, Conv2D, Add, Activation, Lambda
from keras import backend as K
from keras.activations import sigmoid

def attach_attention_module(net, attention_module):
  if attention_module == 'se_block': # SE_block
    net = se_block(net)
  elif attention_module == 'cbam_block': # CBAM_block
    net = cbam_block(net)
  elif attention_module == 'channel_block': # only channel_block
    net = channel_attention(net, ratio=8)
  elif attention_module == 'scan_block': # scan_block
    net = scan_block(net)
  else:
    raise Exception("'{}' is not supported attention module!".format(attention_module))

  return net

def se_block(input_feature, ratio=8):
	"""Contains the implementation of Squeeze-and-Excitation(SE) block.
	As described in https://arxiv.org/abs/1709.01507.
	"""
	
	channel_axis = 1 if K.image_data_format() == "channels_first" else -1
	channel = input_feature._keras_shape[channel_axis]

	se_feature = GlobalAveragePooling2D()(input_feature)
	se_feature = Reshape((1, 1, channel))(se_feature)
	assert se_feature._keras_shape[1:] == (1,1,channel)
	se_feature = Dense(channel // ratio,
					   activation='relu',
					   kernel_initializer='he_normal',
					   use_bias=True,
					   bias_initializer='zeros')(se_feature)
	assert se_feature._keras_shape[1:] == (1,1,channel//ratio)
	se_feature = Dense(channel,
					   activation='sigmoid',
					   kernel_initializer='he_normal',
					   use_bias=True,
					   bias_initializer='zeros')(se_feature)
	assert se_feature._keras_shape[1:] == (1,1,channel)
	if K.image_data_format() == 'channels_first':
		se_feature = Permute((3, 1, 2))(se_feature)

	se_feature = multiply([input_feature, se_feature])
	return se_feature


def cbam_block(cbam_feature, ratio=8):
	"""Contains the implementation of Convolutional Block Attention Module(CBAM) block.
	As described in https://arxiv.org/abs/1807.06521.
	"""

	cbam_feature = channel_attention(cbam_feature, ratio)
	cbam_feature = spatial_attention(cbam_feature)
	return cbam_feature


def scan_block(scan_feature, ratio=8):

	scan_feature = channel_attention(scan_feature, ratio)
	scan_feature = frequency_attention(scan_feature)
	return scan_feature


def channel_attention(input_feature, ratio=8):
	
	channel_axis = 1 if K.image_data_format() == "channels_first" else -1
	channel = input_feature._keras_shape[channel_axis]
	
	shared_layer_one = Dense(channel//ratio,
							 activation='relu',
							 kernel_initializer='he_normal',
							 use_bias=True,
							 bias_initializer='zeros')
	shared_layer_two = Dense(channel,
							 kernel_initializer='he_normal',
							 use_bias=True,
							 bias_initializer='zeros')
	
	avg_pool = GlobalAveragePooling2D()(input_feature)    
	avg_pool = Reshape((1,1,channel))(avg_pool)
	assert avg_pool._keras_shape[1:] == (1,1,channel)
	avg_pool = shared_layer_one(avg_pool)
	assert avg_pool._keras_shape[1:] == (1,1,channel//ratio)
	avg_pool = shared_layer_two(avg_pool)
	assert avg_pool._keras_shape[1:] == (1,1,channel)
	
	max_pool = GlobalMaxPooling2D()(input_feature)
	max_pool = Reshape((1,1,channel))(max_pool)
	assert max_pool._keras_shape[1:] == (1,1,channel)
	max_pool = shared_layer_one(max_pool)
	assert max_pool._keras_shape[1:] == (1,1,channel//ratio)
	max_pool = shared_layer_two(max_pool)
	assert max_pool._keras_shape[1:] == (1,1,channel)
	
	cbam_feature = Add()([avg_pool,max_pool])
	cbam_feature = Activation('sigmoid')(cbam_feature)
	
	if K.image_data_format() == "channels_first":
		cbam_feature = Permute((3, 1, 2))(cbam_feature)
	
	return multiply([input_feature, cbam_feature])

def spatial_attention(input_feature):
	kernel_size = 7
	
	if K.image_data_format() == "channels_first":
		channel = input_feature._keras_shape[1]
		cbam_feature = Permute((2,3,1))(input_feature)
	else:
		channel = input_feature._keras_shape[-1]
		cbam_feature = input_feature
	
	avg_pool = Lambda(lambda x: K.mean(x, axis=3, keepdims=True))(cbam_feature)
	assert avg_pool._keras_shape[-1] == 1
	max_pool = Lambda(lambda x: K.max(x, axis=3, keepdims=True))(cbam_feature)
	assert max_pool._keras_shape[-1] == 1
	concat = Concatenate(axis=3)([avg_pool, max_pool])
	assert concat._keras_shape[-1] == 2
	cbam_feature = Conv2D(filters = 1,
					kernel_size=kernel_size,
					strides=1,
					padding='same',
					activation='sigmoid',
					kernel_initializer='he_normal',
					use_bias=False)(concat)	
	assert cbam_feature._keras_shape[-1] == 1
	
	if K.image_data_format() == "channels_first":
		cbam_feature = Permute((3, 1, 2))(cbam_feature)
		
	return multiply([input_feature, cbam_feature])

def frequency_attention(input_feature):

	scan_feature = Permute((2, 3, 1))(input_feature)
	col = input_feature._keras_shape[1]

	'''
	여기다 채워주세요
	'''

	shared_layer_one = Dense(col,
							 activation='relu',
							 kernel_initializer='he_normal',
							 use_bias=True,
							 bias_initializer='zeros')
	shared_layer_two = Dense(col,
							 kernel_initializer='he_normal',
							 use_bias=True,
							 bias_initializer='zeros')

	avg_pool1 = GlobalAveragePooling2D()(scan_feature)
	avg_pool1 = Reshape((1,1,col))(avg_pool1)
	assert avg_pool1._keras_shape[1:] == (1,1,col)
	avg_pool1 = shared_layer_one(avg_pool1)
	assert avg_pool1._keras_shape[1:] == (1,1,col)
	avg_pool1 = shared_layer_two(avg_pool1)
	assert avg_pool1._keras_shape[1:] == (1,1,col)

	max_pool1 = GlobalMaxPooling2D()(scan_feature)
	max_pool1 = Reshape((1,1,col))(max_pool1)
	assert max_pool1._keras_shape[1:] == (1,1,col)
	max_pool1 = shared_layer_one(max_pool1)
	assert max_pool1._keras_shape[1:] == (1,1,col)
	max_pool1 = shared_layer_two(max_pool1)
	assert max_pool1._keras_shape[1:] == (1,1,col)

	scan_feature1 = Add()([avg_pool1,max_pool1])
	scan_feature1 = Activation('sigmoid')(scan_feature1)

	scan_feature = Permute((2, 3, 1))(scan_feature)

	shared_layer_one = Dense(col,
							 activation='relu',
							 kernel_initializer='he_normal',
							 use_bias=True,
							 bias_initializer='zeros')
	shared_layer_two = Dense(col,
							 kernel_initializer='he_normal',
							 use_bias=True,
							 bias_initializer='zeros')

	avg_pool2 = GlobalAveragePooling2D()(scan_feature)
	avg_pool2 = Reshape((1,1,col))(avg_pool2)
	assert avg_pool2._keras_shape[1:] == (1,1,col)
	avg_pool2 = shared_layer_one(avg_pool2)
	assert avg_pool2._keras_shape[1:] == (1,1,col)
	avg_pool2 = shared_layer_two(avg_pool2)
	assert avg_pool2._keras_shape[1:] == (1,1,col)

	max_pool2 = GlobalMaxPooling2D()(scan_feature)
	max_pool2 = Reshape((1,1,col))(max_pool2)
	assert max_pool2._keras_shape[1:] == (1,1,col)
	max_pool2 = shared_layer_one(max_pool2)
	assert max_pool2._keras_shape[1:] == (1,1,col)
	max_pool2 = shared_layer_two(max_pool2)
	assert max_pool2._keras_shape[1:] == (1,1,col)

	scan_feature2 = Add()([avg_pool2,max_pool2])
	scan_feature2 = Activation('sigmoid')(scan_feature2)


	scan_feature = Permute((2, 3, 1))(scan_feature)
	scan_feature2 = Permute((2, 3, 1))(scan_feature2)
	scan_feature1 = Permute((3, 1, 2))(scan_feature1)
	scan= multiply([scan_feature2, scan_feature1])

	return multiply([scan_feature, scan])

