# Copyright 2018, 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.

from absl.testing import parameterized
import tensorflow as tf

from tensorflow_federated.python.core.backends.native import execution_contexts
from tensorflow_federated.python.core.impl.tensorflow_context import tensorflow_computation
from tensorflow_federated.python.simulation.datasets import client_data as cd


class CheckRandomSeedTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.named_parameters(
      ('integer1', 0),
      ('integer2', 1234),
      ('integer3', 2**32 - 1),
      ('sequence1', [0, 2, 323]),
      ('sequence2', [1024]),
      ('sequence3', [2**31, 2**32 - 1]),
      ('none', None),
  )
  def test_validate_does_not_raise_on_expected_inputs(self, seed):
    cd.check_numpy_random_seed(seed)

  @parameterized.named_parameters(
      ('integer1', -1),
      ('integer2', 2**40),
      ('integer3', 2**32),
      ('sequence1', [0, 2, -1]),
      ('sequence2', [2**33]),
      ('sequence3', [2**31, 2**32 - 1, -500]),
      ('string', 'bad_seed'),
  )
  def test_validate_raises_on_unexpected_inputs(self, seed):
    with self.assertRaises(cd.InvalidRandomSeedError):
      cd.check_numpy_random_seed(seed)


def create_concrete_client_data() -> cd.ConcreteClientData:
  """Creates a simple `ConcreteClientData` instance.

  The resulting `ClientData` has the following clients and datasets (written as
  lists):
  *   client `1`: [0]
  *   client `2`: [0, 1]
  *   client `3`: [0, 1, 2]

  Returns:
    A `ConcreteClientData` instance.
  """
  client_ids = ['1', '2', '3']

  def create_dataset_fn(client_id):
    num_examples = tf.strings.to_number(client_id, out_type=tf.int64)
    return tf.data.Dataset.range(num_examples)

  return cd.ClientData.from_clients_and_tf_fn(
      client_ids=client_ids, serializable_dataset_fn=create_dataset_fn
  )


def dataset_length(dataset):
  return dataset.reduce(0, lambda x, _: x + 1)


class TrainTestClientSplitTest(tf.test.TestCase, parameterized.TestCase):

  def get_even_odd_client_data(self):
    """Creates a `ClientData` where only clients with even IDs have data."""

    def create_dataset_fn(client_id):
      client_id_as_int = tf.strings.to_number(client_id, out_type=tf.int64)
      num_examples = 1 if client_id_as_int % 2 == 0 else 0
      return tf.data.Dataset.range(num_examples)

    client_ids = [str(x) for x in range(10)]
    return cd.ClientData.from_clients_and_tf_fn(
        client_ids=client_ids, serializable_dataset_fn=create_dataset_fn
    )

  def test_split_train_test_selects_nonempty_test_clients(self):
    # Only even client_ids have data:
    client_data = self.get_even_odd_client_data()

    train, test = cd.ClientData.train_test_client_split(
        client_data, num_test_clients=3
    )
    # Test that all clients end up in one of the two ClientData:
    self.assertCountEqual(
        client_data.client_ids, train.client_ids + test.client_ids
    )
    self.assertLen(test.client_ids, 3)
    for client_id in test.client_ids:
      self.assertEqual(int(client_id) % 2, 0)

    train, test = cd.ClientData.train_test_client_split(
        client_data, num_test_clients=5
    )
    self.assertLen(test.client_ids, 5)
    self.assertLen(train.client_ids, 5)

  def test_split_train_test_not_enough_nonempty_clients(self):
    client_data = self.get_even_odd_client_data()
    with self.assertRaisesRegex(ValueError, 'too many clients with no data.'):
      cd.ClientData.train_test_client_split(client_data, num_test_clients=6)

  def test_split_train_test_too_few_clients(self):
    client_data = self.get_even_odd_client_data()
    with self.assertRaisesRegex(ValueError, 'has only 10 clients.*11'):
      cd.ClientData.train_test_client_split(client_data, num_test_clients=11)

  def test_split_train_test_no_test_clients_requested(self):
    client_data = self.get_even_odd_client_data()
    with self.assertRaisesRegex(ValueError, 'Please specify'):
      cd.ClientData.train_test_client_split(client_data, num_test_clients=0)

  def test_split_train_test_fixed_seed(self):
    client_data = self.get_even_odd_client_data()

    train_0, test_0 = cd.ClientData.train_test_client_split(
        client_data, num_test_clients=3, seed=0
    )
    train_1, test_1 = cd.ClientData.train_test_client_split(
        client_data, num_test_clients=3, seed=0
    )

    self.assertEqual(train_0.client_ids, train_1.client_ids)
    self.assertEqual(test_0.client_ids, test_1.client_ids)

  @parameterized.named_parameters(
      ('integer1', 0),
      ('integer2', 1234),
      ('integer3', 2**32 - 1),
      ('sequence1', [0, 2, 323]),
      ('sequence2', [1024]),
      ('sequence3', [2**31, 2**32 - 1]),
      ('none', None),
  )
  def test_split_does_not_raise_on_expected_random_seed(self, seed):
    client_data = client_data = self.get_even_odd_client_data()
    cd.ClientData.train_test_client_split(
        client_data, num_test_clients=3, seed=seed
    )

  @parameterized.named_parameters(
      ('integer1', -1),
      ('integer2', 2**40),
      ('integer3', 2**32),
      ('sequence1', [0, 2, -1]),
      ('sequence2', [2**33]),
      ('sequence3', [2**31, 2**32 - 1, -500]),
      ('string', 'bad_seed'),
  )
  def test_split_raises_on_unexpected_random_seed(self, seed):
    client_data = client_data = self.get_even_odd_client_data()
    with self.assertRaises(cd.InvalidRandomSeedError):
      cd.ClientData.train_test_client_split(
          client_data, num_test_clients=3, seed=seed
      )


