import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import os

n_iterations  = 1000
#fuction for convo layers
def convo(X,inp_channel,out_channel):
    #with tf.variable_scope(scope or "linear"):
    W1 = tf.Variable(tf.random_normal([5, 5,inp_channel, 8], stddev=0.1), name="W1")
    B1 = tf.Variable(tf.random_normal([8], stddev=0.1), name="B1")
    conv1=tf.nn.conv2d(X,W1,strides=[1,1,1,1],padding='SAME')+B1


    conv1 = tf.nn.relu(conv1)
    W2 = tf.Variable(tf.random_normal([5, 5,8, out_channel], stddev=0.1), name="W2")
    B2 = tf.Variable(tf.random_normal([out_channel], stddev=0.1), name="B2")
    conv2=tf.nn.conv2d(conv1,W2,strides=[1,1,1,1],padding='SAME')+B2

    conv2 = tf.nn.relu(conv2)
    return conv2


files = [os.path.join('abc.jpg')]
data=plt.imread(files[0]).astype(np.float32)/255
img=[data.tolist()]

files2 = [os.path.join('blur2.jpeg')]
data2=plt.imread(files2[0]).astype(np.float32)/255
img2=[data2.tolist()]

inp_img= tf.placeholder(tf.float32, shape=[None,804,1920,3], name='X')
out_img = tf.placeholder(tf.float32, shape=[None,804,1920,3], name='Y')


n_channels=[3,3,3,3]

Previous_out=inp_img
#initiating convo layers
for conv_i in range(1,len(n_channels)):

    current_i=convo(X=Previous_out,inp_channel= n_channels[conv_i-1],out_channel= n_channels[conv_i])
    Previous_out= tf.reshape(Previous_out,[1,804,1920,n_channels[conv_i]])
    Previous_out=current_i+Previous_out
Final_out= Previous_out




#cost_function
cost = current_i+inp_img - out_img
loss = tf.reduce_sum(tf.square(cost))
optimizer = tf.train.AdamOptimizer(0.001).minimize(loss)
sess = tf.InteractiveSession()
init = tf.global_variables_initializer().run()

for it_i in range(n_iterations):

    optimizer.run(feed_dict={inp_img:img,out_img:img2})
    if it_i%20 == 0:
        train_accuracy=loss.eval(feed_dict={int_img:img,out_img:img2})
        print(it_i,train_accuracy)
        final_img=Final_out.eval(feed_dict={int_img:img,out_img:img2})
        final_img1=np.squeeze(np.array(final_img))
        plt.imshow(final_img1)
        plt.show()