from keras import activations, initializers, constraints
from keras import regularizers
from keras.engine import Layer
import keras.backend as K
class GCN(Layer):
	def __init__(self,units,support=1,activation=None,use_bias=True,
				kernel_initializer='glorot_uniform',bias_initializer='zeros',
				kernel_regularizer=None,bias_regularizer=None,
				activity_regularizer=None,kernel_constraint=None,
				bias_constraint=None,**kwargs):
		if 'input_shape' not in kwargs  and 'input_dim' in kwargs:
			kwargs['input_shape']=(kwargs.pop('input_dim'),)
		super(GCN,self).__init__(**kwargs)
		self.units=units
		self.activation=activations.get(activation)
		self.use_bias=use_bias
		self.kernel_initializer=initializers.get(kernel_initializer)
		self.kernel_regularizer=regularizers.get(kernel_regularizer)
		self.kernel_constraint=constraints.get(kernel_constraint)
		self.bias_initializer=initializers.get(bias_initializer)
		self.bias_regularizer=regularizers.get(bias_regularizer)
		self.bias_constraint=constraints.get(bias_constraint)
		self.activity_regularizer=regularizers.get(activity_regularizer)
		self.support_masking=True
		self.support=support #support should always be greater than 1


	def compute_output_shape(self,input_shapes):
		feature_shape=input_shapes[0]
		output_shape=(feature_shape[0],self.units)
		return output_shape

	def build(self,input_shapes):
		features_shape=input_shapes[0]
		input_dim=features_shape[1]

		self.kernel=self.add_weight(shape=(input_dim*self.support,self.units),
									initializer=self.kernel_initializer,
									name='kernel',
									regularizer=self.kernel_regularizer,
									constraint=self.kernel_constraint)
		if self.use_bias:
			self.bias=self.add_weight(shape=(self.units,),
									 initializer=self.bias_initializer,
									 name='bias',
									 regularizer=self.bias_regularizer,
									 constraint=self.bias_constraint)
		else:
			self.bias=None
		self.built=True

	def call(self,inputs,mask=None):
		features=inputs[0]
		basis=inputs[1:]
		supports=list()
		for i in range(self.support):
			supports.append(K.dot(basis[i],features))
		supports=K.concatenate(supports,axis=1)
		output=K.dot(supports,self.kernel)
		if self.bias:
			output+=self.bias

		return self.activation(output)

	def get_config(self):
		config={'units':self.units,
				'support':self.support,
				'activation':activations.serialize(self.activation),
				'use_bias':self.use_bias,
				'kernel_initializer':initializers.serialize(self.kernel_initializer),
				'kernel_constraint':constraints.serialize(self.kernel_constraint),
				'kernel_regularizer':regularizers.serialize(self.kernel_regularizer),
				'bias_initializer':initializers.serialize(self.bias_initializer),
				'bias_constraint':constraints.serialize(self.bias_constraint),
				'bias_regularizer':regularizers.serialize(self.bias_regularizer),
				'activity_regularizer':regularizers.serialize(self.activity_regularizer)

		}
		base_config=super(GCN,self).get_config()
		return dict(list(base_config.items())+list(config.items()))