class ConcreteClientDataTest(tf.test.TestCase, parameterized.TestCase):

  def test_raises_on_tff_computation(self):
    @tensorflow_computation.tf_computation
    def dataset_fn():
      return 0

    with self.assertRaises(TypeError):
      cd.ClientData.from_clients_and_tf_fn(
          client_ids=[], serializable_dataset_fn=dataset_fn
      )

  def test_concrete_client_data_create_expected_datasets(self):
    client_data = create_concrete_client_data()
    self.assertEqual(
        client_data.element_type_structure,
        tf.TensorSpec(shape=(), dtype=tf.int64),
    )
    for i in client_data.client_ids:
      client_dataset = client_data.create_tf_dataset_for_client(i)
      self.assertEqual(dataset_length(client_dataset), int(i))

  def test_dataset_computation_lists_all_elements(self):
    client_data = create_concrete_client_data()
    for client_id in client_data.client_ids:
      expected_values = list(range(int(client_id)))
      client_dataset = client_data.dataset_computation(client_id)
      actual_values = list(client_dataset.as_numpy_iterator())
      self.assertEqual(expected_values, actual_values)

  def test_datasets_is_lazy(self):
    client_ids = [1, 2, 3]

    # Note: this is called once on initialization of ClientData
    # with client_ids[0] in order to get the element type.
    # After that, it should be called lazily when `next` is called
    # on a `.datasets()` iterator.
    called_count = 0

    def only_call_me_thrice(client_id):
      nonlocal called_count
      called_count += 1
      if called_count == 1:
        self.assertEqual(client_id, client_ids[0])
      if called_count > 3:
        raise Exception('called too many times')
      num_examples = client_id
      return tf.data.Dataset.range(num_examples)

    client_data = cd.ClientData.from_clients_and_tf_fn(
        client_ids=client_ids, serializable_dataset_fn=only_call_me_thrice
    )
    datasets_iter = client_data.datasets()
    next(datasets_iter)
    next(datasets_iter)
    with self.assertRaisesRegex(Exception, 'called too many times'):
      next(datasets_iter)

  def test_datasets_limit_count(self):
    client_data = create_concrete_client_data()
    dataset_list = list(client_data.datasets(limit_count=1))
    self.assertLen(dataset_list, 1)

  def test_datasets_doesnt_shuffle_client_ids_list(self):
    client_data = create_concrete_client_data()
    client_ids_copy = client_data.client_ids.copy()

    client_data.datasets()
    self.assertEqual(client_data.client_ids, client_ids_copy)
    client_data.datasets()
    self.assertEqual(client_data.client_ids, client_ids_copy)
    client_data.datasets()
    self.assertEqual(client_data.client_ids, client_ids_copy)

  @parameterized.named_parameters(
      ('integer1', 0),
      ('integer2', 1234),
      ('integer3', 2**32 - 1),
      ('sequence1', [0, 2, 323]),
      ('sequence2', [1024]),
      ('sequence3', [2**31, 2**32 - 1]),
      ('none', None),
  )
  def test_datasets_does_not_raise_on_expected_random_seed(self, seed):
    client_data = create_concrete_client_data()
    next(client_data.datasets(seed=seed))

  @parameterized.named_parameters(
      ('integer1', -1),
      ('integer2', 2**40),
      ('integer3', 2**32),
      ('sequence1', [0, 2, -1]),
      ('sequence2', [2**33]),
      ('sequence3', [2**31, 2**32 - 1, -500]),
      ('string', 'bad_seed'),
  )
  def test_datasets_raises_on_unexpected_random_seed(self, seed):
    client_data = create_concrete_client_data()
    with self.assertRaises(cd.InvalidRandomSeedError):
      next(client_data.datasets(seed=seed))

  def test_create_tf_dataset_from_all_clients(self):
    client_data = create_concrete_client_data()
    dataset = client_data.create_tf_dataset_from_all_clients()
    dataset_list = list(dataset.as_numpy_iterator())
    self.assertCountEqual(dataset_list, [0, 0, 0, 1, 1, 2])

  @parameterized.named_parameters(
      ('integer1', 0),
      ('integer2', 1234),
      ('integer3', 2**32 - 1),
      ('sequence1', [0, 2, 323]),
      ('sequence2', [1024]),
      ('sequence3', [2**31, 2**32 - 1]),
      ('none', None),
  )
  def test_create_from_all_does_not_raise_on_expected_random_seed(self, seed):
    client_data = create_concrete_client_data()
    client_data.create_tf_dataset_from_all_clients(seed=seed)

  @parameterized.named_parameters(
      ('integer1', -1),
      ('integer2', 2**40),
      ('integer3', 2**32),
      ('sequence1', [0, 2, -1]),
      ('sequence2', [2**33]),
      ('sequence3', [2**31, 2**32 - 1, -500]),
      ('string', 'bad_seed'),
  )
  def test_create_from_all_raises_on_unexpected_random_seed(self, seed):
    client_data = create_concrete_client_data()
    with self.assertRaises(cd.InvalidRandomSeedError):
      client_data.create_tf_dataset_from_all_clients(seed=seed)

  def test_dataset_from_large_client_list(self):
    client_ids = [str(x) for x in range(1_000_000)]

    def create_dataset(_):
      return tf.data.Dataset.range(100)

    client_data = cd.ClientData.from_clients_and_tf_fn(
        client_ids=client_ids, serializable_dataset_fn=create_dataset
    )
    # Ensure this completes within the test timeout without raising error.
    # Previous implementations caused this to take an very long time via Python
    # list -> generator -> list transformations.
    try:
      client_data.create_tf_dataset_from_all_clients(seed=42)
    except Exception as e:  # pylint: disable=broad-except
      self.fail(e)


