__author__ = 'Charlie'
"""Image coloring by fully convolutional networks - incomplete """
import numpy as np
import tensorflow as tf
import os, sys, inspect
from datetime import datetime
import scipy.misc as misc

lib_path = os.path.realpath(
    os.path.abspath(os.path.join(os.path.split(inspect.getfile(inspect.currentframe()))[0], "..")))
if lib_path not in sys.path:
    sys.path.insert(0, lib_path)

import TensorflowUtils as utils

FLAGS = tf.flags.FLAGS
tf.flags.DEFINE_string("data_dir", "Data_zoo/CIFAR10_data/", """Path to the CIFAR10 data""")
tf.flags.DEFINE_string("mode", "train", "Network mode train/ test")
tf.flags.DEFINE_string("test_image_path", "", "Path to test image - read only if mode is test")
tf.flags.DEFINE_integer("batch_size", "128", "train batch size")
tf.flags.DEFINE_string("logs_dir", "logs/ImageColoring_logs/", """Path to save logs and checkpoint if needed""")

DATA_URL = 'http://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz'

LEARNING_RATE = 1e-3
MAX_ITERATIONS = 100001

NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 20000

IMAGE_SIZE = 32


def read_cifar10(filename_queue):
    class CIFAR10Record(object):
        pass

    result = CIFAR10Record()

    label_bytes = 1
    result.height = IMAGE_SIZE
    result.width = IMAGE_SIZE
    result.depth = 3
    image_bytes = result.height * result.width * result.depth
    record_bytes = label_bytes + image_bytes

    reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
    result.key, value = reader.read(filename_queue)

    record_bytes = tf.decode_raw(value, tf.uint8)

    depth_major = tf.cast(tf.reshape(tf.slice(record_bytes, [label_bytes], [image_bytes]),
                                     [result.depth, result.height, result.width]), tf.float32)

    image = tf.transpose(depth_major, [1, 2, 0])
    # extended_image = tf.reshape(image, (result.height, result.width, result.depth))
    result.color_image = image
    print result.color_image.get_shape()
    print "Converting image to gray scale"
    result.gray_image = 0.21 * result.color_image[ :, :, 2] + 0.72 * result.color_image[ :, :,
                                                                       1] + 0.07 * result.color_image[ :, :, 0]
    result.gray_image = tf.expand_dims(result.gray_image, 2)
    print result.gray_image.get_shape()

    return result


def get_image(image_dir):
    image = misc.imread(image_dir)
    image = np.ndarray.reshape(image.astype(np.float32), ((1,) + image.shape))
    return image


def inputs():
    data_dir = os.path.join(FLAGS.data_dir, 'cifar-10-batches-bin')
    filenames = [os.path.join(data_dir, 'data_batch_%d.bin' % i) for i in xrange(1, 6)]
    for f in filenames:
        if not tf.gfile.Exists(f):
            raise ValueError('Failed to find file: ' + f)

    filename_queue = tf.train.string_input_producer(filenames)
    read_input = read_cifar10(filename_queue)
    num_preprocess_threads = 8
    min_queue_examples = int(0.4 * NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN)
    print "Shuffling"
    input_gray, input_colored = tf.train.shuffle_batch([read_input.gray_image, read_input.color_image],
                                                       batch_size=FLAGS.batch_size,
                                                       num_threads=num_preprocess_threads,
                                                       capacity=min_queue_examples + 3 * FLAGS.batch_size,
                                                       min_after_dequeue=min_queue_examples)
    input_gray = (input_gray - 128) / 255.0
    input_colored = (input_colored - 128) / 255.0
    return input_gray, input_colored


