# 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.
"""Tests for FunctionModel."""

import collections

import numpy as np
import tensorflow as tf

from tensorflow_federated.python.core.backends.native import execution_contexts
from tensorflow_federated.python.core.impl.types import type_conversions
from tensorflow_federated.python.learning import model as model_lib
from tensorflow_federated.python.learning.metrics import aggregator
from tensorflow_federated.python.learning.models import functional
from tensorflow_federated.python.tensorflow_libs import variable_utils


def initial_weights():
  """Returns lists of trainable variables and non-trainable variables."""
  trainable_variables = (np.asarray([[0.0, 0.0, 0.0]], dtype=np.float32),
                         np.asarray([0.0], dtype=np.float32))
  non_trainable_variables = ()
  return (trainable_variables, non_trainable_variables)


@tf.function
def predict_on_batch(model_weights, x, training):
  """Test predict_on_batch implementing linear regression."""
  trainable = model_weights[0]
  w, b = trainable
  # For the sake of testing, only add the bias term when training so that
  # we get different outputs.
  if training:
    return tf.matmul(x, w, transpose_b=True) + b
  else:
    return tf.matmul(x, w, transpose_b=True)


@tf.function
def forward_pass(model_weights, batch_input, training):
  """Test forward_pass implementing linear regression on MSE."""
  x, y = batch_input
  predictions = predict_on_batch(model_weights, x, training)
  residuals = predictions - y
  num_examples = tf.shape(predictions)[0]
  total_loss = tf.reduce_sum(tf.pow(residuals, 2.))
  average_loss = total_loss / tf.cast(num_examples, tf.float32)
  return model_lib.BatchOutput(
      loss=average_loss, predictions=predictions, num_examples=num_examples)


def create_test_keras_model():
  return tf.keras.models.Sequential([
      tf.keras.layers.InputLayer(input_shape=(1,)),
      tf.keras.layers.Dense(
          1, kernel_initializer='zeros', bias_initializer='zeros')
  ])


def create_test_dataset():
  """Create a test dataset."""

  def preprocess(ds):

    def generate_example(i, t):
      del t  # Unused.
      features = tf.random.stateless_uniform(shape=[3], seed=(0, i))
      label = tf.expand_dims(
          tf.reduce_sum(features * tf.constant([1.0, 2.0, 3.0])), axis=-1) + 5.0
      return (features, label)

    return ds.map(generate_example).batch(5, drop_remainder=True)

  num_examples = 25
  return preprocess(tf.data.Dataset.range(num_examples).enumerate())


class FunctionalTest(tf.test.TestCase):