import tensorflow as tf

from functools import reduce

class Convolution3DNetwork(object):

    DEFAULT_LAYER_PADDING = 'VALID'

    DEFAULT_CONV_STRIDE = [1, 1, 1, 1, 1]

    def __init__(self, config):

        self._config = config

        self._strides = self._config.get_strides()

        self._pool_strides = self._config.get_pool_strides()

        self._pool_windows = self._config.get_pool_windows()

        self._init_weights()

        self._init_biases()

    def _init_weights(self):

        self._weights = [

            tf.Variable(init_func, name=name) 

            for name, init_func in self._config.get_fc_weights()

            ]

        self._conv_weights = [

                tf.Variable(init_func, name=name) 

                for name, init_func in self._config.get_conv_weights()

            ]

    def _init_biases(self):

        self._biases = [

                tf.Variable(init_func, name=name) 

                for name, init_func in self._config.get_fc_biases()

            ]

        self._conv_biases = [

                tf.Variable(init_func, name=name) 

                for name, init_func in self._config.get_conv_biases()

            ]

    def weights(self):

        return self._conv_weights + self._weights

    def l2_regularizer(self):

        if self._config.with_l2_norm():

            return reduce(lambda x, y: tf.nn.l2_loss(x) + tf.nn.l2_loss(y),

                          self._weights)

        return 0

    def biases(self):

        return self._conv_biases + self._biases

    # Create some wrappers for simplicity

    def conv3d(self, x, W, b, name, 

               strides=DEFAULT_CONV_STRIDE,

               padding=DEFAULT_LAYER_PADDING):

        with tf.variable_scope(name) as scope:

            # Conv3D wrapper, with bias and relu activation

            x = tf.nn.conv3d(x, W, strides=strides, 

                             padding=padding, name=scope.name)

            x = tf.nn.bias_add(x, b, name='bias')

            return tf.nn.relu(x, name='relu')

    def maxpool3d(self, x, name, k, 

                  strides=DEFAULT_CONV_STRIDE, 

                  padding=DEFAULT_LAYER_PADDING):

        # MaxPool3D wrapper

        return tf.nn.max_pool3d(x, ksize=k, strides=strides, 

                                padding=padding, name=name)

    def fc(self, x, weights, bias, name, dropout=None, with_relu=True):

        with tf.variable_scope(name) as scope:

            fc = tf.add(tf.matmul(x, weights), bias, name=scope.name)

            if with_relu:

                fc = tf.nn.relu(fc, name='relu')

            if dropout:

                fc = tf.nn.dropout(fc, dropout)

            return fc

    # Create model

    def conv_net(self, x, dropout):

        # Convolution Layer

        last_conv_layer = x

        for i, weight in enumerate(self._conv_weights):

             # Convolution Layer

            last_conv_layer = self.conv3d(last_conv_layer, weight, 

                                         self._conv_biases[i],

                                         name="conv" + str(i), 

                                         strides=self._strides[i])

            # Max Pooling (down-sampling)

            if self._pool_windows[i]:

                last_conv_layer = self.maxpool3d(last_conv_layer, 

                                                name="pool" + str(i),

                                                k=self._pool_windows[i],

                                                strides=self._pool_strides[i])

            print("After current layer: ", last_conv_layer.get_shape().as_list())

        if self._config.has_dropout_after_convolutions():

            last_conv_layer = tf.nn.dropout(last_conv_layer, dropout)

        conv_shape = last_conv_layer.get_shape().as_list()

        fully_con_input_size = reduce(lambda x, y: x * y, conv_shape[1:])

        print("SHAPE of the last convolution layer after max pooling: {}, new shape {}".format(

            conv_shape, fully_con_input_size))

        # Fully connected layer

        # Reshape conv output to fit fully connected layer input

        number = conv_shape[0] or -1

        fully_connected = tf.reshape(last_conv_layer, [number, fully_con_input_size])

        for i, weight in enumerate(self._weights[:-1]):

            if self._config.has_fc_dropout(i):

                layer_dropout = dropout

            else:

                layer_dropout = None

            fully_connected = self.fc(fully_connected, 

                                     weight, 

                                     self._biases[i], 

                                     name='fully_connected' + str(i),

                                     dropout=layer_dropout)

        # Output, class prediction

        out = tf.add(tf.matmul(fully_connected, self._weights[-1]),

                     self._biases[-1], name='output_layer')

        return out

def loss_function_with_logits(logits, labels, tensor_name='cost_func'):

    return tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(

        logits=logits, labels=labels), name=tensor_name)

# Sparse sofmtax is used for mutually exclusive classes,

# labels rank must be logits rank - 1

def sparse_loss_with_logits(logits, labels, tensor_name='cost_func'):

    return tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(

        logits=logits, labels=labels), name=tensor_name)