import tensorflow as tf

def conv2d(x, W, b, strides=1):
    # Conv2D wrapper, with bias and relu activation
    x = tf.nn.conv2d(x, W, strides=[1, strides, strides, 1], padding='SAME')
    x = tf.nn.bias_add(x, b)
    return tf.nn.relu(x)

def maxpool2d(x, k=2):
    return tf.nn.max_pool(x, ksize=[1, k, k, 1], strides=[1, k, k, 1],padding='SAME')

def conv2d_transpose(x, W, b, output, strides=2):
    # Conv2D wrapper, with bias and relu activation
    x = tf.nn.conv2d_transpose(x, W, output_shape= output, strides=[1, strides, strides, 1], padding='SAME')
    x = tf.nn.bias_add(x, b)
    return tf.nn.relu(x)

def make_wts_and_bias(weights, biases, input_channels, output_channels, type):
    curr_weights_num = len(weights)
    curr_biases_num = len(biases)
    if type == 'normal':
        new_w = 'wc' + str(curr_weights_num + 1)
        weights[new_w] = tf.get_variable('W' + str(curr_weights_num + 1), shape=(3, 3, input_channels, output_channels), initializer=tf.contrib.layers.xavier_initializer())
        new_b = 'b' + str(curr_biases_num + 1)
        biases[new_b] = tf.get_variable('B' + str(curr_biases_num + 1), shape=(output_channels), initializer=tf.contrib.layers.xavier_initializer())
    elif type == 'transpose':
        new_w = 'wuc' + str(curr_weights_num + 1)
        weights[new_w] = tf.get_variable('W' + str(curr_weights_num + 1), shape=(3, 3, output_channels, input_channels), initializer=tf.contrib.layers.xavier_initializer())
        new_b = 'ub' + str(curr_biases_num + 1)
        biases[new_b] = tf.get_variable('B' + str(curr_biases_num + 1), shape=(output_channels), initializer=tf.contrib.layers.xavier_initializer())
    return weights, biases, new_w, new_b

def residual_block(x, weights, biases, output_channels):
    input_channels =  x.shape[3]
    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, input_channels, output_channels, 'normal')
    y1 = conv2d(x, weights[new_w], biases[new_b])

    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, output_channels, output_channels, 'normal')
    y2 = conv2d(y1, weights[new_w], biases[new_b])

    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, output_channels, output_channels, 'normal')
    y3 = conv2d(y2, weights[new_w], biases[new_b])

    y3 = y3 + x
    return y1, y2, y3, weights, biases


def conv_net(x):
    weights = {}
    biases = {}
    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, 1, 32, 'normal')
    conv0 = conv2d(x, weights[new_w], biases[new_b])

    # block1
    conv1, conv2, conv3, weights, biases = residual_block(conv0, weights, biases, 32)
    mp1 = maxpool2d(conv3, 2)

    # block2
    conv4, conv5, conv6, weights, biases = residual_block(mp1, weights, biases, 32)
    mp2 = maxpool2d(conv6, 2)

    # block3
    conv7, conv8, conv9, weights, biases = residual_block(mp2, weights, biases, 32)
    mp3 = maxpool2d(conv9, 2)

    # block4
    conv10, conv11, conv12, weights, biases = residual_block(mp3, weights, biases, 32)
    mp4 = maxpool2d(conv12, 2)

    # block5
    conv13, conv14, conv15, weights, biases = residual_block(mp4, weights, biases, 32)

    # sideblock1
    conv16, conv17, conv18, weights, biases = residual_block(conv3, weights, biases, 32)

    # sideblock2
    conv19, conv20, conv21, weights, biases = residual_block(conv6, weights, biases, 32)

    # sideblock3
    conv22, conv23, conv24, weights, biases = residual_block(conv9, weights, biases, 32)

    # sideblock4
    conv25, conv26, conv27, weights, biases = residual_block(conv12, weights, biases, 32)

    # up1
    mp5 = maxpool2d(conv27, 2)
    conv15 = conv15 + mp5
    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, 32, 32, 'transpose')
    conv28 = conv2d_transpose(conv15, weights[new_w], biases[new_b], [tf.shape(conv12)[0], 108, 82, 32], strides=2)
    conv29, conv30, conv31, weights, biases = residual_block(conv28, weights, biases, 32)

    # up2
    mp6 = maxpool2d(conv24, 2)
    conv31 = conv31 + mp6
    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, 32, 32, 'transpose')
    conv32 = conv2d_transpose(conv31, weights[new_w], biases[new_b], [tf.shape(conv12)[0], 215, 163, 32], strides=2)
    conv33, conv34, conv35, weights, biases = residual_block(conv32, weights, biases, 32)

    # up3
    mp7 = maxpool2d(conv21, 2)
    conv35 = conv35 + mp7
    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, 32, 32, 'transpose')
    conv36 = conv2d_transpose(conv35, weights[new_w], biases[new_b], [tf.shape(conv12)[0], 430, 325, 32], strides=2)
    conv37, conv38, conv39, weights, biases = residual_block(conv36, weights, biases, 32)

    # up4
    mp8 = maxpool2d(conv18, 2)
    conv39 = conv39 + mp8
    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, 32, 32, 'transpose')
    conv40 = conv2d_transpose(conv39, weights[new_w], biases[new_b], [tf.shape(conv12)[0], 860, 650, 32], strides=2)
    conv41, conv42, conv43, weights, biases = residual_block(conv40, weights, biases, 32)

    weights, biases, new_w, new_b = make_wts_and_bias(weights, biases, 32, 1, 'normal')
    conv44 = conv2d(conv43, weights[new_w], biases[new_b])
    out = conv44
    return out