#!/usr/bin/env python
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#    http://www.apache.org/licenses/LICENSE-2.0
#  Unless required by applicable law or agreed to in writing,
#  software distributed under the License is distributed on an
#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#  KIND, either express or implied.  See the License for the
#  specific language governing permissions and limitations
#  under the License.


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

import os.path

import tensorflow as tf
from tensorflow.contrib.slim.python.slim.nets.inception_v3 import inception_v3_base

slim = tf.contrib.slim


class ModelWrapper(object):
    """
        Model wrapper class to perform image captioning with a ShowAndTellModel
    """

    def __init__(self):
        super(ModelWrapper, self).__init__()

    def build_graph(self, checkpoint_path):
        """Builds the inference graph"""

        tf.logging.info("Building model.")
        ShowAndTellModel().build()
        saver = tf.train.Saver()

        return self._create_restore_fn(checkpoint_path, saver)

    def _create_restore_fn(self, checkpoint_path, saver):
        """Creates a function that restores a model from checkpoint file"""

        if tf.gfile.IsDirectory(checkpoint_path):
            checkpoint_path = tf.train.latest_checkpoint(checkpoint_path)
            if not checkpoint_path:
                raise ValueError("No checkpoint file found in: %s" % checkpoint_path)

        def _restore_fn(sess):
            tf.logging.info("Loading model from checkpoint: %s", checkpoint_path)
            saver.restore(sess, checkpoint_path)
            tf.logging.info("Successfully loaded checkpoint: %s",
                            os.path.basename(checkpoint_path))

        return _restore_fn

    def feed_image(self, sess, encoded_image):
        initial_state = sess.run(fetches="lstm/initial_state:0",
                                 feed_dict={"image_feed:0": encoded_image})
        return initial_state

    def inference_step(self, sess, input_feed, state_feed):
        softmax_output, state_output = sess.run(
            fetches=["softmax:0", "lstm/state:0"],
            feed_dict={
                "input_feed:0": input_feed,
                "lstm/state_feed:0": state_feed,
            })
        return softmax_output, state_output


