#!/usr/bin/env python
# coding: utf-8

# In[1]:


import tensorflow as tf
import imageio
import numpy as np
import os


# In[2]:


CROP_SIZE = 32
DOWN_SIZE = 16
CHANNEL = 3
BATCH_SIZE = 128


# In[3]:


TRAINING_FOLDER = "./291"
TEST_FOLDER = "./Set5"


# In[4]:


image_list = os.listdir(TEST_FOLDER)
test_training_set = []
test_image_shape = []
test_test_set = []
test_file_name = []
for img in image_list :
    try :
        image = imageio.imread(os.path.join(TEST_FOLDER, img))
        shape = image.shape
        test_image_shape.append(shape)
        test_file_name.append(img)
        image = tf.image.rgb_to_grayscale(image)

        downsize_image = tf.image.resize_images(image, [int(shape[0]/2), int(shape[1]/2)])
        restored_image = tf.image.resize_images(downsize_image, [shape[0], shape[1]])

        # restored_image = tf.Session().run(restored_image) / 255
        # image = tf.Session().run(restored_image) / 255
        restored_image = tf.Session().run(restored_image / 255)
        image = tf.Session().run(image / 255)

        test_training_set.append(restored_image.tolist())
        test_test_set.append(image.tolist())

    except Exception as e :
        print(e)


# In[6]:


X = tf.placeholder(tf.float32, [None, None, None, None])
Y = tf.placeholder(tf.float32, [None, None, None, None])


# In[7]:


W1 = tf.get_variable("W1", shape=[3,3,2,32], initializer=tf.contrib.layers.xavier_initializer())
W2 = tf.get_variable("W2", shape=[3,3,32,32], initializer=tf.contrib.layers.xavier_initializer())
W3 = tf.get_variable("W3", shape=[3,3,32,1], initializer=tf.contrib.layers.xavier_initializer())


# In[8]:


b1 = tf.Variable(tf.zeros([32]))
b2 = tf.Variable(tf.zeros([32]))
b3 = tf.Variable(tf.zeros([1]))


# In[9]:


Wh1 = tf.get_variable("Wh1", shape=[3,3,32,32], initializer=tf.contrib.layers.xavier_initializer())
Wh2 = tf.get_variable("Wh2", shape=[3,3,32,32], initializer=tf.contrib.layers.xavier_initializer())
Wz = tf.get_variable("Wz", shape=[3,3,32,32], initializer=tf.contrib.layers.xavier_initializer())


# In[10]:


tf.summary.histogram("W1", W1)
tf.summary.histogram("W2", W2)
tf.summary.histogram("W3", W3)
tf.summary.histogram("Wh1", Wh1)
tf.summary.histogram("Wh2", Wh2)
tf.summary.histogram("Wz", Wz)


# In[11]:

layer11 = tf.nn.conv2d(tf.concat([X, X], 3), W1, strides=[1,1,1,1], padding="SAME")
h0 = tf.nn.tanh(tf.nn.conv2d(layer11, Wz, strides=[1,1,1,1], padding="SAME"))
layer12 = tf.nn.relu(tf.nn.conv2d(h0, Wh2, strides=[1,1,1,1], padding="SAME"))
layer13 = tf.nn.conv2d(layer12, W3, strides=[1,1,1,1], padding="SAME")

layer21 = tf.nn.conv2d(tf.concat([X, layer13], 3), W1, strides=[1,1,1,1], padding="SAME")
h1 = tf.nn.tanh(tf.nn.conv2d(h0, Wh1, strides=[1,1,1,1], padding="SAME") + tf.nn.conv2d(layer21, Wz, strides=[1,1,1,1], padding="SAME"))
layer22 = tf.nn.relu(tf.nn.conv2d(h1, Wh2, strides=[1,1,1,1], padding="SAME"))
layer23 = tf.nn.conv2d(layer22, W3, strides=[1,1,1,1], padding="SAME")

layer31 = tf.nn.conv2d(tf.concat([X, layer23], 3), W1, strides=[1,1,1,1], padding="SAME")
h2 = tf.nn.tanh(tf.nn.conv2d(h1, Wh1, strides=[1,1,1,1], padding="SAME") + tf.nn.conv2d(layer31, Wz, strides=[1,1,1,1], padding="SAME"))
layer32 = tf.nn.relu(tf.nn.conv2d(h2, Wh2, strides=[1,1,1,1], padding="SAME"))
layer33 = tf.nn.conv2d(layer32, W3, strides=[1,1,1,1], padding="SAME")
hypothesis = layer33


# In[12]:


with tf.name_scope("Cost") :
    # cost = -tf.reduce_mean(Y * tf.log(hypothesis) + (1 - Y) * tf.log(1 - hypothesis))
    # cost = tf.reduce_mean(tf.square(hypothesis - Y))
    cost = tf.reduce_sum(tf.square(hypothesis - Y))
    tf.summary.scalar("Cost", cost)


# In[13]:


with tf.name_scope("Train") :
    optimizer = tf.train.AdamOptimizer(learning_rate=0.0001, beta1=0.9, beta2=0.999, epsilon=1e-08).minimize(cost)


# In[14]:


# accuracy = tf.reduce_mean(tf.cast(tf.equal(hypothesis, Y), dtype=tf.float32))
accuracy = tf.reduce_mean(tf.image.psnr(hypothesis, Y, max_val=1.0))
tf.summary.scalar("Accuracy", accuracy)


saver = tf.train.Saver()

# In[15]:


with tf.Session() as sess:

    # First let's load meta graph and restore weights
    saver.restore(sess, "srcnn")

    x_data = []
    y_data = []
    for idx_ in range(len(test_training_set)):
        shape = test_image_shape[idx_]
        widx, hidx = [np.random.randint(shape[0] - CROP_SIZE - 1, size=1)[0],
                      np.random.randint(shape[1] - CROP_SIZE - 1, size=1)[0]]
        cropped_image = []
        test_cropped_image = []
        for idx in range(widx, widx + CROP_SIZE):
            cropped_image.append(test_training_set[idx_][idx][hidx:hidx + CROP_SIZE])
            test_cropped_image.append(test_test_set[idx_][idx][hidx:hidx + CROP_SIZE])
        x_data.append(cropped_image)
        y_data.append(test_cropped_image)

    for idx in range(len(test_file_name)):
        h, a, c = sess.run(
            [hypothesis, accuracy, cost], feed_dict={X: [test_training_set[idx]], Y: [test_test_set[idx]]}
        )
        print("<Image {}> accuracy {}".format(test_file_name[idx], a))
        imageio.imwrite("result/" + "train" + test_file_name[idx], np.array(test_training_set[idx]))
        imageio.imwrite("result/" + "test" + test_file_name[idx], np.array(test_test_set[idx]))
        imageio.imwrite("result/" + "result" + test_file_name[idx],
                        hypothesis.eval(({X: [test_training_set[idx]], Y: [test_test_set[idx]]}))[0])


# In[ ]:




