# Lint as: python3
# Copyright 2020 The Orbit 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.
# ==============================================================================
"""An abstraction that users can easily handle their custom training loops."""

import abc
from typing import Any, Dict, Optional, Text
import dataclasses
from orbit import runner
from orbit import utils
import tensorflow as tf


@dataclasses.dataclass(frozen=True)
class TrainerOverrides:
  """Advanced overrides for Orbit trainers.

  Attributes:
    use_tf_while_loop: A boolean indicates whether to wrap the train step with
      a `tf.while_loop`.
    use_tf_function: A boolean indicates whether a `tf.function` will be used.
      If False, training will run on pure eager mode.
    use_tpu_summary_optimization: A boolean indicates whether to enable the
      performance optimization for summaries in TPUs. In TPUs, writing
      summaries with outside compilation inside train step is slow. If True,
      it creates two `tf.function` with two XLA programs: one with summaries
      and one without, and run the program with summaries (slow one) only if
      necessary.
  """
  use_tf_while_loop: bool = True
  use_tf_function: bool = True
  use_tpu_summary_optimization: bool = False


class StandardTrainer(runner.AbstractTrainer, metaclass=abc.ABCMeta):
  """Implements the standard functionality of AbstractTrainer APIs."""

  def __init__(self,
               train_dataset,
               use_tf_while_loop=True,
               use_tf_function=True,
               use_tpu_summary_optimization=False):
    """Construct a `StandardTrainer` object.

    Args:
      train_dataset: A tf.nest-compatible structure of tf.data.Dataset or
        DistributedDataset.
      use_tf_while_loop: A boolean indicates whether to wrap the train step with
        a `tf.while_loop`.
      use_tf_function: A boolean indicates whether a `tf.function` will be used.
        If False, training will run on pure eager mode.
      use_tpu_summary_optimization: A boolean indicates whether to enable the
        performance optimization for summaries in TPUs. In TPUs, writing
        summaries with outside compilation inside train step is slow. If True,
        it creates two `tf.function` with two XLA programs: one with summaries
          and one without, and run the program with summaries (slow one) only if
          necessary.
    """
    if use_tf_while_loop and not use_tf_function:
      raise ValueError("`use_tf_while_loop=True` and `use_tf_function=False` "
                       "is not supported")
    if use_tpu_summary_optimization and not use_tf_while_loop:
      raise ValueError("`use_tpu_summary_optimization=True` and "
                       "`use_tf_while_loop=False` is not supported")
    self._use_tf_while_loop = use_tf_while_loop
    self._use_tf_function = use_tf_function
    self._train_dataset = train_dataset
    self._train_iter = None
    self._train_loop_fn = None
    self._use_tpu_summary_optimization = use_tpu_summary_optimization

  def train(self,
            num_steps: Optional[tf.Tensor]) -> Optional[Dict[Text, tf.Tensor]]:
    """See base class."""
    self.train_loop_begin()

    if self._train_iter is None:
      self._train_iter = tf.nest.map_structure(iter, self.train_dataset)

    if self._train_loop_fn is None:
      train_fn = self.train_step
      if self._use_tf_while_loop:
        self._train_loop_fn = utils.create_tf_while_loop_fn(train_fn)
        if self._use_tpu_summary_optimization:
          self._train_loop_fn = utils.train_function_with_summaries(
              self._train_loop_fn)
        else:
          self._train_loop_fn = tf.function(self._train_loop_fn)
      else:
        if self._use_tf_function:
          train_fn = tf.function(train_fn)
        self._train_loop_fn = utils.create_loop_fn(train_fn)

    self._train_loop_fn(self._train_iter, num_steps)
    return self.train_loop_end()

  def train_loop_begin(self):
    """Called once at the beginning of the training loop.

    This method is called before dataset iterators creation.
    This is a good place to reset metrics that accumulate values over multiple
    steps of training.
    """
    pass

  @abc.abstractmethod
  def train_step(self, iterator):
    """Implements one step of training.

    What a "step" consists of is up to the implementer. If using distribution
    strategies, the call to this method should take place in the "cross-replica
    context" for generality, to allow e.g. multiple iterator dequeues and calls
    to `strategy.run`.

    Note that if `use_tf_function=True`, all the code inside `train_step` should
    be tf.function compatible, as they will be traced with tf.function. This
    means you cannot put arbitrary python code in this function. If users have
    any numpy operations, they should be put in `train_loop_begin` or
    `train_loop_end` functions.

    Args:
      iterator: A tf.nest-compatible structure of tf.data Iterator or
        DistributedIterator.
    """
    pass

  def train_loop_end(self) -> Optional[Dict[Text, tf.Tensor]]:
    """Called at the end of the training loop.

    This is a good place to get metric results. The value returned from this
    function will be returned as-is from the train() method.

    Returns:
      The function may return a dictionary of `Tensors`, which will be
      written to logs and as TensorBoard summaries. It can also be a
      nested dictionary, yielding a hierarchy of summary directories.
    """
    pass

  @property
  def train_dataset(self):
    """Returns the train_dataset instance."""
    return self._train_dataset

  @train_dataset.setter
  def train_dataset(self, train_dataset):
    """Set a new train dataset and replace with the existing one.

    Any unfinished work in the previous dataset will be discarded.

    Args:
      train_dataset: A tf.nest-compatible structure of tf.data.Dataset or
        DistributedDataset.
    """
    self._train_dataset = train_dataset
    self._train_iter = None


