import tensorflow as tf 
import numpy as np 

# Choose the number of classes
num_classes = 3
image_shape = (576,600)

# # Load the previously saved meta graph in the default graph
saver = tf.train.import_meta_graph('./models/model.ckpt.meta')
#saver = tf.train.import_meta_graph('../saved/model4.meta')

graph = tf.get_default_graph()

image_input = graph.get_tensor_by_name("Placeholder:0")

what_is_this = graph.get_tensor_by_name("Placeholder_1:0")

trained_output = graph.get_tensor_by_name("cond/DeepLab_v3/logits/Conv2D:0")

final_conv_layer = tf.layers.conv2d(trained_output,num_classes,1,
                kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001),
                kernel_initializer=tf.truncated_normal_initializer(mean=0, stddev=0.01),name='final_conv_layer')
final_layer = net = tf.image.resize_bilinear(final_conv_layer, image_shape,name='final_layer')


print(image_input.shape)
print(what_is_this.shape)
print(trained_output.shape)
print("Shape of final Conv layer: {}".format(final_conv_layer.shape))
print("Shape of final Output: {}".format(final_layer.shape))

t = tf.constant([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]], name='test_name')
print(t.name)  # [2, 2, 3]

with tf.Session() as sess:

    sess.run(tf.global_variables_initializer())
    saver.restore(sess, save_path="./models/model.ckpt")
    #saver.restore(sess, save_path="../saved/model4")

    # # Look at the parameters and the names
    # resnet_params = [(v.name, v.get_shape().as_list())
    #     for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)]
    #
    # print("Resnet Parameters")
    # for p in resnet_params:
    #     print("Variable Name {0}: number of params = {1}".format(p[0], p[1]))
    # print()

    op = sess.graph.get_operations()
    op_list = [m.name for m in op]

    for operation in op_list:
        print(operation)