# Copyright 2020 DeepMind Technologies Limited. 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.
# ==============================================================================
"""Helper functions for template rendering."""

import collections
import itertools


ReductionAxes = collections.namedtuple('ReductionAxes',
                                       ['n_axes', 'all_axes', 'reduction_axes',
                                        'remaining_n_axes', 'remaining_axes'])

TransposeAxes = collections.namedtuple('TransposeAxes',
                                       ['n_axes', 'all_axes', 'transpose_axes',
                                        'result_axes'])


def axis_list(n_axes: int, reverse=False) -> str:
  """Returns a comma-separated list of axis TypeVar short names.

  Args:
    n_axes: Maximum number of axes to include in the list.
            n_axes=1 -> 'A1', n_axes=2 -> 'A1, A2', etc.
    reverse: If False, the returned list starts from A1 and counts up to An.
             If True, starts at An and counts down to A1.
  Returns:
    A string containing the list of axes.

  For example, get_axis_list(2) -> 'A1, A2'.
  """
  axes = range(1, n_axes + 1)
  if reverse:
    axes = reversed(axes)
  return ', '.join(f'A{i}' for i in axes)


def jax_array_type(n_axes: int) -> str:
  """Returns the generic JAX array type, parameterised by a number of axes.

  For example, get_jax_array_type(2) -> 'Array2[A1, A2]'.

  Args:
    n_axes: Rank of array type to return.
  Returns:
    A string containing array type.
  """
  return f'Array{n_axes}[{axis_list(n_axes)}]'


def transpose_axes(n_axes: int):
  """A generator that yields input and output axes of transpose.

  Args:
    n_axes: Rank of array whose possible transposes to consider.

  Yields:
    A `TransposeAxes` object for each possible transpose.

  For example, calculate_transpose_axes(2) would yield `TransposeAxes` objects
  encoding:
    Transpose shape [A1, A2] with axes=[0, 1] -> Shape[A1, A2]
                                       [1, 0] -> Shape[A2, A1]
  """
  assert n_axes >= 1

  # [A1, A2, ..., An]
  all_axes = list(range(1, n_axes + 1))
  all_axes_str = [f'A{i}' for i in all_axes]
  all_axes_str = ', '.join(all_axes_str)

  for transpose_axes in itertools.permutations(range(n_axes)):
    transpose_axes_str = (f'L{i}' for i in transpose_axes)
    transpose_axes_str = ', '.join(transpose_axes_str)
    transpose_axes_str = f'Tuple[{transpose_axes_str}]'

    result_axes = (all_axes[i] for i in transpose_axes)
    if result_axes:
      result_axes_str = (f'A{i}' for i in result_axes)
      result_axes_str = ', '.join(result_axes_str)
    else:
      result_axes_str = ''

    yield TransposeAxes(n_axes=n_axes,
                        all_axes=all_axes_str,
                        transpose_axes=transpose_axes_str,
                        result_axes=result_axes_str)


def reduction_axes(n_axes: int):
  """A generator that yields input and output axes of reduction operations.

  Args:
    n_axes: Rank of array whose possible reductions to consider.

  Yields:
    A `ReductionAxes` object for each possible reduction (where axes
    increase in value - e.g. we don't consider `axes=(1, 0)` - to cut down
    on the size of the stubs generated).

  For example, calculate_reduction_axes(2) would yield `ReductionAxes` objects
  encoding:
    Reduce shape [A1, A2] over axes    0 -> shape[A2]
                                       1 -> shape[A1]
                                      -1 -> shape[A1]
                                    0, 1 -> shape[]
  """
  assert n_axes >= 1
  final_axis = n_axes - 1

  # [A1, A2, ..., An]
  all_axes_str = [f'A{i}' for i in range(1, n_axes + 1)]
  all_axes_str = ', '.join(all_axes_str)

  n_reduction_axes_iter = range(1, n_axes + 1)
  for n_reduction_axes in n_reduction_axes_iter:

    # First, make a list of all possible permutations of reduction axes.
    reduction_axes_list = []
    for reduction_axes_combo in itertools.combinations(
        range(n_axes),
        n_reduction_axes
    ):
      # Thanks to `combinations`, `reduction_axes_combo` is already sorted.
      reduction_axes_list.append(reduction_axes_combo)
      # Also consider permutations where we refer to the final axis as '-1'.
      if final_axis in reduction_axes_combo:
        reduction_axes_list.append((*reduction_axes_combo[:-1], -1))

    # Second, for each permutation of reduction axes, yield a ReductionAxes.
    for reduction_axes in reduction_axes_list:
      reduction_axes_str = (
          f'L{i}' if i != -1 else 'LN1'
          for i in reduction_axes
      )
      reduction_axes_str = ', '.join(reduction_axes_str)
      if len(reduction_axes) != 1:
        reduction_axes_str = f'Tuple[{reduction_axes_str}]'

      concrete_reduction_axes = [
          (n_axes - 1) if axis == -1 else axis
          for axis in reduction_axes
      ]
      remaining_axes = set(range(n_axes)) - set(concrete_reduction_axes)
      remaining_axes = sorted(tuple(remaining_axes))
      remaining_n_axes = len(remaining_axes)
      if remaining_axes:
        remaining_axes_str = (f'A{i + 1}' for i in remaining_axes)
        remaining_axes_str = ', '.join(remaining_axes_str)
        remaining_axes_str = ', ' + remaining_axes_str
      else:
        remaining_axes_str = ''

      yield ReductionAxes(n_axes,
                          all_axes_str,
                          reduction_axes_str,
                          remaining_n_axes,
                          remaining_axes_str)
