import tensorflow as tf

def conv(x, filter_height, filter_width, num_filters, stride_y, stride_x, name,
         padding='SAME', groups=1):
    """Create a convolution layer.
    Adapted from: https://github.com/ethereon/caffe-tensorflow
    """
    # Get number of input channels
    input_channels = int(x.get_shape()[-1])

    # Create lambda function for the convolution
    convolve = lambda i, k: tf.nn.conv2d(i, k,
                                         strides=[1, stride_y, stride_x, 1],
                                         padding=padding)

    with tf.variable_scope(name) as scope:
        # Create tf variables for the weights and biases of the conv layer
        weights = tf.get_variable('weights', shape=[filter_height,
                                                    filter_width,
                                                    input_channels/groups,
                                                    num_filters])
        biases = tf.get_variable('biases', shape=[num_filters])

    if groups == 1:
        conv = convolve(x, weights)

    # In the cases of multiple groups, split inputs & weights and
    else:
        # Split input and weights and convolve them separately
        input_groups = tf.split(axis=3, num_or_size_splits=groups, value=x)
        weight_groups = tf.split(axis=3, num_or_size_splits=groups,
                                 value=weights)
        output_groups = [convolve(i, k) for i, k in zip(input_groups, weight_groups)]

        # Concat the convolved output together again
        conv = tf.concat(axis=3, values=output_groups)

    # Add biases
    bias = tf.reshape(tf.nn.bias_add(conv, biases), tf.shape(conv))

    # Apply relu function
    relu = tf.nn.relu(bias, name=scope.name)

    return relu
def fc(x, num_in, num_out, name, relu=True):
    """Create a fully connected layer."""
    with tf.variable_scope(name) as scope:

        # Create tf variables for the weights and biases
        weights = tf.get_variable('weights', shape=[num_in, num_out],
                                  trainable=True)
        biases = tf.get_variable('biases', [num_out], trainable=True)

        # Matrix multiply weights and inputs and add bias
        act = tf.nn.xw_plus_b(x, weights, biases, name=scope.name)

    if relu:
        # Apply ReLu non linearity
        relu = tf.nn.relu(act)
        return relu
    else:
        return act
def max_pool(x, filter_height, filter_width, stride_y, stride_x, name,
             padding='SAME'):
    """Create a max pooling layer."""
    return tf.nn.max_pool(x, ksize=[1, filter_height, filter_width, 1],
                          strides=[1, stride_y, stride_x, 1],
                          padding=padding, name=name)
def lrn(x, radius, alpha, beta, name, bias=1.0):
    """Create a local response normalization layer."""
    return tf.nn.local_response_normalization(x, depth_radius=radius,
                                              alpha=alpha, beta=beta,
                                              bias=bias, name=name)
def dropout(x, keep_prob):
    """Create a dropout layer."""
    return tf.nn.dropout(x, keep_prob)

def pretrained_classifier(inputae, n_label, reuse=False,  name='classifier', isTrain = False, verbose=False):
    if verbose:
        print("Classifier", isTrain)
    if isTrain == True:
        DROPOUT_KEEP_PROB = 0.5
    else:
        DROPOUT_KEEP_PROB = 1.0
    with tf.variable_scope(name) as scope:
        if reuse:
            tf.get_variable_scope().reuse_variables()
        else:
            assert tf.get_variable_scope().reuse is False
        """Create the network graph."""
        # 1st Layer: Conv (w ReLu) -> Lrn -> Pool
        if verbose:
            print('input: ', inputae)
        conv1 = conv(inputae, 11, 11, 96, 4, 4, padding='VALID', name='conv1')
        if verbose:
            print('conv1: ', conv1)
        norm1 = lrn(conv1, 2, 2e-05, 0.75, name='norm1')
        pool1 = max_pool(norm1, 3, 3, 2, 2, padding='VALID', name='pool1')
        if verbose:
            print('pool1: ', pool1)
        # 2nd Layer: Conv (w ReLu)  -> Lrn -> Pool with 2 groups
        conv2 = conv(pool1, 5, 5, 256, 1, 1, groups=2, name='conv2')
        if verbose:
            print('conv2: ', conv2)
        norm2 = lrn(conv2, 2, 2e-05, 0.75, name='norm2')
        pool2 = max_pool(norm2, 3, 3, 2, 2, padding='VALID', name='pool2')
        if verbose:
            print('pool2: ', pool2)
        # 3rd Layer: Conv (w ReLu)
        conv3 = conv(pool2, 3, 3, 384, 1, 1, name='conv3')
        if verbose:
            print('conv3: ', conv3)
        # 4th Layer: Conv (w ReLu) splitted into two groups
        conv4 = conv(conv3, 3, 3, 384, 1, 1, groups=2, name='conv4')
        if verbose:
            print('conv4: ', conv4)
        # 5th Layer: Conv (w ReLu) -> Pool splitted into two groups
        conv5 = conv(conv4, 3, 3, 256, 1, 1, groups=2, name='conv5')
        if verbose:
            print('conv5: ', conv5)
        pool5 = max_pool(conv5, 3, 3, 1, 1, padding='VALID', name='pool5')
        if verbose:
            print('pool5: ', pool5)
        # 6th Layer: Flatten -> FC (w ReLu) -> Dropout
        flattened = tf.reshape(pool5, [-1, 4*4*256])
        if verbose:
            print('flattened: ', flattened)
        fc6 = fc(flattened, 4*4*256, 4096, name='fc6')
        dropout6 = dropout(fc6, DROPOUT_KEEP_PROB)
        if verbose:
            print('fc6: ', fc6)
        # 7th Layer: FC (w ReLu) -> Dropout
        fc7 = fc(dropout6, 4096, 4096, name='fc7')
        dropout7 = dropout(fc7, DROPOUT_KEEP_PROB)
        if verbose:
            print('fc7: ', fc7)
        # 8th Layer: FC and return unscaled activations
        pred1 = fc(dropout7, 4096, n_label, relu=False, name='fc8')
        #cnn_output = tf.nn.softmax(fc_layer_3)
        #print("logit: ", pred)
        if isTrain == False:
            if verbose:
                print(isTrain)
            pred = tf.stop_gradient(pred1) 
        #prediction = tf.nn.sigmoid(pred1)
        #pred_y = tf.argmax(prediction, 1)
        return pred1