class PreprocessClientDataTest(tf.test.TestCase, parameterized.TestCase):

  def test_preprocess_creates_expected_client_datasets(self):
    client_data = create_concrete_client_data()

    def preprocess_fn(dataset):
      return dataset.map(lambda x: 2 * x)

    preprocess_client_data = client_data.preprocess(preprocess_fn)
    for client_id in client_data.client_ids:
      expected_dataset = [2 * a for a in range(int(client_id))]
      actual_dataset = preprocess_client_data.create_tf_dataset_for_client(
          client_id
      )
      self.assertEqual(
          expected_dataset, list(actual_dataset.as_numpy_iterator())
      )

  def test_preprocess_with_take_one(self):
    client_data = create_concrete_client_data()
    preprocess_fn = lambda x: x.take(1)

    preprocess_client_data = client_data.preprocess(preprocess_fn)
    for client_id in client_data.client_ids:
      dataset = preprocess_client_data.create_tf_dataset_for_client(client_id)
      self.assertEqual(dataset_length(dataset), 1)

    self.assertLen(
        client_data.client_ids,
        dataset_length(
            preprocess_client_data.create_tf_dataset_from_all_clients()
        ),
    )

  def test_preprocess_creates_expected_client_datasets_with_dataset_comp(self):
    client_data = create_concrete_client_data()

    def preprocess_fn(dataset):
      return dataset.map(lambda x: 2 * x)

    preprocess_client_data = client_data.preprocess(preprocess_fn)
    for client_id in client_data.client_ids:
      expected_dataset = [2 * a for a in range(int(client_id))]
      actual_dataset = preprocess_client_data.dataset_computation(client_id)
      self.assertEqual(
          expected_dataset, list(actual_dataset.as_numpy_iterator())
      )

  def test_preprocess_creates_expected_amalgamated_dataset(self):
    client_data = create_concrete_client_data()

    def preprocess_fn(dataset):
      return dataset.map(lambda x: 2 * x)

    preprocess_client_data = client_data.preprocess(preprocess_fn)
    expected_amalgamated_dataset = [0, 0, 2, 0, 2, 4]
    actual_amalgamated_dataset = (
        preprocess_client_data.create_tf_dataset_from_all_clients()
    )
    self.assertCountEqual(
        expected_amalgamated_dataset,
        list(actual_amalgamated_dataset.as_numpy_iterator()),
    )

  def test_preprocess_raises_on_tff_computation(self):
    client_data = create_concrete_client_data()

    @tensorflow_computation.tf_computation
    def foo():
      return 1

    with self.assertRaises(cd.IncompatiblePreprocessFnError):
      client_data.preprocess(foo)


if __name__ == '__main__':
  execution_contexts.set_sync_local_cpp_execution_context()
  tf.test.main()
