#@title 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
#
# https://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.

"""
Explanation: Copyright 2021 The TensorFlow Authors.
End of explanation
"""


import tensorflow as tf
import tensorflow.compat.v1 as tf1

"""
Explanation: Migrate metrics and optimizers
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/guide/migrate/metrics_optimizers">
    <img src="https://www.tensorflow.org/images/tf_logo_32px.png" />
    View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/metrics_optimizers.ipynb">
    <img src="https://www.tensorflow.org/images/colab_logo_32px.png" />
    Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/metrics_optimizers.ipynb">
    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
    View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/metrics_optimizers.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>

In TF1, tf.metrics is the API namespace for all the metric functions. Each of the metrics is a function that takes label and prediction as input parameters and returns the corresponding metrics tensor as result. In TF2, tf.keras.metrics contains all the metric functions and objects. The Metric object can be used with tf.keras.Model and tf.keras.layers.layer to calculate metric values.
Setup
Let's start with a couple of necessary TensorFlow imports,
End of explanation
"""


features = [[1., 1.5], [2., 2.5], [3., 3.5]]
labels = [0, 0, 1]
eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]
eval_labels = [0, 1, 1]

"""
Explanation: and prepare some simple data for demonstration:
End of explanation
"""


def _input_fn():
  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)

def _eval_input_fn():
  return tf1.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

def _model_fn(features, labels, mode):
  logits = tf1.layers.Dense(2)(features)
  predictions = tf.math.argmax(input=logits, axis=1)
  loss = tf1.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  accuracy = tf1.metrics.accuracy(labels=labels, predictions=predictions)
  return tf1.estimator.EstimatorSpec(mode, 
                                     predictions=predictions,
                                     loss=loss, 
                                     train_op=train_op,
                                     eval_metric_ops={'accuracy': accuracy})

estimator = tf1.estimator.Estimator(model_fn=_model_fn)
estimator.train(_input_fn)

estimator.evaluate(_eval_input_fn)

"""
Explanation: TF1: tf.compat.v1.metrics with Estimator
In TF1, the metrics can be added to EstimatorSpec as the eval_metric_ops, and the op is generated via all the metrics functions defined in tf.metrics. You can follow the example to see how to use tf.metrics.accuracy.
End of explanation
"""


def mean_squared_error(labels, predictions):
  labels = tf.cast(labels, predictions.dtype)
  return {"mean_squared_error": 
          tf1.metrics.mean_squared_error(labels=labels, predictions=predictions)}

estimator = tf1.estimator.add_metrics(estimator, mean_squared_error)
estimator.evaluate(_eval_input_fn)

"""
Explanation: Also, metrics could be added to estimator directly via tf.estimator.add_metrics().
End of explanation
"""


dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

inputs = tf.keras.Input((2,))
logits = tf.keras.layers.Dense(2)(inputs)
predictions = tf.math.argmax(input=logits, axis=1)
model = tf.keras.models.Model(inputs, predictions)
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer, loss='mse', metrics=[tf.keras.metrics.Accuracy()])

model.evaluate(eval_dataset, return_dict=True)

"""
Explanation: TF2: Keras Metrics API with tf.keras.Model
In TF2, tf.keras.metrics contains all the metrics classes and functions. They are designed in a OOP style and integrate closely with other tf.keras API. All the metrics can be found in tf.keras.metrics namespace, and there is usually a direct mapping between tf.compat.v1.metrics with tf.keras.metrics. 
In the following example, the metrics are added in model.compile() method. Users only need to create the metric instance, without specifying the label and prediction tensor. The Keras model will route the model output and label to the metrics object.
End of explanation
"""


accuracy = tf.keras.metrics.Accuracy()

accuracy.update_state(y_true=[0, 0, 1, 1], y_pred=[0, 0, 0, 1])
accuracy.result().numpy()


accuracy.update_state(y_true=[0, 0, 1, 1], y_pred=[0, 0, 0, 0])
accuracy.update_state(y_true=[0, 0, 1, 1], y_pred=[1, 1, 0, 0])
accuracy.result().numpy()

"""
Explanation: With eager execution enabled, tf.keras.metrics.Metric instances can be directly used to evaluate numpy data or eager tensors. tf.keras.metrics.Metric objects are stateful containers. The metric value can be updated via metric.update_state(y_true, y_pred), and the result can be retrieved by metrics.result().
End of explanation
"""