class ShowAndTellModel(object):
    """
        Image captioning implementation based on the paper,

        "Show and Tell: A Neural Image Caption Generator"
        Oriol Vinyals, Alexander Toshev, Samy Bengio, Dumitru Erhan

        For more details, please visit : http://arxiv.org/abs/1411.4555
    """

    def __init__(self):

        # scale used to initialize model variables
        self.initializer_scale = 0.08

        # dimensions of Inception v3 input images
        self.image_height = 299
        self.image_width = 299

        # LSTM input and output dimensionality, respectively
        self.embedding_size = 512
        self.num_lstm_units = 512

        # number of unique words in the vocab (plus 1, for <UNK>)
        # the default value is larger than the expected actual vocab size to allow
        # for differences between tokenizer versions used in preprocessing, there is
        # no harm in using a value greater than the actual vocab size, but using a
        # value less than the actual vocab size will result in an error
        self.vocab_size = 12000

        # reader for the input data
        self.reader = tf.TFRecordReader()

        # to match the "Show and Tell" paper we initialize all variables with a
        # random uniform initializer
        self.initializer = tf.random_uniform_initializer(
            minval=-self.initializer_scale,
            maxval=self.initializer_scale)

        # a float32 Tensor with shape [batch_size, height, width, channels]
        self.images = None

        # an int32 Tensor with shape [batch_size, padded_length]
        self.input_seqs = None

        # an int32 Tensor with shape [batch_size, padded_length]
        self.target_seqs = None

        # an int32 0/1 Tensor with shape [batch_size, padded_length]
        self.input_mask = None

        # a float32 Tensor with shape [batch_size, embedding_size]
        self.image_embeddings = None

        # a float32 Tensor with shape [batch_size, padded_length, embedding_size]
        self.seq_embeddings = None

        # collection of variables from the inception submodel
        self.inception_variables = []

        # global step Tensor
        self.global_step = None

    def process_image(self, encoded_image, resize_height=346, resize_width=346, thread_id=0):
        """Decodes and processes an image string"""

        # helper function to log an image summary to the visualizer. Summaries are
        # only logged in thread 0
        def image_summary(name, img):
            if not thread_id:
                tf.summary.image(name, tf.expand_dims(img, 0))

        # decode image into a float32 Tensor of shape [?, ?, 3] with values in [0, 1)
        with tf.name_scope("decode", values=[encoded_image]):
            image = tf.image.decode_jpeg(encoded_image, channels=3)

        image = tf.image.convert_image_dtype(image, dtype=tf.float32)
        image_summary("original_image", image)

        # resize image
        assert (resize_height > 0) == (resize_width > 0)
        if resize_height:
            image = tf.image.resize_images(image,
                                           size=[resize_height, resize_width],
                                           method=tf.image.ResizeMethod.BILINEAR)

        # central crop, assuming resize_height > height, resize_width > width
        image = tf.image.resize_image_with_crop_or_pad(image, self.image_height, self.image_width)

        image_summary("resized_image", image)

        image_summary("final_image", image)

        # rescale to [-1,1] instead of [0, 1]
        image = tf.subtract(image, 0.5)
        image = tf.multiply(image, 2.0)
        return image

    def build_inputs(self):
        """Input prefetching, preprocessing and batching"""

        image_feed = tf.placeholder(dtype=tf.string, shape=[], name="image_feed")
        input_feed = tf.placeholder(dtype=tf.int64,
                                    shape=[None],  # batch_size
                                    name="input_feed")

        # process image and insert batch dimensions
        images = tf.expand_dims(self.process_image(image_feed), 0)
        input_seqs = tf.expand_dims(input_feed, 1)

        # no target sequences or input mask in inference mode
        target_seqs = None
        input_mask = None

        self.images = images
        self.input_seqs = input_seqs
        self.target_seqs = target_seqs
        self.input_mask = input_mask

    def build_image_embeddings(self):
        """Builds the image model(Inception V3) subgraph and generates image embeddings"""

        # parameter initialization
        batch_norm_params = {
            "is_training": False,
            "trainable": False,
            # decay for the moving averages
            "decay": 0.9997,
            # epsilon to prevent 0s in variance
            "epsilon": 0.001,
            # collection containing the moving mean and moving variance
            "variables_collections": {
                "beta": None,
                "gamma": None,
                "moving_mean": ["moving_vars"],
                "moving_variance": ["moving_vars"],
            }
        }

        stddev = 0.1,
        dropout_keep_prob = 0.8

        with tf.variable_scope("InceptionV3", "InceptionV3", [self.images]) as scope:
            with slim.arg_scope(
                    [slim.conv2d, slim.fully_connected],
                    weights_regularizer=None,
                    trainable=False):
                with slim.arg_scope(
                        [slim.conv2d],
                        weights_initializer=tf.truncated_normal_initializer(stddev=stddev),
                        activation_fn=tf.nn.relu,
                        normalizer_fn=slim.batch_norm,
                        normalizer_params=batch_norm_params):
                    net, end_points = inception_v3_base(self.images, scope=scope)
                    with tf.variable_scope("logits"):
                        shape = net.get_shape()
                        net = slim.avg_pool2d(net, shape[1:3], padding="VALID", scope="pool")
                        net = slim.dropout(
                            net,
                            keep_prob=dropout_keep_prob,
                            is_training=False,
                            scope="dropout")
                        net = slim.flatten(net, scope="flatten")

        # add summaries
        for v in end_points.values():
            tf.contrib.layers.summaries.summarize_activation(v)

        self.inception_variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="InceptionV3")

        # map inception output(net) into embedding space
        with tf.variable_scope("image_embedding") as scope:
            image_embeddings = tf.contrib.layers.fully_connected(
                inputs=net,
                num_outputs=self.embedding_size,
                activation_fn=None,
                weights_initializer=self.initializer,
                biases_initializer=None,
                scope=scope)

        # save the embedding size in the graph
        tf.constant(self.embedding_size, name="embedding_size")

        self.image_embeddings = image_embeddings

    def build_seq_embeddings(self):
        """Builds the input sequence embeddings"""

        with tf.variable_scope("seq_embedding"), tf.device("/cpu:0"):
            embedding_map = tf.get_variable(
                name="map",
                shape=[self.vocab_size, self.embedding_size],
                initializer=self.initializer)
            seq_embeddings = tf.nn.embedding_lookup(embedding_map, self.input_seqs)

        self.seq_embeddings = seq_embeddings

    def build_model(self):

        # this LSTM cell has biases and outputs tanh(new_c) * sigmoid(o), but the
        # modified LSTM in the "Show and Tell" paper has no biases and outputs
        # new_c * sigmoid(o).

        lstm_cell = tf.contrib.rnn.BasicLSTMCell(
            num_units=self.num_lstm_units, state_is_tuple=True)

        with tf.variable_scope("lstm", initializer=self.initializer) as lstm_scope:
            # feed the image embeddings to set the initial LSTM state
            zero_state = lstm_cell.zero_state(
                batch_size=self.image_embeddings.get_shape()[0], dtype=tf.float32)
            _, initial_state = lstm_cell(self.image_embeddings, zero_state)

            # allow the LSTM variables to be reused
            lstm_scope.reuse_variables()

            # because this is inference mode,
            # use concatenated states for convenient feeding and fetching
            tf.concat(axis=1, values=initial_state, name="initial_state")

            # placeholder for feeding a batch of concatenated states
            state_feed = tf.placeholder(dtype=tf.float32,
                                        shape=[None, sum(lstm_cell.state_size)],
                                        name="state_feed")
            state_tuple = tf.split(value=state_feed, num_or_size_splits=2, axis=1)

            # run a single LSTM step
            lstm_outputs, state_tuple = lstm_cell(
                inputs=tf.squeeze(self.seq_embeddings, axis=[1]),
                state=state_tuple)

            # concatentate the resulting state
            tf.concat(axis=1, values=state_tuple, name="state")

        # stack batches vertically
        lstm_outputs = tf.reshape(lstm_outputs, [-1, lstm_cell.output_size])

        with tf.variable_scope("logits") as logits_scope:
            logits = tf.contrib.layers.fully_connected(
                inputs=lstm_outputs,
                num_outputs=self.vocab_size,
                activation_fn=None,
                weights_initializer=self.initializer,
                scope=logits_scope)

        tf.nn.softmax(logits, name="softmax")

    def setup_global_step(self):
        """Sets up the global step Tensor"""

        global_step = tf.Variable(
            initial_value=0,
            name="global_step",
            trainable=False,
            collections=[tf.GraphKeys.GLOBAL_STEP, tf.GraphKeys.GLOBAL_VARIABLES])

        self.global_step = global_step

    def build(self):
        self.build_inputs()
        self.build_image_embeddings()
        self.build_seq_embeddings()
        self.build_model()
        self.setup_global_step()