"""Routines used for visualizing data or a learned representation.

We used the off-the-shelf plugin of tensorboard to visualize vectors
in an embedding space.

The subclass inheriting from `Visualize` should implement `get_data`,
`compute` and `config_embedding`.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import tensorflow as tf

from nets_base.arg_scope import nets_arg_scope

slim = tf.contrib.slim


class Visualize(object):
    """Implement an abstract class for data visualization.

    In its subclass `get_data`, `compute` and `config_embedding`
    must be implemented. For an example see `VisualizeImages`.
    """

    __metaclass__ = abc.ABCMeta

    def visualize(self,
                  tfrecord_dir,
                  checkpoint_dirs,
                  log_dir=None,
                  batch_size=300,
                  split_name='train',
                  shuffle=True,
                  use_batch_norm=True,
                  batch_stat=False,
                  **kwargs):
        """Visualize data or its representation in an embedding space.

        Args:
            tfrecord_dir: The directory that contains the dataset tfreocrds
                (which can be generated by `convert_TFrecord` scripts).
            checkpoints_dir: The directorys containing checkpoints of
                the models used for visualization.
            log_dir: The directory to log the checkpoint.
            batch_size: Number of samples to read and to show
                in the embedding.
            split_name: The part of the dataset to use.
            shuffle: Whether to shuffle the data for the evaluation.
            use_batch_norm: Passes to `self.used_arg_scope` to decide
                whether to use batch normalization.
            batch_stat: Whether to use batch statistics or moving
                mean/variance for batch normalization.
            **kwargs: Arguments pass to the `self.compute`.
        """
        if log_dir is not None and not tf.gfile.Exists(log_dir):
            tf.gfile.MakeDirs(log_dir)

        if (checkpoint_dirs is not None
                and not isinstance(checkpoint_dirs, (tuple, list))):
            checkpoint_dirs = [checkpoint_dirs]

        with tf.Graph().as_default():

            with tf.name_scope('Data_provider'):
                self.get_data(split_name, tfrecord_dir, batch_size, shuffle)

            with slim.arg_scope(self.used_arg_scope(
                    batch_stat, use_batch_norm)):
                self.compute(**kwargs)

            with tf.Session() as sess:
                with slim.queues.QueueRunners(sess):
                    self.init_model(sess, checkpoint_dirs)
                    self.config_embedding(sess, log_dir)

    @abc.abstractmethod
    def get_data(self, split_name, tfrecord_dir, batch_size, shuffle):
        """Read data from some directory and load them in batches.

        Args:
            tfrecord_dir: The directory where the tfrecords of the
                dataset are stored.
            batch_size: The number of elements contained in each batch.
            shuffle: Whether to shuffle input data or not.
        """
        pass

    def used_arg_scope(self, batch_stat, use_batch_norm):
        """The slim argument scope that is used for main computations.

        Args:
            use_batch_norm: Whether to do batch normalization or not.
            renorm: Whether to do batch renormalization or not. I've in
                fact never used it.

        Returns:
            An argument scope to be used for model computations.
        """
        return nets_arg_scope(
            is_training=batch_stat, use_batch_norm=use_batch_norm)

    @abc.abstractmethod
    def compute(self, **kwargs):
        """Compute necessary values of the model.

        The content of this function can vary a lot from case to case and
        we don't have some fixed arguments for this functions.
        For example, we may want to compute just one representation of
        input data of we may want to visualize activations all the layers
        in the network.
        """
        pass

    def init_model(self, sess, checkpoint_dirs):
        """Initialize the model from a/several directory(ies).

        Args:
            sess: The session in which we run the initialization.
            checkpoint_dirs: The directory(ies) containing the necessay
                models used for computation.
        """
        if checkpoint_dirs is not None:
            assert len(checkpoint_dirs) == 1
            checkpoint_path = tf.train.latest_checkpoint(checkpoint_dirs[0])
            saver = tf.train.Saver(tf.model_variables())
            saver.restore(sess, checkpoint_path)
        else:
            sess.run(tf.global_variables_initializer())

    @abc.abstractmethod
    def config_embedding(self, sess, log_dir):
        """Configurations for the plugin to visualize data.

        Args:
            sess: Session used for computation.
            log_dir: Where to save the checkpoint for visualization.
        """
        pass


def visualize(visualize_class,
              used_architecture,
              tfrecord_dir,
              checkpoint_dirs,
              log_dir,
              **kwargs):
    visualize_instance = visualize_class(used_architecture)
    for key in kwargs.copy():
        if hasattr(visualize_instance, key):
            setattr(visualize_instance, key, kwargs[key])
            del kwargs[key]
    visualize_instance.visualize(
        tfrecord_dir, checkpoint_dirs, log_dir, **kwargs)
