# Copyright 2021, 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.simulation.baselines import baseline_task
from tensorflow_federated.python.simulation.baselines import client_spec
from tensorflow_federated.python.simulation.baselines.stackoverflow import word_prediction_tasks


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

  def test_constructs_with_eval_client_spec(self):
    train_client_spec = client_spec.ClientSpec(
        num_epochs=2, batch_size=10, max_elements=3, shuffle_buffer_size=5
    )
    eval_client_spec = client_spec.ClientSpec(
        num_epochs=1, batch_size=2, max_elements=5, shuffle_buffer_size=10
    )
    baseline_task_spec = word_prediction_tasks.create_word_prediction_task(
        train_client_spec,
        vocab_size=10,
        eval_client_spec=eval_client_spec,
        use_synthetic_data=True,
    )
    self.assertIsInstance(baseline_task_spec, baseline_task.BaselineTask)

  def test_constructs_without_eval_client_spec(self):
    train_client_spec = client_spec.ClientSpec(
        num_epochs=2, batch_size=10, max_elements=3, shuffle_buffer_size=5
    )
    baseline_task_spec = word_prediction_tasks.create_word_prediction_task(
        train_client_spec, vocab_size=10, use_synthetic_data=True
    )
    self.assertIsInstance(baseline_task_spec, baseline_task.BaselineTask)

  @parameterized.named_parameters(
      ('sequence_length1', 1),
      ('sequence_length8', 8),
      ('sequence_length10', 10),
      ('sequence_length50', 50),
  )
  def test_constructs_with_different_sequence_lengths(self, sequence_length):
    train_client_spec = client_spec.ClientSpec(
        num_epochs=2, batch_size=10, max_elements=3, shuffle_buffer_size=5
    )
    baseline_task_spec = word_prediction_tasks.create_word_prediction_task(
        train_client_spec,
        vocab_size=10,
        sequence_length=sequence_length,
        use_synthetic_data=True,
    )
    self.assertIsInstance(baseline_task_spec, baseline_task.BaselineTask)

  @parameterized.named_parameters(
      ('sequence_length0', 0),
      ('sequence_length_minus1', -1),
      ('sequence_length_minus5', -5),
  )
  def test_raises_on_bad_sequence_lengths(self, sequence_length):
    train_client_spec = client_spec.ClientSpec(
        num_epochs=2, batch_size=10, max_elements=3, shuffle_buffer_size=5
    )
    with self.assertRaises(ValueError):
      word_prediction_tasks.create_word_prediction_task(
          train_client_spec,
          vocab_size=10,
          sequence_length=sequence_length,
          use_synthetic_data=True,
      )

  @parameterized.named_parameters(
      ('vocab_size1', 1),
      ('vocab_size8', 8),
      ('vocab_size10', 10),
      ('vocab_size50', 50),
  )
  def test_constructs_with_different_vocab_sizes(self, vocab_size):
    train_client_spec = client_spec.ClientSpec(
        num_epochs=2, batch_size=10, max_elements=3, shuffle_buffer_size=5
    )
    baseline_task_spec = word_prediction_tasks.create_word_prediction_task(
        train_client_spec, vocab_size=vocab_size, use_synthetic_data=True
    )
    self.assertIsInstance(baseline_task_spec, baseline_task.BaselineTask)

  @parameterized.named_parameters(
      ('vocab_size0', 0),
      ('vocab_size_minus1', -1),
      ('vocab_size_minus5', -5),
  )
  def test_raises_on_bad_vocab_size(self, vocab_size):
    train_client_spec = client_spec.ClientSpec(
        num_epochs=2, batch_size=10, max_elements=3, shuffle_buffer_size=5
    )
    with self.assertRaises(ValueError):
      word_prediction_tasks.create_word_prediction_task(
          train_client_spec, vocab_size=vocab_size, use_synthetic_data=True
      )

  def test_model_is_compatible_with_preprocessed_data(self):
    train_client_spec = client_spec.ClientSpec(num_epochs=1, batch_size=10)
    baseline_task_spec = word_prediction_tasks.create_word_prediction_task(
        train_client_spec, use_synthetic_data=True
    )
    centralized_dataset = (
        baseline_task_spec.datasets.get_centralized_test_data()
    )
    sample_batch = next(iter(centralized_dataset))
    model = baseline_task_spec.model_fn()
    model.forward_pass(sample_batch)


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