import tensorflow as tf
import tensorflow.contrib as U
import os
import numpy as np
import preprocessing as pp
import module as mm
import glob
import time
import cv2
import tensorflow.contrib.layers as layers

Height = 256
Width = 256
batch_size = 8

dPath_data = ('D:/dataset/places365challenge/**/')
dPath_test = ('D:/dataset/landscape/')

name_data = glob.glob(dPath_data + '*.jpg', recursive=True)
name_test = glob.glob(dPath_test + '*.jpg')
num_data = len(name_data)
total_batch = num_data // batch_size // 2

save_path = './mid_test/v3'
model_path = './model/v3'

restore = False
restore_point = 15000
Checkpoint = model_path + '/cVG iter ' + str(restore_point) + '/'
WeightName = Checkpoint + 'Train_' + str(restore_point) + '.meta'

if restore == False:
    restore_point = 0

saving_iter = 5000
Max_iter = 1000000

# ------- variables

X1 = tf.placeholder(tf.float32, [batch_size, Height, Width, 3])
X2 = tf.placeholder(tf.float32, [batch_size, Height, Width, 3])

MASK = tf.placeholder(tf.float32, [1, 1, Width*2, 1])

# ------- structure

img_input = tf.concat([tf.concat([X1, X2], 2), tf.concat([X2, X1], 2)], 0)

img_gen = mm.GNet(img_input, reuse=False, name='G_net')

d_real, f_real, v_dep = mm.DNet(tf.concat([X1, X2], 0), reuse=False, name='D_net')
d_fake, f_fake, _ = mm.DNet(img_gen, reuse=True, name='D_net')

_, _, W, _ = f_fake.shape.as_list()

MASK_temp = tf.image.resize(MASK, [1, W], method=1)

v1 = mm.Projection_v(X1, v_dep, reuse=False, name='VP')
v2 = mm.Projection_v(X2, v_dep, reuse=True, name='VP')
v = tf.concat([v1 * MASK_temp + v2 * (1 - MASK_temp), v2 * MASK_temp + v1 * (1 - MASK_temp)], axis=0)

d_real += tf.reduce_sum(f_real * tf.concat([v1, v2], 0), axis=3, keepdims=True)
d_fake += tf.reduce_sum(f_fake * v, axis=3, keepdims=True)

# ------- Loss

Loss_D = tf.reduce_mean(tf.nn.relu(1+d_fake)) + tf.reduce_mean(tf.nn.relu(1-d_real))

Loss_gan = -tf.reduce_mean(d_fake)

# MASK_temp = MASK[:, :, :Width, :] * 2 - 1
#
# MASK_temp = tf.concat([MASK_temp, tf.image.flip_left_right(MASK_temp)], 2)

img_input_av = layers.avg_pool2d(img_input, 64, stride=1, padding='SAME')

img_gen_av = layers.avg_pool2d(img_gen, 64, stride=1, padding='SAME')

Loss_re = tf.reduce_mean(tf.abs((img_input_av - img_gen_av)))

Loss_G = Loss_re + 0.1*Loss_gan

# --------------------- variable & optimizer

var_D = [v for v in tf.global_variables() if v.name.startswith('D_net') or v.name.startswith('VP')]
var_G = [v for v in tf.global_variables() if v.name.startswith('G_net')]

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

with tf.control_dependencies(update_ops):
    optimize_D = tf.train.AdamOptimizer(learning_rate=0.0004, beta1=0.5, beta2=0.9).minimize(Loss_D, var_list=var_D)
    optimize_G = tf.train.AdamOptimizer(learning_rate=0.0001, beta1=0.5, beta2=0.9).minimize(Loss_G, var_list=var_G)

# --------- Run

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.8
config.gpu_options.allow_growth = False

sess = tf.Session(config=config)

init = tf.global_variables_initializer()
sess.run(init)
saver = tf.train.Saver()

if restore == True:
    print('Weight Restoring.....')
    Restore = tf.train.import_meta_graph(WeightName)
    Restore.restore(sess, tf.train.latest_checkpoint(Checkpoint))
    print('Weight Restoring Finish!')
else:
    print('Weight Inception restoring')
    restore_val = []
    var_in = [v for v in tf.global_variables() if v.name.startswith('Inception')]
    for it in range(len(var_in)):
        vname = var_in[it].name
        if "biase" in vname:
            pass
        else:
            restore_val.append(var_in[it])
    temp = U.framework.list_variables('./inception_v3.ckpt')
    saverG = tf.train.Saver(restore_val)
    saverG.restore(sess, './inception_v3.ckpt')


start_time = time.time()
for iter_count in range(restore_point, Max_iter + 1):

    i = iter_count % total_batch
    e = iter_count // total_batch

    if i == 0:
        np.random.shuffle(name_data)

    # image_time = time.time()
    img1, img2 = pp.MakeImageBlock(name_data, Height, Width, i, batch_size, resize=False)
    # consume_time = time.time() - image_time
    # print('time = %.4f' % consume_time)

    mask = pp.MakeMask(Width)

    _, Loss1, check = sess.run([optimize_D, Loss_D, img_input_av], feed_dict={X1: img1, X2: img2, MASK: mask})
    # check = 255.0 * ((check[0] + 1) / 2.0)
    # cv2.imshow('a', check.astype(np.uint8))
    # cv2.waitKey(-1)
    _, Loss2 = sess.run([optimize_G, Loss_G], feed_dict={X1: img1, X2: img2, MASK: mask})

    if iter_count % 100 == 0:
        consume_time = time.time() - start_time
        print('%d\tEpoch : %d\tD Loss = %.5f\tG Loss = %.5f\ttime = %.4f' % (iter_count, e, Loss1, Loss2, consume_time))
        start_time = time.time()

    if iter_count % saving_iter == 0:

        print('SAVING MODEL')
        Temp = model_path + '/cVG iter %s/' % iter_count

        if not os.path.exists(Temp):
            os.makedirs(Temp)

        SaveName = (Temp + 'Train_%s' % (iter_count))
        saver.save(sess, SaveName)
        print('SAVING MODEL Finish')

        for isave in range(5):

            img_t1, img_t2 = pp.MakeRandomImageBlock(name_test, Height, Width, batch_size, resize=True)

            img_sample = sess.run(img_gen, feed_dict={X1: img_t1, X2: img_t2})

            for kk in range(batch_size):
                temp_img_gen = img_sample[kk]
                temp_img1 = img_t1[kk]
                temp_img2 = img_t2[kk]

                img_ge = 255.0 * ((temp_img_gen + 1) / 2.0)
                img_o1 = 255.0 * ((temp_img1 + 1) / 2.0)
                img_o2 = 255.0 * ((temp_img2 + 1) / 2.0)

                Bigpaper = np.zeros((Height, 4 * Width + 30, 3))
                Bigpaper[:, 0:Width, :] = img_o1
                Bigpaper[:, Width: 2 * Width, :] = img_o2
                Bigpaper[:, 2 * Width + 30:, :] = img_ge

                save_name = save_path + '/%04d' % iter_count
                name = save_name + '/img_%02d.png' % (isave * batch_size + kk)
                if not os.path.exists(save_name):
                    os.makedirs(save_name)
                cv2.imwrite(name, Bigpaper)

