# Copyright 2020, 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.
import collections

from absl.testing import absltest
import tensorflow as tf

from tensorflow_federated.python.core.api import computations
from tensorflow_federated.python.core.impl.federated_context import intrinsics
from tensorflow_federated.python.core.impl.types import computation_types
from tensorflow_federated.python.core.impl.types import placements
from tensorflow_federated.python.core.templates import iterative_process
from tensorflow_federated.python.simulation import iterative_process_compositions


def _create_whimsy_iterative_process():

  @computations.tf_computation()
  def init():
    return []

  @computations.tf_computation(init.type_signature.result)
  def next_fn(x):
    return x

  return iterative_process.IterativeProcess(initialize_fn=init, next_fn=next_fn)


def _create_federated_int_dataset_identity_iterative_process():

  @computations.tf_computation()
  def create_dataset():
    return tf.data.Dataset.range(5)

  @computations.federated_computation()
  def init():
    return intrinsics.federated_eval(create_dataset, placements.CLIENTS)

  @computations.federated_computation(init.type_signature.result)
  def next_fn(x):
    return x

  return iterative_process.IterativeProcess(initialize_fn=init, next_fn=next_fn)


def _create_stateless_int_dataset_reduction_iterative_process():

  @computations.tf_computation()
  def make_zero():
    return tf.cast(0, tf.int64)

  @computations.federated_computation()
  def init():
    return intrinsics.federated_eval(make_zero, placements.SERVER)

  @computations.tf_computation(computation_types.SequenceType(tf.int64))
  def reduce_dataset(x):
    return x.reduce(tf.cast(0, tf.int64), lambda x, y: x + y)

  @computations.federated_computation(
      (init.type_signature.result,
       computation_types.FederatedType(
           computation_types.SequenceType(tf.int64), placements.CLIENTS)))
  def next_fn(server_state, client_data):
    del server_state  # Unused
    return intrinsics.federated_sum(
        intrinsics.federated_map(reduce_dataset, client_data))

  return iterative_process.IterativeProcess(initialize_fn=init, next_fn=next_fn)


def _create_stateless_int_vector_unknown_dim_dataset_reduction_iterative_process(
):
  # Tests handling client data of unknown shape and summing to fixed shape.

  @computations.tf_computation()
  def make_zero():
    return tf.reshape(tf.cast(0, tf.int64), shape=[1])

  @computations.federated_computation()
  def init():
    return intrinsics.federated_eval(make_zero, placements.SERVER)

  @computations.tf_computation(
      computation_types.SequenceType(
          computation_types.TensorType(tf.int64, shape=[None])))
  def reduce_dataset(x):
    return x.reduce(
        tf.cast(tf.constant([0]), tf.int64), lambda x, y: x + tf.reduce_sum(y))

  @computations.federated_computation(
      computation_types.FederatedType(
          computation_types.TensorType(tf.int64, shape=[None]),
          placements.SERVER),
      computation_types.FederatedType(
          computation_types.SequenceType(
              computation_types.TensorType(tf.int64, shape=[None])),
          placements.CLIENTS))
  def next_fn(server_state, client_data):
    del server_state  # Unused
    return intrinsics.federated_sum(
        intrinsics.federated_map(reduce_dataset, client_data))

  return iterative_process.IterativeProcess(initialize_fn=init, next_fn=next_fn)


@computations.tf_computation(tf.string)
def int_dataset_computation(x):
  del x  # Unused
  return tf.data.Dataset.range(5)


@computations.tf_computation(tf.string)
def vector_int_dataset_computation(x):
  del x  # Unused
  return tf.data.Dataset.range(5).map(lambda x: tf.reshape(x, shape=[1]))


@computations.tf_computation(tf.string)
def float_dataset_computation(x):
  del x  # Unused
  return tf.data.Dataset.range(5, output_type=tf.float32)


@computations.tf_computation(tf.int32)
def int_identity(x):
  return x


@computations.federated_computation(
    tf.int32,
    computation_types.FederatedType(
        computation_types.SequenceType(tf.int64), placements.CLIENTS),
    tf.float32,
)
def test_int64_sequence_struct_computation(a, dataset, b):
  return a, dataset, b


@computations.federated_computation(
    tf.int32,
    computation_types.StructType([
        tf.int64,
        computation_types.at_clients(computation_types.SequenceType(tf.int64)),
        tf.float64,
    ]),
    tf.float32,
)
def test_int64_sequence_nested_struct_computation(a, dataset, b):
  return a, dataset, b


@computations.federated_computation(
    computation_types.StructType([
        computation_types.at_clients(computation_types.SequenceType(tf.int64)),
    ]),
    computation_types.at_clients(computation_types.SequenceType(tf.int64)),
)
def test_int64_sequence_multiple_matching_federated_types_computation(a, b):
  return a, b


@computations.federated_computation(
    computation_types.FederatedType(
        computation_types.SequenceType(tf.int64), placements.CLIENTS))
def test_int64_sequence_computation(dataset):
  del dataset
  return intrinsics.federated_value(5, placements.SERVER)


class ConstructDatasetsOnClientsComputationTest(absltest.TestCase):

  def test_raises_non_computation_dataset_comp(self):
    fn = lambda _: []
    with self.assertRaises(TypeError):
      iterative_process_compositions.compose_dataset_computation_with_computation(
          fn, test_int64_sequence_struct_computation)

  def test_raises_non_computation_outer_comp(self):
    non_comp = lambda x: x
    with self.assertRaises(TypeError):
      iterative_process_compositions.compose_dataset_computation_with_computation(
          int_dataset_computation, non_comp)

  def test_raises_computation_not_returning_dataset(self):
    with self.assertRaises(TypeError):
      iterative_process_compositions.compose_dataset_computation_with_computation(
          int_identity, test_int64_sequence_struct_computation)

  def test_raises_computation_no_dataset_parameter(self):
    no_dataset_comp = computations.federated_computation(
        lambda x: x, [tf.int32])
    with self.assertRaises(
        iterative_process_compositions.SequenceTypeNotFoundError):
      iterative_process_compositions.compose_dataset_computation_with_computation(
          int_dataset_computation, no_dataset_comp)

  def test_raises_mismatched_dataset_comp_return_type_and_sequence_type(self):
    with self.assertRaises(
        iterative_process_compositions.SequenceTypeNotAssignableError):
      iterative_process_compositions.compose_dataset_computation_with_computation(
          float_dataset_computation, test_int64_sequence_struct_computation)

  def test_mutates_comp_accepting_only_dataset(self):
    expected_new_next_type_signature = computation_types.FunctionType(
        parameter=computation_types.FederatedType(tf.string,
                                                  placements.CLIENTS),
        result=computation_types.FederatedType(tf.int32, placements.SERVER))
    new_comp = iterative_process_compositions.compose_dataset_computation_with_computation(
        int_dataset_computation, test_int64_sequence_computation)
    expected_new_next_type_signature.check_equivalent_to(