def inference(image):
    W1 = utils.weight_variable_xavier_initialized([9, 9, 1, 32])
    b1 = utils.bias_variable([32])
    tf.histogram_summary("W1", W1)
    tf.histogram_summary("b1", b1)
    h_conv1 = tf.nn.relu(utils.conv2d_basic(image, W1, b1))

    W2 = utils.weight_variable_xavier_initialized([3, 3, 32, 64])
    b2 = utils.bias_variable([64])
    tf.histogram_summary("W2", W2)
    tf.histogram_summary("b2", b2)
    h_conv2 = tf.nn.relu(utils.conv2d_strided(h_conv1, W2, b2))

    W3 = utils.weight_variable_xavier_initialized([3, 3, 64, 128])
    b3 = utils.bias_variable([128])
    tf.histogram_summary("W3", W3)
    tf.histogram_summary("b3", b3)
    h_conv3 = tf.nn.relu(utils.conv2d_strided(h_conv2, W3, b3))

    # upstrides
    W4 = utils.weight_variable_xavier_initialized([3, 3, 64, 128])
    b4 = utils.bias_variable([64])
    tf.histogram_summary("W4", W4)
    tf.histogram_summary("b4", b4)
    h_conv4 = tf.nn.relu(utils.conv2d_transpose_strided(h_conv3, W4, b4))

    W5 = utils.weight_variable_xavier_initialized([3, 3, 32, 64])
    b5 = utils.bias_variable([32])
    tf.histogram_summary("W5", W5)
    tf.histogram_summary("b5", b5)
    h_conv5 = tf.nn.relu(utils.conv2d_transpose_strided(h_conv4, W5, b5))

    W6 = utils.weight_variable_xavier_initialized([9, 9, 32, 3])
    b6 = utils.bias_variable([3])
    tf.histogram_summary("W6", W6)
    tf.histogram_summary("b6", b6)
    pred_image = tf.nn.tanh(utils.conv2d_basic(h_conv5, W6, b6))

    return pred_image


def loss(pred, colored):
    rmse = tf.sqrt(2 * tf.nn.l2_loss(tf.sub(colored, pred))) / FLAGS.batch_size
    tf.scalar_summary("RMSE", rmse)
    return rmse


def train(loss_val, step):
    learning_rate = tf.train.exponential_decay(LEARNING_RATE, step, 0.4 * MAX_ITERATIONS, 0.99)
    train_op = tf.train.AdamOptimizer(learning_rate).minimize(loss_val, global_step=step)
    return train_op


def main(argv=None):
    utils.maybe_download_and_extract(FLAGS.data_dir, DATA_URL, is_tarfile=True)
    print "Setting up model..."
    global_step = tf.Variable(0,trainable=False)
    gray, color = inputs()
    pred = 255 * inference(gray) + 128
    tf.image_summary("Gray", gray, max_images=1)
    tf.image_summary("Ground_truth", color, max_images=1)
    tf.image_summary("Prediction", pred, max_images=1)

    image_loss = loss(pred, color)
    train_op = train(image_loss, global_step)

    summary_op = tf.merge_all_summaries()
    with tf.Session() as sess:
        print "Setting up summary writer, queue, saver..."
        sess.run(tf.initialize_all_variables())
        
        summary_writer = tf.train.SummaryWriter(FLAGS.logs_dir, sess.graph)
        saver = tf.train.Saver()

        ckpt = tf.train.get_checkpoint_state(FLAGS.logs_dir)
        if ckpt and ckpt.model_checkpoint_path:
            print "Restoring model from checkpoint..."
            saver.restore(sess, ckpt.model_checkpoint_path)
        tf.train.start_queue_runners(sess)
        for step in xrange(MAX_ITERATIONS):
            if step % 400 == 0:
                loss_val, summary_str = sess.run([image_loss, summary_op])
                print "Step %d, Loss: %g" % (step, loss_val)
                summary_writer.add_summary(summary_str, global_step=step)

            if step % 1000 == 0:
                saver.save(sess, FLAGS.logs_dir + "model.ckpt", global_step=step)
                print "%s" % datetime.now()

            sess.run(train_op)

if __name__ == "__main__":
    tf.app.run()
