import tensorflow as tf
import numpy as np
import pickle as pkl
import networkx as nx
import scipy.sparse as sp
from sklearn.gaussian_process.kernels import RBF

_LAYER_UIDS = {}

def zeros(shape, name=None):
    """All zeros."""
    initial = tf.zeros(shape, dtype=tf.float32)
    return tf.Variable(initial, name=name)

def dot(x, y, sparse=False):
    """Wrapper for tf.matmul (sparse vs dense)."""
    res = tf.layers.conv2d(x,y[1],[1 ,y[0]])
    return res[:, :, 0, :]
def get_layer_uid(layer_name=''):
    """Helper function, assigns unique layer IDs."""
    if layer_name not in _LAYER_UIDS:
        _LAYER_UIDS[layer_name] = 1
        return 1
    else:
        _LAYER_UIDS[layer_name] += 1
        return _LAYER_UIDS[layer_name]

class Layer(object):

    def __init__(self, **kwargs):
        allowed_kwargs = {'name'}
        for kwarg in kwargs.keys():
            assert kwarg in allowed_kwargs, 'Invalid keyword argument: ' + kwarg
        name = kwargs.get('name')
        if not name:
            layer = self.__class__.__name__.lower()
            name = name = layer + '_' + str(get_layer_uid(layer))
        self.name = name
        self.vars = {}

class GraphConvolution(Layer):
    """Graph convolution layer."""
    def __init__(self, input, adj_matrix , output_dim, dropout=0.,act=tf.nn.relu, bias=False,
                 **kwargs):
        super(GraphConvolution, self).__init__(**kwargs)

        self.dropout = dropout

        self.act = act
        self.adj_matrix = adj_matrix
        
        self.bias = bias
        self.input = input
        self.output_dim = output_dim

        with tf.variable_scope(self.name + '_vars'):
            if self.bias:
                self.vars['bias'] = zeros([output_dim], name='bias')
                
    def call(self):                                      
        with tf.name_scope(self.name):
             outputs = self._call(self.input)
        return outputs

    def _call(self, inputs):
        #print(inputs)
        x = inputs
        x = tf.nn.dropout(x, self.dropout)
                                          
        # convolve
        x = tf.matmul(self.adj_matrix,x)
        print(x.shape)
        pre_sup = dot(tf.expand_dims(x,-1),[int(self.input.shape[-1]),int(self.output_dim)])
        output = pre_sup

        # bias
        if self.bias:
            output += self.vars['bias']

        return self.act(output)

class Model(object):
    def __init__(self, **kwargs):
        allowed_kwargs = {'name'}
        for kwarg in kwargs.keys():
            assert kwarg in allowed_kwargs, 'Invalid keyword argument: ' + kwarg
        name = kwargs.get('name')
        if not name:
            name = self.__class__.__name__.lower()
        self.name = name
    
        self.vars = {}
        self.placeholders = {}


        self.layers = []
        self.activations = []

        self.inputs = None
        self.outputs = None

    def build(self):
        raise NotImplementedError

class GCN(Model):
    def __init__(self, x, adj_matrix,output_dim, dropout = 0.5,**kwargs):
        super(GCN, self).__init__(**kwargs)

        self.input = x
        self.adj_matrix = adj_matrix
        self.dropout = dropout
        self.output_dim = output_dim
        
    def build(self):
        if len(self.output_dim) == 0:
            return self.input

        outputs = GraphConvolution(input = self.input,
                        adj_matrix = self.adj_matrix,
                        output_dim=self.output_dim[0],
                        act=tf.nn.relu,
                        dropout=self.dropout).call()

        for i in range(1,len(self.output_dim)):
            outputs = GraphConvolution(input=outputs,
                            adj_matrix = self.adj_matrix,
                            output_dim=self.output_dim[i],
                            act=tf.nn.relu,
                            dropout=self.dropout).call()
        
        return outputs




