"""Optimization utilities."""

import inspect

import tensorflow as tf
import tensorflow_addons as tfa

from tensorflow_addons.optimizers.weight_decay_optimizers import DecoupledWeightDecayExtension

from opennmt.utils import misc


_OPTIMIZERS_REGISTRY = misc.ClassRegistry(base_class=tf.keras.optimizers.Optimizer)

register_optimizer = _OPTIMIZERS_REGISTRY.register  # pylint: disable=invalid-name

def get_optimizer_class(name):
  """Returns the optimizer class.

  Args:
    name: The optimizer name.

  Returns:
    A class extending ``tf.keras.optimizers.Optimizer``.

  Raises:
    ValueError: if :obj:`name` can not be resolved to an optimizer class.
  """
  optimizer_class = None
  if optimizer_class is None:
    optimizer_class = getattr(tf.keras.optimizers, name, None)
  if optimizer_class is None:
    optimizer_class = getattr(tfa.optimizers, name, None)
  if optimizer_class is None:
    optimizer_class = _OPTIMIZERS_REGISTRY.get(name)
  if optimizer_class is None:
    raise ValueError("Unknown optimizer class: %s" % name)
  return optimizer_class

def make_optimizer(name, learning_rate, **kwargs):
  """Creates the optimizer.

  Args:
    name: The name of the optimizer class in ``tf.keras.optimizers`` or
      ``tfa.optimizers`` as a string.
    learning_rate: The learning rate or learning rate schedule to use.
    **kwargs: Additional optimizer arguments. If ``weight_decay`` is set, the
      optimizer will be extended with decoupled weight decay.

  Returns:
    A ``tf.keras.optimizers.Optimizer`` instance.

  Raises:
    ValueError: if :obj:`name` can not be resolved to an optimizer class.
  """
  optimizer_class = get_optimizer_class(name)
  if "weight_decay" in kwargs:
    if DecoupledWeightDecayExtension not in inspect.getmro(optimizer_class):
      optimizer_class = tfa.optimizers.extend_with_decoupled_weight_decay(optimizer_class)
  optimizer = optimizer_class(learning_rate=learning_rate, **kwargs)
  return optimizer


class GradientAccumulator(object):
  """Gradient accumulation utility.

  When used with a distribution strategy, the accumulator should be called in a
  replica context. Gradients will be accumulated locally on each replica and
  without synchronization. Users should then call ``.gradients``, scale the
  gradients if required, and pass the result to ``apply_gradients``.
  """

  # We use the ON_READ synchronization policy so that no synchronization is
  # performed on assignment. To get the value, we call .value() which returns the
  # value on the current replica without synchronization.

  def __init__(self):
    """Initializes the accumulator."""
    self._gradients = []
    self._accum_steps = None

  @property
  def step(self):
    """Number of accumulated steps."""
    if self._accum_steps is None:
      self._accum_steps = tf.Variable(
          tf.constant(0, dtype=tf.int64),
          trainable=False,
          synchronization=tf.VariableSynchronization.ON_READ,
          aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
    return self._accum_steps.value()

  @property
  def gradients(self):
    """The accumulated gradients on the current replica."""
    if not self._gradients:
      raise ValueError("The accumulator should be called first to initialize the gradients")
    return list(gradient.value() for gradient in self._gradients)

  def __call__(self, gradients):
    """Accumulates :obj:`gradients` on the current replica."""
    if not self._gradients:
      _ = self.step  # Create the step variable.
      self._gradients.extend([
          tf.Variable(
              tf.zeros_like(gradient),
              trainable=False,
              synchronization=tf.VariableSynchronization.ON_READ)
          for gradient in gradients])
    if len(gradients) != len(self._gradients):
      raise ValueError("Expected %s gradients, but got %d" % (
          len(self._gradients), len(gradients)))

    for accum_gradient, gradient in zip(self._gradients, gradients):
      accum_gradient.assign_add(gradient, read_value=False)
    self._accum_steps.assign_add(1)

  def reset(self):
    """Resets the accumulated gradients on the current replica."""
    if not self._gradients:
      return
    self._accum_steps.assign(0)
    for gradient in self._gradients:
      gradient.assign(tf.zeros(gradient.shape, dtype=gradient.dtype), read_value=False)
