# Copyright 2023 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.

"""Neural network operations commonly shared by the architectures."""

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

import functools

import tensorflow as tf, tf_keras


class NormActivation(tf_keras.layers.Layer):
  """Combined Normalization and Activation layers."""

  def __init__(self,
               momentum=0.997,
               epsilon=1e-4,
               trainable=True,
               init_zero=False,
               use_activation=True,
               activation='relu',
               fused=True,
               name=None):
    """A class to construct layers for a batch normalization followed by a ReLU.

    Args:
      momentum: momentum for the moving average.
      epsilon: small float added to variance to avoid dividing by zero.
      trainable: `bool`, if True also add variables to the graph collection
        GraphKeys.TRAINABLE_VARIABLES. If False, freeze batch normalization
        layer.
      init_zero: `bool` if True, initializes scale parameter of batch
        normalization with 0. If False, initialize it with 1.
      use_activation: `bool`, whether to add the optional activation layer after
        the batch normalization layer.
      activation: 'string', the type of the activation layer. Currently support
        `relu` and `swish`.
      fused: `bool` fused option in batch normalziation.
      name: `str` name for the operation.
    """
    super(NormActivation, self).__init__(trainable=trainable)
    if init_zero:
      gamma_initializer = tf_keras.initializers.Zeros()
    else:
      gamma_initializer = tf_keras.initializers.Ones()
    self._normalization_op = tf_keras.layers.BatchNormalization(
        momentum=momentum,
        epsilon=epsilon,
        center=True,
        scale=True,
        trainable=trainable,
        fused=fused,
        gamma_initializer=gamma_initializer,
        name=name)
    self._use_activation = use_activation
    if activation == 'relu':
      self._activation_op = tf.nn.relu
    elif activation == 'swish':
      self._activation_op = tf.nn.swish
    else:
      raise ValueError('Unsupported activation `{}`.'.format(activation))

  def __call__(self, inputs, is_training=None):
    """Builds the normalization layer followed by an optional activation layer.

    Args:
      inputs: `Tensor` of shape `[batch, channels, ...]`.
      is_training: `boolean`, if True if model is in training mode.

    Returns:
      A normalized `Tensor` with the same `data_format`.
    """
    # We will need to keep training=None by default, so that it can be inherit
    # from keras.Model.training
    if is_training and self.trainable:
      is_training = True
    inputs = self._normalization_op(inputs, training=is_training)

    if self._use_activation:
      inputs = self._activation_op(inputs)
    return inputs


def norm_activation_builder(momentum=0.997,
                            epsilon=1e-4,
                            trainable=True,
                            activation='relu',
                            **kwargs):
  return functools.partial(
      NormActivation,
      momentum=momentum,
      epsilon=epsilon,
      trainable=trainable,
      activation=activation,
      **kwargs)