@dataclasses.dataclass(frozen=True)
class EvaluatorOverrides:
  """Advanced overrides for Orbit evaluators.

  Attributes:
    use_tf_function: A boolean indicates whether a `tf.function` will be used.
      If False, training will run on pure eager mode.
  """
  use_tf_function: bool = True


class StandardEvaluator(runner.AbstractEvaluator, metaclass=abc.ABCMeta):
  """Implements the standard functionality of AbstractEvaluator APIs."""

  def __init__(self, eval_dataset, use_tf_function=True):
    """Construct a `StandardEvaluator` object.

    Args:
      eval_dataset: A tf.nest-compatible structure of tf.data.Dataset or
        DistributedDataset.
      use_tf_function: A boolean indicates whether a `tf.function` will be used.
        If False, evaluation will run on pure eager mode.
    """
    self._eval_use_tf_function = use_tf_function
    self._eval_dataset = eval_dataset
    self._eval_loop_fn = None

  def evaluate(
      self, num_steps: Optional[tf.Tensor]) -> Optional[Dict[Text, tf.Tensor]]:
    """See base class."""
    outputs = self.eval_begin()  # pylint: disable=assignment-from-no-return

    eval_iter = tf.nest.map_structure(iter, self._eval_dataset)
    if self._eval_loop_fn is None:
      eval_fn = self.eval_step
      if self._eval_use_tf_function:
        eval_fn = tf.function(eval_fn)
      self._eval_loop_fn = utils.create_loop_fn(eval_fn)

    outputs = self._eval_loop_fn(
        eval_iter, num_steps, state=outputs, reduce_fn=self.eval_reduce)
    if outputs is None:
      return self.eval_end()
    else:
      return self.eval_end(outputs)

  def eval_begin(self) -> Any:
    """Called once at the beginning of the evaluation.

    This method is called before dataset iterators creation.
    This is a good place to reset metrics that accumulate values over the entire
    evaluation.

    Returns:
      An output which is passed as `state` argument into `eval_reduce` function.
    """
    pass

  @abc.abstractmethod
  def eval_step(self, iterator) -> Any:
    """Implements one step of evaluation.

    What a "step" consists of is up to the implementer. If using distribution
    strategies, the call to this method should take place in the "cross-replica
    context" for generality, to allow e.g. multiple iterator dequeues and calls
    to `strategy.run`.

    Note that if `use_tf_function=True`, all the code inside `eval_step` should
    be tf.function compatible, as they will be traced with tf.function. This
    means you cannot put arbitrary python code in this function. If users have
    any numpy operations, they should be put in `eval_begin`, `eval_end` or
    `eval_reduce` functions.

    Args:
      iterator: A tf.nest-compatible structure of tf.data Iterator or
        DistributedIterator.

    Returns:
      An output which is passed as `step_outputs` argument into `eval_reduce`
      function.
    """
    pass

  def eval_end(self, *args) -> Optional[Dict[Text, tf.Tensor]]:
    """Called at the end of the evaluation.

    This is a good place to get metric results. The value returned from this
    function will be returned as-is from the evaluate() method.

    Args:
      *args: the outputs from `eval_reduce` for the last eval step.

    Returns:
      The function may return a dictionary of `Tensors`, which will be
      written to logs and as TensorBoard summaries. It can also be a
      nested dictionary, yielding a hierarchy of summary directories.
    """
    pass

  def eval_reduce(self, state=None, step_outputs=None) -> Any:
    """A function to do the reduction on the evaluation outputs per step.

    This is useful for passing states throughout evaluation. E.g. it can be used
    to maintain the output losses from all the evaluation steps, and compute the
    mean loss in `eval_end` function.

    Args:
      state: A maintained state throughout the evaluation.
      step_outputs: Outputs from the current evaluation step.

    Returns:
      An output which is passed as `state` argument into `eval_reduce` function
      for the next step. After evaluation is finished, the output from last step
      will be passed into `eval_end` function.
    """
    pass

  @property
  def eval_dataset(self):
    """Returns the train_datase instance."""
    return self._eval_dataset

  @eval_dataset.setter
  def eval_dataset(self, eval_dataset):
    """Set a new eval dataset and replace with the existing one.

    Args:
      eval_dataset: A tf.nest-compatible structure of tf.data.Dataset or
        DistributedDataset.
    """
    self._eval_dataset = eval_dataset