# Copyright 2017 The Sonnet 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.
# ============================================================================

"""A minimal interface mlp module."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

from six.moves import xrange  # pylint: disable=redefined-builtin
from sonnet.python.modules import base
from sonnet.python.modules import basic
from sonnet.python.modules import util

import tensorflow as tf


class MLP(base.AbstractModule, base.Transposable):
  """A Multi-Layer perceptron module."""

  def __init__(self,
               output_sizes,
               activation=tf.nn.relu,
               activate_final=False,
               initializers=None,
               partitioners=None,
               regularizers=None,
               use_bias=True,
               custom_getter=None,
               name="mlp"):
    """Constructs an MLP module.

    Args:
      output_sizes: An iterable of output dimensionalities as defined in
        `basic.Linear`. Output size can be defined either as number or via a
        callable. In the latter case, since the function invocation is deferred
        to graph construction time, the user must only ensure that entries can
        be called when build is called. Each entry in the iterable defines
        properties in the corresponding linear layer.
      activation: An activation op. The activation is applied to intermediate
        layers, and optionally to the output of the final layer.
      activate_final: Boolean determining if the activation is applied to
        the output of the final layer. Default `False`.
      initializers: Optional dict containing ops to initialize the linear
        layers' weights (with key 'w') or biases (with key 'b').
      partitioners: Optional dict containing partitioners to partition the
        linear layers' weights (with key 'w') or biases (with key 'b').
      regularizers: Optional dict containing regularizers for the linear layers'
        weights (with key 'w') and the biases (with key 'b'). As a default, no
        regularizers are used. A regularizer should be a function that takes
        a single `Tensor` as an input and returns a scalar `Tensor` output, e.g.
        the L1 and L2 regularizers in `tf.contrib.layers`.
      use_bias: Whether to include bias parameters in the linear layers.
        Default `True`.