# Copyright 2017 The TensorFlow Authors All Rights Reserved.
#
# Licensed 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.
# ==============================================================================

"""Model implementations."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from abc import ABCMeta
from abc import abstractmethod
import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim.python.slim.nets import inception
from tensorflow.contrib.slim.python.slim.nets import resnet_v2 as resnet_v2
from tensorflow.contrib.slim.python.slim.nets import resnet_utils as resnet_utils


def get_embedder(
    embedder_strategy, config, images, is_training, reuse=False,
    l2_normalize_embedding=True):
  """Returns an embedder based on config.

  Args:
    embedder_strategy: String, name of embedder version to return.
    config: LuaTable object, training config.
    images: 4-D float `Tensor` containing batch images.
    is_training: Boolean or placeholder for boolean,
      indicator for whether or not we're training.
    reuse: Boolean: Reuse embedder variable scope.
    l2_normalize_embedding: Boolean, whether or not to l2 normalize the
      embedding.
  Returns:
    embedder: An `Embedder` object.
  Raises:
    ValueError: if unknown embedder_strategy specified.
  """
  if embedder_strategy == 'inception_baseline':
    pretrained_ckpt = config.inception_conv_ss_fc.pretrained_checkpoint
    return InceptionBaselineEmbedder(
        images,
        pretrained_ckpt,
        config.random_projection,
        config.random_projection_dim)

  strategy_to_embedder = {
      'inception_conv_ss_fc': InceptionConvSSFCEmbedder,
      'resnet': ResnetEmbedder,
  }
  if embedder_strategy not in strategy_to_embedder:
    raise ValueError('unknown embedder_strategy', embedder_strategy)

  embedding_size = config.embedding_size
  l2_reg_weight = config.learning.l2_reg_weight
  embedder = strategy_to_embedder[embedder_strategy](
      config[embedder_strategy], images, embedding_size,
      is_training, embedding_l2=l2_normalize_embedding,
      l2_reg_weight=l2_reg_weight, reuse=reuse)
  return embedder


def build_inceptionv3_graph(images, endpoint, is_training, checkpoint,
                            reuse=False):
  """Builds an InceptionV3 model graph.

  Args:
    images: A 4-D float32 `Tensor` of batch images.
    endpoint: String, name of the InceptionV3 endpoint.
    is_training: Boolean, whether or not to build a training or inference graph.
    checkpoint: String, path to the pretrained model checkpoint.
    reuse: Boolean, whether or not we are reusing the embedder.
  Returns:
    inception_output: `Tensor` holding the InceptionV3 output.
    inception_variables: List of inception variables.
    init_fn: Function to initialize the weights (if not reusing, then None).
  """
  with slim.arg_scope(inception.inception_v3_arg_scope()):