import tensorflow as tf
import numpy as np

dim = 256
kernel_dim = 3
dilation_rate = np.array([2, 2])
input_img_np = np.random.random((1, dim, dim, 1)).astype(np.float32)
kernel =  np.random.random((kernel_dim,kernel_dim,1,1)).astype(np.float32)

with tf.Session() as sess:
    concrete_input_op = tf.constant(input_img_np)
    concrete_output_op = tf.nn.convolution(concrete_input_op, kernel, padding='SAME', dilation_rate=dilation_rate)
    concrete_output = sess.run(concrete_output_op)
    
    print('convolution + CONCRETE + SAME')
    print('concrete_input_op: ', concrete_input_op.get_shape())
    print('concrete_output_op: ', concrete_output_op.get_shape())
    print('concrete_output:', concrete_output.shape)
    assert(concrete_input_op.get_shape() == concrete_output_op.get_shape())


    undef_input_op = tf.placeholder(tf.float32, shape=(None, dim, dim, 1))
    undef_output_op = tf.nn.convolution(undef_input_op, kernel, padding='SAME', dilation_rate=dilation_rate)
    undef_output = sess.run(undef_output_op, feed_dict={undef_input_op: input_img_np})
    
    print('convolution + UNDEF + SAME')
    print('undef_input_op: ', undef_input_op.get_shape())
    print('undef_output_op: ', undef_output_op.get_shape())
    print('undef_output:', undef_output.shape)
    # This assert will correctly fail even though the shapes are ok because shapes are only partially known
    # assert(undef_input_op.get_shape() == undef_output_op.get_shape())

    valid_concrete_input_op = tf.constant(input_img_np)
    valid_concrete_output_op = tf.nn.convolution(valid_concrete_input_op, kernel, padding='VALID', dilation_rate=dilation_rate)
    valid_concrete_output = sess.run(valid_concrete_output_op)
    
    print('convolution + CONCRETE + VALID')
    print('valid_concrete_input_op: ', valid_concrete_input_op.get_shape())
    print('valid_concrete_output_op: ', valid_concrete_output_op.get_shape())
    print('valid_concrete_output:', valid_concrete_output.shape)


    valid_undef_input_op = tf.placeholder(tf.float32, shape=(None, dim, dim, 1))
    valid_undef_output_op = tf.nn.convolution(valid_undef_input_op, kernel, padding='VALID', dilation_rate=dilation_rate)
    valid_undef_output = sess.run(valid_undef_output_op, feed_dict={valid_undef_input_op: input_img_np})
    
    print('convolution + UNDEF + VALID')
    print('valid_undef_input_op: ',  valid_undef_input_op.get_shape())
    print('valid_undef_output_op: ', valid_undef_output_op.get_shape())
    print('valid_undef_output:', valid_undef_output.shape)
    # This assert will correctly fail even though the shapes are ok because shapes are only partially known
    # assert(undef_input_op.get_shape() == undef_output_op.get_shape())
    ############################################################################
    # Now atrous
    concrete_input_op = tf.constant(input_img_np)
    concrete_output_op = tf.nn.atrous_conv2d(concrete_input_op, kernel, padding='SAME', rate=2)
    concrete_output = sess.run(concrete_output_op)
    
    print('atrous_conv2d + CONCRETE + SAME')
    print('concrete_input_op: ', concrete_input_op.get_shape())
    print('concrete_output_op: ', concrete_output_op.get_shape())
    print('concrete_output_op: ', concrete_output_op.get_shape())
    print('concrete_output:', concrete_output.shape)
    assert(concrete_input_op.get_shape() == concrete_output_op.get_shape())


    undef_input_op = tf.placeholder(tf.float32, shape=(None, dim, dim, 1))
    undef_output_op = tf.nn.atrous_conv2d(undef_input_op, kernel, padding='SAME', rate=2)
    undef_output = sess.run(undef_output_op, feed_dict={undef_input_op: input_img_np})
    
    print('atrous_conv2d + UNDEF + SAME')
    print('undef_input_op: ', undef_input_op.get_shape())
    print('undef_output_op: ', undef_output_op.get_shape())
    print('undef_output:', undef_output.shape)
    # This assert will correctly fail even though the shapes are ok because shapes are only partially known
    # assert(undef_input_op.get_shape() == undef_output_op.get_shape())

    valid_concrete_input_op = tf.constant(input_img_np)
    valid_concrete_output_op = tf.nn.atrous_conv2d(valid_concrete_input_op, kernel, padding='VALID', rate=2)
    valid_concrete_output = sess.run(valid_concrete_output_op)
    
    print('atrous_conv2d + CONCRETE + VALID')
    print('valid_concrete_input_op: ', valid_concrete_input_op.get_shape())
    print('valid_concrete_output_op: ', valid_concrete_output_op.get_shape())
    print('valid_concrete_output:', valid_concrete_output.shape)


    valid_undef_input_op = tf.placeholder(tf.float32, shape=(None, dim, dim, 1))
    valid_undef_output_op = tf.nn.atrous_conv2d(valid_undef_input_op, kernel, padding='VALID', rate=2)
    valid_undef_output = sess.run(valid_undef_output_op, feed_dict={valid_undef_input_op: input_img_np})
    
    print('atrous_conv2d + UNDEF + VALID')
    print('valid_undef_input_op: ',  valid_undef_input_op.get_shape())
    print('valid_undef_output_op: ', valid_undef_output_op.get_shape())
    print('valid_undef_output:', valid_undef_output.shape)
    # This assert will correctly fail even though the shapes are ok because shapes are only partially known
    # assert(undef_input_op.get_shape() == undef_output_op.get_shape())