tf.reset_default_graph()
with tf.Graph().as_default():
    # hyper-params
    learning_rate = 0.0002
    epochs = 250
    batch_size = 16
    N_w = 11 #number of frames concatenated together
    channels = 9*N_w
    drop_out = [0.5, 0.5, 0.5, 0, 0, 0, 0, 0]

    def conv_down(x, N, stride, count): #Conv [4x4, str_2] > Batch_Normalization > Leaky_ReLU
        with tf.variable_scope("conv_down_{}_{}".format(N, count)) as scope: #N == depth of tensor
            with tf.variable_scope("conv_down_4x4_str{}".format(stride)) : #this's used for downsampling
                x = tf.layers.conv2d(x, N, kernel_size=4, strides=stride, padding='same', kernel_initializer=tf.truncated_normal_initializer(stddev=np.sqrt(0.2)), name=scope)
                x = tf.contrib.layers.batch_norm(x)
                x = tf.nn.leaky_relu(x) #for conv_down, implement leakyReLU
        return x

    def conv_up(x, N, drop_rate, stride, count): #Conv_transpose [4x4, str_2] > Batch_Normalizaiton > DropOut > ReLU
        with tf.variable_scope("{}".format(count)) as scope:
            x = tf.layers.conv2d_transpose(x, N, kernel_size=4, strides=stride, padding='same', kernel_initializer=tf.truncated_normal_initializer(stddev=np.sqrt(0.2)), name=scope)
            x = tf.contrib.layers.batch_norm(x)
            if drop_rate is not 0:
                x = tf.nn.dropout(x, keep_prob=drop_rate)
            x = tf.nn.relu(x)
        return x

    def conv_refine(x, N, drop_rate): #Conv [3x3, str_1] > Batch_Normalization > DropOut > ReLU
        x = tf.layers.conv2d(x, N, kernel_size=3, strides=1, padding='same', kernel_initializer=tf.truncated_normal_initializer(stddev=np.sqrt(0.2)))
        x = tf.contrib.layers.batch_norm(x)
        if drop_rate is not 0:
            x = tf.nn.dropout(x, keep_prob=drop_rate)
        x = tf.nn.relu(x)
        return x

    def conv_upsample(x, N, drop_rate, stride, count):
        with tf.variable_scope("conv_upsamp_{}_{}".format(N,count)) :
            with tf.variable_scope("conv_up_{}".format(count)):
                x = conv_up(x, 2*N, drop_rate, stride,count)
            with tf.variable_scope("refine1"):
                x = conv_refine(x, N, drop_rate)
            with tf.variable_scope("refine2"):
                x = conv_refine(x, N, drop_rate)
        return x

    def biLinearDown(x, N):
        return tf.image.resize_images(x, [N, N])

    def finalTanH(x):
        return tf.nn.tanh(x)

    def T(x):
        #channel_output_structure
        down_channel_output = [64, 128, 256, 512, 512, 512, 512, 512]
        up_channel_output= [512, 512, 512, 512, 256, 128, 64, 3]
        biLinearDown_output= [32, 64, 128] #for skip-connection

        #down_sampling
        conv1 = conv_down(x, down_channel_output[0], 2, 1)
        conv2 = conv_down(conv1, down_channel_output[1], 2, 2)
        conv3 = conv_down(conv2, down_channel_output[2], 2, 3)
        conv4 = conv_down(conv3, down_channel_output[3], 1, 4)
        conv5 = conv_down(conv4, down_channel_output[4], 1, 5)
        conv6 = conv_down(conv5, down_channel_output[5], 1, 6)
        conv7 = conv_down(conv6, down_channel_output[6], 1, 7)
        conv8 = conv_down(conv7, down_channel_output[7], 1, 8)

        #upsampling
        dconv1 = conv_upsample(conv8, up_channel_output[0], drop_out[0], 1, 1)
        dconv2 = conv_upsample(dconv1, up_channel_output[1], drop_out[1], 1, 2)
        dconv3 = conv_upsample(dconv2, up_channel_output[2], drop_out[2], 1, 3)
        dconv4 = conv_upsample(dconv3, up_channel_output[3], drop_out[3], 1, 4)
        dconv5 = conv_upsample(dconv4, up_channel_output[4], drop_out[4], 1, 5)
        dconv6 = conv_upsample(tf.concat([dconv5, biLinearDown(x, biLinearDown_output[0])], axis=3), up_channel_output[5], drop_out[5], 2, 6)
        dconv7 = conv_upsample(tf.concat([dconv6, biLinearDown(x, biLinearDown_output[1])], axis=3), up_channel_output[6], drop_out[6], 2, 7)
        dconv8 = conv_upsample(tf.concat([dconv7, biLinearDown(x, biLinearDown_output[2])], axis=3), up_channel_output[7], drop_out[7], 2, 8)

        #final_tanh
        T_x = finalTanH(dconv8)

        return T_x

    # input_tensor X
    x = tf.placeholder(tf.float32, [batch_size, 256, 256, channels]) # batch_size x Height x Width x N_w

    # define sheudo_input for testing
    sheudo_input = np.float32(np.random.uniform(low=-1., high=1., size=[16, 256,256, 99]))

    # initialize_
    init_g = tf.global_variables_initializer()
    init_l = tf.local_variables_initializer()
    with tf.Session() as sess:
        sess.run(init_g)
        sess.run(init_l)
        sess.run(T(x),  feed_dict={x: sheudo_input})