#:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::#
#                                                                               #
# Part of solution to exercise set for week 13                                  #
# INF5860 - Machine Learning for Image analysis                                 #
# University of Oslo                                                            #
#                                                                               #
#                                                                               #
# Ole-Johan Skrede    olejohas at ifi dot uio dot no                            #
# 2018.04.12                                                                    #
#                                                                               #
#:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::#

"""An example to show how to use tensorflow embeddings"""

import os

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.contrib.tensorboard.plugins import projector
import matplotlib.pyplot as plt
import numpy as np

def config():
    """Return a dict of configuration settings used in the program"""

    conf = {}

    # Paths for embedding metadata and sprites
    conf['job_dir'] = "/tmp/embeddings/"
    # Relevant dataset will be put in this location after download
    conf['data_dir'] = "/tmp/mnist_data"

    conf['sprite_file'] = os.path.join(conf['job_dir'], 'mnist_sprite.png')
    conf['metadata_file'] = os.path.join(conf['job_dir'], 'mnist_metadata.tsv')

    # Create directories
    if not os.path.exists(conf['job_dir']):
        os.makedirs(conf['job_dir'])

    # Number of nodes in the input (the dimensions of an input example)
    conf['height'] = 28
    conf['width'] = 28

    # The number of images to be part of the visualization
    conf['batch_size'] = 10000

    return conf

def create_sprite_image(images):
    """Returns a sprite image containing all images"""
    print(images.shape)
    batch_size, height, width = images.shape
    num_plots = int(np.ceil(np.sqrt(batch_size)))

    sprite_image = np.ones((num_plots * height, num_plots * width))

    for ind_i in range(num_plots):
        for ind_j in range(num_plots):
            list_ind = ind_i * num_plots + ind_j
            if list_ind < batch_size:
                image = images[list_ind]
                sprite_image[ind_i * height:(ind_i + 1) * height,
                             ind_j * width:(ind_j + 1) * width] = image

    return sprite_image

def create_and_save_sprite(image_batch, height, width, save_path):
    """Create a sprite image of the image_batch and save it"""

    image_batch = image_batch.reshape((-1, height, width))
    # Invert the colors so that the background is white and the digits are black
    image_batch = 1.0 - image_batch
    sprite_image = create_sprite_image(image_batch)

    plt.imsave(save_path, sprite_image, cmap='gray')

def create_metadata_file(label_batch, save_path):
    """Create a tsv (tab separated value) file so that tensorboard can know the classes of the
    datapoints
    """
    with open(save_path, 'w') as metadata_file:
        metadata_file.write("Index\tLabel\n")
        for ind, label in enumerate(label_batch):
            metadata_file.write("{}\t{}\n".format(ind, label))



def main():
    """Main"""
    print("Start program")
    conf = config()
    mnist = input_data.read_data_sets(conf['data_dir'], one_hot=False)

    image_batch, label_batch = mnist.test.next_batch(conf['batch_size'])

    images = tf.get_variable(name="images",
                             shape=[conf['batch_size'], conf['height'] * conf['width']],
                             dtype=tf.float32,
                             initializer=tf.constant_initializer(image_batch),
                             trainable=False)
    summary_writer = tf.summary.FileWriter(conf['job_dir'])

    projector_config = projector.ProjectorConfig()
    embedding = projector_config.embeddings.add()
    embedding.tensor_name = images.name

    embedding.metadata_path = conf['metadata_file']

    embedding.sprite.image_path = conf['sprite_file']
    embedding.sprite.single_image_dim.extend([conf['height'], conf['width']])

    projector.visualize_embeddings(summary_writer, projector_config)

    sess = tf.Session()
    sess.run(tf.global_variables_initializer())
    saver = tf.train.Saver()
    saver.save(sess, os.path.join(conf['job_dir'], "model.ckpt"), 1)

    create_and_save_sprite(image_batch, conf['height'], conf['width'], conf['sprite_file'])

    create_metadata_file(label_batch, conf['metadata_file'])



if __name__ == "__main__":
    main()