# 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.
# ==============================================================================
"""DenseNet implementation with TPU support.

Original paper: (https://arxiv.org/abs/1608.06993)
"""

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

from absl import flags
import tensorflow.compat.v1 as tf
from tensorflow.contrib import layers as contrib_layers

# Learning hyperaparmeters
_BATCH_NORM_DECAY = 0.997
_BATCH_NORM_EPSILON = 1e-5

FLAGS = flags.FLAGS
flags.DEFINE_bool("use_bottleneck", False, "Use bottleneck convolution layers")


def conv(image, filters, strides=1, kernel_size=3):
  """Convolution with default options from the densenet paper."""
  # Use initialization from https://arxiv.org/pdf/1502.01852.pdf

  return tf.layers.conv2d(
      inputs=image,
      filters=filters,
      kernel_size=kernel_size,
      strides=strides,
      activation=tf.identity,
      use_bias=False,
      padding="same",
      kernel_initializer=tf.variance_scaling_initializer(),
  )


def dense_block(image, filters, is_training):
  """Standard BN+Relu+conv block for DenseNet."""
  image = tf.layers.batch_normalization(
      inputs=image,
      axis=-1,
      training=is_training,
      fused=True,
      center=True,
      scale=True,
      momentum=_BATCH_NORM_DECAY,
      epsilon=_BATCH_NORM_EPSILON,
  )

  if FLAGS.use_bottleneck:
    # Add bottleneck layer to optimize computation and reduce HBM space
    image = tf.nn.relu(image)
    image = conv(image, 4 * filters, strides=1, kernel_size=1)
    image = tf.layers.batch_normalization(
        inputs=image,
        axis=-1,
        training=is_training,
        fused=True,
        center=True,
        scale=True,
        momentum=_BATCH_NORM_DECAY,
        epsilon=_BATCH_NORM_EPSILON,
    )

  image = tf.nn.relu(image)
  return conv(image, filters)


def transition_layer(image, filters, is_training):
  """Construct the transition layer with specified growth rate."""

  image = tf.layers.batch_normalization(
      inputs=image,
      axis=-1,
      training=is_training,
      fused=True,
      center=True,
      scale=True,
      momentum=_BATCH_NORM_DECAY,
      epsilon=_BATCH_NORM_EPSILON,
  )
  image = tf.nn.relu(image)
  conv_img = conv(image, filters=filters, kernel_size=1)
  return tf.layers.average_pooling2d(
      conv_img, pool_size=2, strides=2, padding="same")


def _int_shape(layer):
  return layer.get_shape().as_list()


# Definition of the CIFAR-10 network
def densenet_cifar_model(image,
                         k,
                         layers,
                         is_training=True,
                         num_blocks=3,
                         num_classes=10):
  """Construct a DenseNet with the specified growth size and layers."""
  layers_per_block = int((layers - 4) / num_blocks)

  v = conv(image, filters=2*k, strides=(1, 1), kernel_size=(3, 3))
  for i in range(num_blocks):
    with tf.variable_scope("block-%d" % i):
      for j in range(layers_per_block):
        with tf.variable_scope("conv-%d-%d" % (i, j)):
          dv = dense_block(v, k, is_training)
          v = tf.concat([v, dv], axis=3)
    if i != num_blocks - 1:
      with tf.variable_scope("transition-%d" % i):
        v = transition_layer(v, _int_shape(v)[3], is_training)

  global_pool = tf.reduce_sum(v, axis=(2, 3), name="global_pool")
  logits = tf.layers.dense(
      global_pool,
      units=num_classes,
      activation=tf.identity,
      kernel_initializer=tf.random_normal_initializer(
          stddev=2.0 / (_int_shape(global_pool)[1] * 10)),
      kernel_regularizer=contrib_layers.l2_regularizer(1e-4),
      bias_regularizer=contrib_layers.l2_regularizer(1e-4),
  )
  return logits


# Definition of the Imagenet network
def densenet_imagenet_model(image, k, depths, num_classes, is_training=True):
  """Construct a DenseNet with the specified growth size and layers."""

  num_channels = 2 * k
  v = conv(image, filters=2 * k, strides=2, kernel_size=7)
  v = tf.layers.batch_normalization(
      inputs=v,
      axis=-1,
      training=is_training,
      fused=True,
      center=True,
      scale=True,
      momentum=_BATCH_NORM_DECAY,
      epsilon=_BATCH_NORM_EPSILON,
  )
  v = tf.nn.relu(v)
  v = tf.layers.max_pooling2d(v, pool_size=3, strides=2, padding="same")
  for i, depth in enumerate(depths):
    with tf.variable_scope("block-%d" % i):
      for j in range(depth):
        with tf.variable_scope("denseblock-%d-%d" % (i, j)):
          output = dense_block(v, k, is_training)
          v = tf.concat([v, output], axis=3)
          num_channels += k
      if i != len(depths) - 1:
        num_channels /= 2
        v = transition_layer(v, int(num_channels), is_training)

  global_pool = tf.reduce_mean(v, axis=(1, 2), name="global_pool")
  dense_layer = tf.layers.dense(global_pool, units=num_classes)
  logits = tf.identity(dense_layer, "final_dense")

  return logits


def densenet_imagenet_121(inputs, is_training=True, num_classes=1001):
  """DenseNet 121."""
  depths = [6, 12, 24, 16]
  growth_rate = 32
  return densenet_imagenet_model(inputs, growth_rate, depths, num_classes,
                                 is_training)


def densenet_imagenet_169(inputs, is_training=True, num_classes=1001):
  """DenseNet 121."""
  depths = [6, 12, 32, 32]
  growth_rate = 32
  return densenet_imagenet_model(inputs, growth_rate, depths, num_classes,
                                 is_training)


def densenet_imagenet_201(inputs, is_training=True, num_classes=1001):
  """DenseNet 121."""
  depths = [6, 12, 48, 32]
  growth_rate = 32
  return densenet_imagenet_model(inputs, growth_rate, depths, num_classes,
                                 is_training)
