# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Tensorflow based mnist model and input datasets helper."""

_NUM_CLASSES = 10


def get_input_datasets(tf, use_bfloat16=False):
  """Downloads the MNIST dataset and creates train and eval dataset objects.

  Args:
    tf: version specific tensorflow (either tf1 or tf2).
    use_bfloat16: Boolean to determine if input should be cast to bfloat16

  Returns:
    Train dataset, eval dataset, test dataset and input shape.

  """
  # input image dimensions
  img_rows, img_cols = 28, 28
  cast_dtype = tf.bfloat16 if use_bfloat16 else tf.float32

  # the data, split between train and test sets
  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

  if tf.keras.backend.image_data_format() == 'channels_first':
    x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
    x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
    input_shape = (1, img_rows, img_cols)
  else:
    x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
    x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
    input_shape = (img_rows, img_cols, 1)

  x_train = x_train.astype('float32')
  x_test = x_test.astype('float32')
  x_train /= 255
  x_test /= 255

  # convert class vectors to binary class matrices
  y_train = tf.keras.utils.to_categorical(y_train, _NUM_CLASSES)
  y_test = tf.keras.utils.to_categorical(y_test, _NUM_CLASSES)

  # train dataset
  train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
  train_ds = train_ds.repeat()
  train_ds = train_ds.map(lambda x, y: (tf.cast(x, cast_dtype), y))
  train_ds = train_ds.batch(64, drop_remainder=True)

  # eval dataset
  eval_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
  eval_ds = eval_ds.repeat()
  eval_ds = eval_ds.map(lambda x, y: (tf.cast(x, cast_dtype), y))
  eval_ds = eval_ds.batch(64, drop_remainder=True)

  # test dataset
  test_ds = tf.data.Dataset.from_tensor_slices((x_test))
  test_ds = test_ds.repeat()
  test_ds = test_ds.map(lambda x: tf.cast(x, cast_dtype))
  test_ds = test_ds.batch(64, drop_remainder=True)

  return train_ds, eval_ds, test_ds, input_shape


def get_model(tf, input_shape):
  """Builds a Sequential CNN model to recognize MNIST digits.

  Args:
   tf: version specific tensorflow (either tf1 or tf2).
   input_shape: Shape of the input depending on the `image_data_format`.

  Returns:
    a Keras model

  """
  # Define a CNN model to recognize MNIST digits.
  model = tf.keras.models.Sequential()
  model.add(
      tf.keras.layers.Conv2D(
          32, kernel_size=(3, 3), activation='relu', input_shape=input_shape))
  model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
  model.add(tf.keras.layers.Dropout(0.25))
  model.add(tf.keras.layers.Flatten())
  model.add(tf.keras.layers.Dense(128, activation='relu'))
  model.add(tf.keras.layers.Dropout(0.5))
  model.add(tf.keras.layers.Dense(_NUM_CLASSES, activation='softmax'))
  return model
