# coding=utf-8
# Copyright 2021 The Trax 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
#
#     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.

# Lint as: python3
"""Trax pooling layers."""

from trax import fastmath
from trax.layers.base import Fn


# pylint: disable=invalid-name
def MaxPool(pool_size=(2, 2), strides=None, padding='VALID'):
  """Reduces each multi-dimensional window to the max of the window's values.

  Windows, as specified by `pool_size` and `strides`, involve all axes of an
  n-dimensional array except the first and last: :math:`(d_1, ..., d_{n-2})`
  from shape :math:`(d_0, d_1, ..., d_{n-2}, d_{n-1})`.

  Args:
    pool_size: Shape of window that gets reduced to a single vector value.
        If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
        must be a tuple of length :math:`n-2`.
    strides: Offsets from the location of one window to the locations of
        neighboring windows along each axis. If specified, must be a tuple of
        the same length as `pool_size`. If None, then offsets of 1 along each
        window axis, :math:`(1, ..., 1)`, will be used.
    padding: 'VALID' or 'SAME'. If 'VALID', no padding is done, and only
        full windows get reduced; partial windows are discarded. If 'SAME',
        padding is added at array edges as needed to avoid partial windows
        but does not otherwise affect the selection of max values.

  Returns:
    N-dimensional array in which each valid (or padded-valid) window position
    in the input is reduced to / replaced by the max value from that window.
    An output array has the same number of dimensions as its input, but has
    fewer elements.
  """
  layer_name = f'MaxPool{pool_size}'.replace(' ', '')
  def f(x):
    return fastmath.max_pool(
        x, pool_size=pool_size, strides=strides, padding=padding)
  return Fn(layer_name, f)


def SumPool(pool_size=(2, 2), strides=None, padding='VALID'):
  """Reduces each multi-dimensional window to the sum of the window's values.

  Windows, as specified by `pool_size` and `strides`, involve all axes of an
  n-dimensional array except the first and last: :math:`(d_1, ..., d_{n-2})`
  from shape :math:`(d_0, d_1, ..., d_{n-2}, d_{n-1})`.

  Args:
    pool_size: Shape of window that gets reduced to a single vector value.
        If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
        must be a tuple of length :math:`n-2`.
    strides: Offsets from the location of one window to the locations of
        neighboring windows along each axis. If specified, must be a tuple of
        the same length as `pool_size`. If None, then offsets of 1 along each
        window axis, :math:`(1, ..., 1)`, will be used.
    padding: 'VALID' or 'SAME'. If 'VALID', no padding is done, and only
        full windows get reduced; partial windows are discarded. If 'SAME',
        padding is added at array edges as needed to avoid partial
        windows but does not otherwise affect the computation of sums.

  Returns:
    N-dimensional array in which each valid (or padded-valid) window position
    in the input is reduced to / replaced by the sum of values in that window.
    An output array has the same number of dimensions as its input, but has
    fewer elements.
  """
  layer_name = f'SumPool{pool_size}'.replace(' ', '')
  def f(x):
    return fastmath.sum_pool(
        x, pool_size=pool_size, strides=strides, padding=padding)
  return Fn(layer_name, f)


def AvgPool(pool_size=(2, 2), strides=None, padding='VALID'):
  """Reduces each multi-dimensional window to the mean of the window's values.

  Windows, as specified by `pool_size` and `strides`, involve all axes of an
  n-dimensional array except the first and last: :math:`(d_1, ..., d_{n-2})`
  from shape :math:`(d_0, d_1, ..., d_{n-2}, d_{n-1})`.

  Args:
    pool_size: Shape of window that gets reduced to a single vector value.
        If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
        must be a tuple of length :math:`n-2`.
    strides: Offsets from the location of one window to the locations of
        neighboring windows along each axis. If specified, must be a tuple of
        the same length as `pool_size`. If None, then offsets of 1 along each
        window axis, :math:`(1, ..., 1)`, will be used.
    padding: 'VALID' or 'SAME'. If 'VALID', no padding is done, and only
        full windows get reduced; partial windows are discarded. If 'SAME',
        padding is added at array edges as needed but is not counted in the
        computation of averages.

  Returns:
    N-dimensional array in which each valid (or padded-valid) window position
    in the input is reduced to / replaced by the mean of values in that window.
    An output array has the same number of dimensions as its input, but has
    fewer elements.
  """
  layer_name = f'AvgPool{pool_size}'.replace(' ', '')
  def f(x):
    return fastmath.avg_pool(
        x, pool_size=pool_size, strides=strides, padding=padding)
  return Fn(layer_name, f)
