# coding=utf-8
# Copyright 2020 The TF-Agents Authors.
#
# 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
#
#     https://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.

"""Keras layer mirroring tf.contrib.layers.bias_add."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import


class BiasLayer(tf.keras.layers.Layer):
  """Keras layer that only adds a bias to the input.

  `BiasLayer` implements the operation:
  `output = input + bias`

  Args:
      bias_initializer: Initializer for the bias vector.
  Input shape:
      nD tensor with shape: `(batch_size, ..., input_dim)`. The most common
        situation would be a 2D input with shape `(batch_size, input_dim)`. Note
        a rank of at least 2 is required.
  Output shape:
      nD tensor with shape: `(batch_size, ..., input_dim)`. For instance, for a
        2D input with shape `(batch_size, input_dim)`, the output would have
        shape `(batch_size, input_dim)`.
  """

  def __init__(self, bias_initializer='zeros', **kwargs):
    if 'input_shape' not in kwargs and 'input_dim' in kwargs:
      kwargs['input_shape'] = (kwargs.pop('input_dim'),)

    super(BiasLayer, self).__init__(**kwargs)
    self.bias_initializer = tf.keras.initializers.get(bias_initializer)

    self.supports_masking = True

  def build(self, input_shape):
    input_shape = tf.TensorShape(input_shape)
    if input_shape.rank == 1:
      shape = (1,)
    else:
      shape = (tf.compat.dimension_value(input_shape[-1]),)

    self.bias = self.add_weight(
        'bias',
        shape=shape,
        initializer=self.bias_initializer,
        dtype=self.dtype,
        trainable=True)
    self.built = True

  def call(self, inputs):
    if inputs.shape.rank == 1:
      expanded_inputs = tf.expand_dims(inputs, -1)
      with_bias = tf.nn.bias_add(expanded_inputs, self.bias)
      return with_bias[..., 0]
    return tf.nn.bias_add(inputs, self.bias)

  def compute_output_shape(self, input_shape):
    return input_shape

  def get_config(self):
    config = {
        'bias_initializer':
            tf.keras.initializers.serialize(self.bias_initializer),
    }
    base_config = super(BiasLayer, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))
