# 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.
"""An example of personalization strategy."""

import collections
from collections.abc import Callable
from typing import Any, Optional

import tensorflow as tf
import tensorflow_federated as tff

_OptimizerFn = Callable[[], tf.keras.optimizers.Optimizer]
_PersonalizeFn = Callable[
    [tff.learning.models.VariableModel, tf.data.Dataset, tf.data.Dataset, Any],
    collections.OrderedDict[str, tf.Tensor],
]
_EVAL_BATCH_SIZE = 1  # Batch size used when evaluating a dataset.
_SHUFFLE_BUFFER_SIZE = 1000  # Buffer size used when shuffling a dataset.


def build_personalize_fn(
    optimizer_fn: _OptimizerFn,
    batch_size: int,
    num_epochs: int,
    num_epochs_per_eval: int,
    shuffle: bool = True,
) -> _PersonalizeFn:
  """Builds a `tf.function` that represents a personalization strategy.

  The returned `tf.function` represents the optimization algorithm to run on
  a client in order to personalize a given model. It takes a
  `tff.learning.models.VariableModel` (with weights already initialized to the
  desired initial model weights), an unbatched training dataset, an unbatched
  test dataset, and an optional `context` (e.g., extra datasets) as input,
  trains a personalized model on the training dataset for `num_epochs`,
  evaluates the model on the test dataset every `num_epochs_per_eval`, and
  returns the evaluation metrics. The evaluation metrics are computed by
  `evaluate_fn` (see its docstring for more details).

  This builder function only serves as an example. Customers are allowed to
  write any personalization strategy as long as it satisfies the function
  signature specified by `_PERSONALIZE_FN_TYPE`.

  Args:
    optimizer_fn: A no-argument function that returns a
      `tf.keras.optimizers.Optimizer`.
    batch_size: An `int` specifying the batch size used in training.
    num_epochs: An `int` specifying the number of epochs used in training a
      personalized model.
    num_epochs_per_eval: An `int` specifying the frequency that a personalized
      model gets evaluated during the process of training.
    shuffle: A `bool` specifying whether to shuffle train data in every epoch.

  Returns:
    A `tf.function` that trains a personalized model, evaluates the model every
    `num_epochs_per_eval` epochs, and returns the evaluation metrics.
  """
  # Create the `optimizer` here instead of inside the `tf.function` below,
  # because a `tf.function` generally does not allow creating new variables.
  optimizer = optimizer_fn()

  @tf.function
  def personalize_fn(
      model: tff.learning.models.VariableModel,
      train_data: tf.data.Dataset,
      test_data: tf.data.Dataset,
      context: Optional[Any] = None,
  ) -> collections.OrderedDict[str, Any]:
    """A personalization strategy that trains a model and returns the metrics.

    Args:
      model: A `tff.learning.models.VariableModel`.
      train_data: An unbatched `tf.data.Dataset` used for training.
      test_data: An unbatched `tf.data.Dataset` used for evaluation.
      context: An optional object (e.g., extra dataset) used in personalization.

    Returns:
      An `OrderedDict` that maps metric names to `tf.Tensor`s or structures of
      `tf.Tensor`s containing the training and evaluation metrics.
    """
    del context  # This example does not use extra context.

    def train_one_batch(num_examples_sum, batch):
      """Run gradient descent on a batch."""
      with tf.GradientTape() as tape:
        output = model.forward_pass(batch)
      grads = tape.gradient(output.loss, model.trainable_variables)
      optimizer.apply_gradients(
          zip(
              tf.nest.flatten(grads), tf.nest.flatten(model.trainable_variables)
          )
      )
      # Update the number of examples.
      return num_examples_sum + output.num_examples

    # Start training.
    training_state = 0  # Number of total examples used in training.
    metrics_dict = collections.OrderedDict()
    for epoch_idx in range(1, num_epochs + 1):
      if shuffle:
        train_data = train_data.shuffle(_SHUFFLE_BUFFER_SIZE)
      training_state = train_data.batch(batch_size).reduce(
          initial_state=training_state, reduce_func=train_one_batch
      )
      # Evaluate the trained model every `num_epochs_per_eval` epochs.
      if (epoch_idx % num_epochs_per_eval == 0) or (epoch_idx == num_epochs):
        metrics_dict[f'epoch_{epoch_idx}'] = evaluate_fn(model, test_data)

    # Save the training statistics.
    metrics_dict['num_train_examples'] = training_state

    return metrics_dict

  return personalize_fn


@tf.function
def evaluate_fn(
    model: tff.learning.models.VariableModel, dataset: tf.data.Dataset
) -> collections.OrderedDict[str, tf.Tensor]:
  """Evaluates a model on the given dataset.

  The returned metrics include those given by
  `model.report_local_unfinalized_metrics`. These are specified by the `loss`
  and `metrics` arguments when the model is created by
  `tff.learning.models.from_keras_model`. The returned metrics also contain an
  integer metric with name 'num_test_examples'.

  Args:
    model: A `tff.learning.models.VariableModel` created by
      `tff.learning.models.from_keras_model`.
    dataset: An unbatched `tf.data.Dataset`.

  Returns:
    An `OrderedDict` of metric names to scalar `tf.Tensor`s.
  """
  # Resets the model's local variables. This is necessary because
  # `model.report_local_unfinalized_metrics()` aggregates the metrics from *all*
  # previous calls to `forward_pass` (which include the metrics computed in
  # training).
  # Resetting ensures that the returned metrics are computed on test data.
  # Similar to the `reset_states` method of `tf.keras.metrics.Metric`.
  model.reset_metrics()

  def reduce_fn(num_examples_sum, batch):
    output = model.forward_pass(batch, training=False)
    return num_examples_sum + output.num_examples

  # Runs `reduce_fn` over the input dataset. The final metrics can be accessed
  # by `model.report_local_unfinalized_metrics()`.
  num_examples_sum = dataset.batch(_EVAL_BATCH_SIZE).reduce(
      initial_state=0, reduce_func=reduce_fn
  )
  eval_metrics = collections.OrderedDict()
  eval_metrics['num_test_examples'] = num_examples_sum
  local_outputs = model.report_local_unfinalized_metrics()
  # Postprocesses the metric values. This is needed because the values returned
  # by `model.report_local_unfinalized_metrics()` are values of the state
  # variables in each `tf.keras.metrics.Metric`. These values should be
  # processed in the same way as the `result()` method of a
  # `tf.keras.metrics.Metric`.
  for name, metric in local_outputs.items():
    if not isinstance(metric, list):
      raise TypeError(
          'The metric value returned by `report_local_unfinalized_metrics` is '
          'expected to be a list, but found an instance of '
          f'{type(metric)}. Please check that your TFF model is '
          'built from a keras model.'
      )
    if len(metric) == 2:
      # The loss and accuracy metrics used in this p13n example has two values:
      # one represents `sum`, and the other represents `count`.
      eval_metrics[name] = metric[0] / metric[1]
    elif len(metric) == 1:
      eval_metrics[name] = metric[0]
    else:
      raise ValueError(
          'The metric value returned by `report_local_unfinalized_metrics` '
          'is expected to be a list of length 1 or 2, but found '
          f'one with length {len(metric)}.'
      )
  return eval_metrics
