# Copyright 2019, The TensorFlow Federated 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.
#
# pytype: skip-file
# This modules disables the Pytype analyzer, see
# https://github.com/tensorflow/federated/blob/main/docs/pytype.md for more
# information.
"""A library of construction functions for building block structures."""

import functools
import random
import string
from typing import AbstractSet, Any, Callable, Iterator, List, Optional, Sequence, Tuple, Union

import tensorflow as tf

from tensorflow_federated.proto.v0 import computation_pb2 as pb
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.common_libs import serialization_utils
from tensorflow_federated.python.common_libs import structure
from tensorflow_federated.python.core.impl.compiler import building_blocks
from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
from tensorflow_federated.python.core.impl.compiler import transformation_utils
from tensorflow_federated.python.core.impl.types import computation_types
from tensorflow_federated.python.core.impl.types import placements
from tensorflow_federated.python.core.impl.types import type_analysis
from tensorflow_federated.python.core.impl.types import type_conversions
from tensorflow_federated.python.core.impl.types import type_serialization
from tensorflow_federated.python.core.impl.types import type_transformations
from tensorflow_federated.python.core.impl.utils import tensorflow_utils

Index = Union[str, int]
Path = Union[Index, Tuple[Index, ...]]


def select_output_from_lambda(
    comp: building_blocks.Lambda,
    paths: Union[Path, List[Path]]) -> building_blocks.Lambda:
  """Constructs a new function with result of selecting `paths` from `comp`.

  Args:
    comp: Lambda computation with result type `tff.StructType` from which we
      wish to select the sub-results at `paths`.
    paths: Either a `Path` or list of `Path`s specifying the indices we wish to
      select from the result of `comp`. Each path must be a `tuple` of `str` or
      `int` indices from which to select an output. If `paths` is a list, the
      returned computation will have a `tff.StructType` result holding each of
      the specified selections.

  Returns:
    A version of `comp` with result value the selection from the result of
    `comp` specified by `paths`.
  """
  comp.check_lambda()
  comp.type_signature.result.check_struct()

  def _select_path(result, path: Path):
    if not isinstance(path, tuple):
      path = (path,)
    for index in path:
      if result.is_struct():
        result = result[index]
      elif isinstance(index, str):
        result = building_blocks.Selection(result, name=index)
      elif isinstance(index, int):
        result = building_blocks.Selection(result, index=index)
      else:
        raise TypeError('Invalid selection type: expected `str` or `int`, '
                        f'found value `{index}` of type `{type(index)}`.')
    return result

  if isinstance(paths, list):
    # Avoid duplicating `comp.result` by binding it to a local.
    result_name = next(unique_name_generator(comp))
    result_ref = building_blocks.Reference(result_name,
                                           comp.result.type_signature)
    elements = [_select_path(result_ref, path) for path in paths]
    result = building_blocks.Block([(result_name, comp.result)],
                                   building_blocks.Struct(elements))
  else:
    result = _select_path(comp.result, paths)
  return building_blocks.Lambda(comp.parameter_name, comp.parameter_type,
                                result)


def unique_name_generator(comp: building_blocks.ComputationBuildingBlock,
                          prefix: str = '_var') -> Iterator[str]:
  """Yields a new unique name that does not exist in `comp`.

  Args:
    comp: The compuation building block to use as a reference.
    prefix: The prefix to use when generating unique names. If `prefix` is
      `None` or if `comp` contains any name with this prefix, then a unique
      prefix will be generated from random lowercase ascii characters.
  """
  if comp is not None:
    names = transformation_utils.get_unique_names(comp)
  else:
    names = set()
  while prefix is None or any(n.startswith(prefix) for n in names):
    characters = string.ascii_lowercase
    prefix = '_{}'.format(''.join(random.choice(characters) for _ in range(3)))
  index = 1
  while True:
    yield '{}{}'.format(prefix, index)
    index += 1


@functools.lru_cache()
def create_compiled_no_arg_empty_tuple_computation(
) -> building_blocks.CompiledComputation:
  """Returns graph representing a function that returns an empty tuple.

  Returns:
    An instance of `building_blocks.CompiledComputation`, a noarg function
    which returns an empty tuple.
  """
  proto, type_signature = tensorflow_computation_factory.create_empty_tuple()
  return building_blocks.CompiledComputation(
      proto, type_signature=type_signature)


@functools.lru_cache()
def create_compiled_empty_tuple() -> building_blocks.Call:
  """Returns called graph representing the empty tuple.

  Returns:
    An instance of `building_blocks.Call`, calling a noarg function
    which returns an empty tuple. This function is an instance of
    `building_blocks.CompiledComputation`.
  """
  compiled = create_compiled_no_arg_empty_tuple_computation()
  return building_blocks.Call(compiled, None)


@functools.lru_cache()
def create_identity(
    type_signature: computation_types.Type) -> building_blocks.Lambda:
  return building_blocks.Lambda(
      'id_arg', type_signature,
      building_blocks.Reference('id_arg', type_signature))


@functools.lru_cache()
def create_compiled_identity(
    type_signature: computation_types.Type,
    name: Optional[str] = None) -> building_blocks.CompiledComputation:
  """Creates CompiledComputation representing identity function.

  Args:
    type_signature: A `computation_types.Type`.
    name: An optional string name to use as the name of the computation.

  Returns:
    An instance of `building_blocks.CompiledComputation`
    representing the identity function taking an argument of type
    `type_signatu
re` and returning the same value.

  Raises:
    TypeError: If `type_signature` contains any types which cannot appear in
      TensorFlow bindings.
  """
  proto, function_type = tensorflow_computation_factory.create_identity(
      type_signature)
  return building_blocks.CompiledComputation(
      proto, name, type_signature=function_type)


class SelectionSpec(object):
  """Data class representing map from input tuple 
to selection of result.

  Attributes:
    tuple_index: The index of the source of the selection sequence in the
      desired result of the generated TensorFlow. If this `SelectionSpec`
      appears at index i of a list of `SelectionSpec`s, index j is the source
      for the result of the generated function at index i.
    selection_sequence: A list or tuple representing the selections to make from
      `tuple_index`, so that the list `[0]` for example would represent the
      output is the 0th element of `tuple_index`, while `[0, 0]` would represent
      that the output is the 0th element of the 0th element of `tuple_index`.
  """

  def __init__(self, tuple_index: int, selection_sequence: Sequence[int]):
    self._tuple_index = tuple_index
    self._selection_sequence = selection_sequence

  @property
  def tuple_index(self):
    return self._tuple_index

  @property
  def selection_sequence(self):
    return self._selection_se