import tensorflow as tf

def fc_layer(layer_name, x, out_nodes):
    '''
    Wrapper for fully connected layers with RELU activation as default
    '''
    shape = x.get_shape()
    if len(shape) == 5: # FC 3D
        size = shape[1].value*shape[2].value*shape[3].value*shape[4].value
    elif len(shape) == 4:
        size = shape[1].value*shape[2].value*shape[3].value
    else:
        size = shape[-1].value

    with tf.variable_scope(layer_name):
        w = tf.get_variable(name='weight',
                            shape=[size, out_nodes],
                            initializer=tf.constant_initializer(0.0))
        b = tf.get_variable(name='bias',
                            shape=[out_nodes],
                            initializer=tf.constant_initializer(0.0))
        # batch?
        flat_x = tf.reshape(x, [-1,size])

        x = tf.nn.bias_add(tf.matmul(flat_x,w), b)
        x = tf.nn.relu(x)
