# Lint as: python3
# 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.
"""An example of an MNIST model function for use with TensorFlow Federated."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

import tensorflow as tf
import tensorflow_federated as tff


class NumExamplesCounter(tf.keras.metrics.Sum):
  """A `tf.keras.metrics.Metric` that counts the number of examples seen."""

  def __init__(self, name='num_examples', dtype=tf.int64):  # pylint: disable=useless-super-delegation
    super(NumExamplesCounter, self).__init__(name, dtype)

  def update_state(self, y_true, y_pred, sample_weight=None):
    return super(NumExamplesCounter,
                 self).update_state(tf.shape(y_pred)[0], sample_weight)


def create_simple_keras_model(learning_rate=0.1):
  """Returns an instance of `tf.Keras.Model` with just one dense layer.

  Args:
    learning_rate: The learning rate to use with the SGD optimizer.

  Returns:
    An instance of `tf.Keras.Model`.
  """
  model = tf.keras.models.Sequential([
      tf.keras.layers.Flatten(input_shape=(784,)),
      tf.keras.layers.Dense(10, tf.nn.softmax, kernel_initializer='zeros')
  ])

  model.compile(
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      optimizer=tf.keras.optimizers.SGD(learning_rate),
      metrics=[
          tf.keras.metrics.SparseCategoricalAccuracy(),
          NumExamplesCounter()
      ])
  return model


def keras_dataset_from_emnist(dataset):
  """Converts `dataset` for use with the output of `create_simple_keras_model`.

  Args:
    dataset: An instance of `tf.data.Dataset` to read from.

  Returns:
    An instance of `tf.data.Dataset` after conversion.
  """

  def map_fn(example):
    return collections.OrderedDict([('x', tf.reshape(example['pixels'], [-1])),
                                    ('y', example['label'])])

  return dataset.map(map_fn)


def create_keras_model(compile_model=False):
  """Returns an instance of `tf.keras.Model` for use with the MNIST example.

  This code is based on the following target, which unfortunately cannot be
  imported as it is a Python binary, not a library:

  https://github.com/tensorflow/models/blob/master/official/mnist/mnist.py

  Args:
    compile_model: If True, compile the model with a basic optimizer and loss.

  Returns:
    A `tf.keras.Model`.
  """
  # TODO(b/120157713): Find a way to import this code.
  data_format = 'channels_last'
  input_shape = [28, 28, 1]
  l = tf.keras.layers
  initializer = tf.keras.initializers.RandomNormal(seed=0)
  max_pool = l.MaxPooling2D((2, 2), (2, 2),
                            padding='same',
                            data_format=data_format)
  model = tf.keras.Sequential([
      l.Reshape(target_shape=input_shape, input_shape=(28 * 28,)),
      l.Conv2D(
          32,
          5,
          padding='same',
          data_format=data_format,
          activation=tf.nn.relu,
          kernel_initializer=initializer), max_pool,
      l.Conv2D(
          64,
          5,
          padding='same',
          data_format=data_format,
          activation=tf.nn.relu,
          kernel_initializer=initializer), max_pool,
      l.Flatten(),
      l.Dense(1024, activation=tf.nn.relu, kernel_initializer=initializer),
      l.Dropout(0.4, seed=1),
      l.Dense(10, kernel_initializer=initializer)
  ])
  if compile_model:
    model.compile(
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        optimizer=tf.keras.optimizers.SGD(learning_rate=0.1))
  return model


Batch = collections.namedtuple('Batch', ['x', 'y'])  # pylint: disable=invalid-name


def create_random_batch():
  """Returns an instance of `Batch` populated with random tensors."""
  return Batch(
      x=tf.random.uniform(tf.TensorShape([1, 784]), dtype=tf.float32),
      y=tf.constant(1, dtype=tf.int64, shape=[1, 1]))


def model_fn():
  """Constructs the MNIST model wrapped for use with TensorFlow Federated.

  The model constructed by this function can be passed as an argument to
  `tff.learning.build_federated_averaging_process` to create a federated
  training process.

  Returns:
    An instance of `tff.learning.Model` that represents a trainable model.
  """
  keras_model = create_keras_model(compile_model=True)
  dummy_batch = create_random_batch()
  return tff.learning.from_compiled_keras_model(keras_model, dummy_batch)

tensorflow.TensorShape tensorflow.constant tensorflow.shape tensorflow.keras.losses.SparseCategoricalCrossentropy tensorflow.keras.layers.Dense tensorflow.reshape tensorflow.keras.metrics.SparseCategoricalAccuracy tensorflow.keras.initializers.RandomNormal tensorflow.keras.layers.Flatten tensorflow.keras.optimizers.SGD
[(129, 'collections.namedtuple', 'collections.namedtuple', (['"""Batch"""', "['x', 'y']"], {}), False, 'import collections\n'), (97, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'seed': '(0)'}), True, 'import tensorflow as tf\n'), (151, 'tensorflow_federated.learning.from_compiled_keras_model', 'tff.learning.from_compiled_keras_model', (['keras_model', 'dummy_batch'], {}), True, 'import tensorflow_federated as tff\n'), (48, 'tensorflow.keras.layers.Flatten', 'tf.keras.layers.Flatten', ([], {'input_shape': '(784,)'}), True, 'import tensorflow as tf\n'), (49, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(10)', 'tf.nn.softmax'], {'kernel_initializer': '"""zeros"""'}), True, 'import tensorflow as tf\n'), (53, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {}), True, 'import tensorflow as tf\n'), (54, 'tensorflow.keras.optimizers.SGD', 'tf.keras.optimizers.SGD', (['learning_rate'], {}), True, 'import tensorflow as tf\n'), (136, 'tensorflow.constant', 'tf.constant', (['(1)'], {'dtype': 'tf.int64', 'shape': '[1, 1]'}), True, 'import tensorflow as tf\n'), (35, 'tensorflow.shape', 'tf.shape', (['y_pred'], {}), True, 'import tensorflow as tf\n'), (56, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (124, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {}), True, 'import tensorflow as tf\n'), (125, 'tensorflow.keras.optimizers.SGD', 'tf.keras.optimizers.SGD', ([], {'learning_rate': '(0.1)'}), True, 'import tensorflow as tf\n'), (135, 'tensorflow.TensorShape', 'tf.TensorShape', (['[1, 784]'], {}), True, 'import tensorflow as tf\n'), (73, 'tensorflow.reshape', 'tf.reshape', (["example['pixels']", '[-1]'], {}), True, 'import tensorflow as tf\n')]