# coding=utf-8
# Copyright 2020 The Mesh TensorFlow 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.

"""Mesh Tensorflow Optimizers."""


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import re
import gin
from mesh_tensorflow import ops_with_redefined_builtins as mtf
import tensorflow.compat.v1 as tf


def make_optimizer(hparams, lr):
  if hparams.optimizer == "SGD":
    return SgdOptimizer(lr)
  elif hparams.optimizer == "Adafactor":
    return adafactor_optimizer_from_hparams(hparams, lr)
  else:
    raise ValueError("Unknown Optimizer")


class Optimizer(object):
  """Base optimizer class.

  Constructor of subclasses must take `learning_rate` as an argument.
  """

  def apply_grads(self, grads, variables):
    """Apply gradients to variables.

    Call this function externally instead of apply_grad().  This causes the
    operations to be combined, which is necessary for stacking variables
    see mtf.rewrite_stack_variables().

    Args:
      grads: a list of Tensor
      variables: a list of Variables
    Returns:
      a list of Operations
    """
    ops = []
    for grad, var in zip(grads, variables):
      ops.extend(self.apply_grad(grad, var))
    if not ops:
      return ops
    return variables[0].graph.combine_assignments(ops)

  def apply_grad(self, grad, var):
    """Update variable and accumulators.

    Args:
      grad: a Tensor
      var: a Variablle
    Returns:
      a list of Operations
    """
    raise ValueError("apply_grad not implemented %s %s" % (grad, var))


@gin.configurable
class SgdOptimizer(Optimizer):
  """Optimizer implementing SGD."""

  def __init__(self, learning_rate):
    self._lr = learning_rate

  @property
  def lr(self):
    return self._lr

  def apply_grad(self, grad, var):
    if grad is None:
      tf.logging.warning("Gradient is None for variable %s" % var.name)
      return []
    # It is critical to use assign_sub instead of mtf.assign(var - ...)
    #  for the case of bfloat16 activations, so as to avoid repeatedly rounding
    #  the slice value, which results in poor quality.
    return [mtf.assign_sub(var, grad * self.lr)]


@gin.configurable
class MomentumOptimizer(Optimizer):
  """SGD with momentum."""

  def __init__(self, learning_rate, momentum):
    self._lr = learning_rate
    self._momentum = momentum

  @property
  def lr(self):
    return self._lr

  @property
  def momentum(self):
    return self._momentum

  def apply_grad(self, grad, var):
    if grad is None:
      tf.logging.warning("Gradient is None for variable %s" % var.name)
      return []

    updates = []
    v = mtf.get_variable(
        var.mesh, var.name + "_momentum_v", var.shape,
        dtype=var.dtype, initializer=tf.zeros_initializer(), trainable=False)

    with tf.variable_scope(var.name + "/sgd_momentum"):
      updates.append(mtf.assign(v, grad * self.lr + v * self.momentum))
      updates.append(mtf.assign_sub(var, v))

    return updates


@gin.configurable
class AdamWeightDecayOptimizer(Optimizer):
  """A basic Adam optimizer that includes "correct" L2 weight decay."""

  def __init__(self,
               learning_rate,
               weight_decay_rate=0.0,
               beta_1=0.9,
               beta_2=0.999,
               epsilon=1e-6,
               exclude_from_weight_decay=None):
    """Constructs a AdamWeightDecayOptimizer."""

    self.learning_rate = learning_rate
    self.weight_decay_rate = weight_decay_rate
    self.beta_1 = beta_1
    self.beta_2 = beta_2
    self.epsilon = epsilon
    self.exclude_from_weight_decay = exclude_from_weight_decay

  def apply_grad(self, grad, var):
    """See base class."""
    if grad is None:
      tf.logging.warning("Gradient is None for variable %s" % var.name)
      return []
    grad = mtf.to_float(grad)

    assignments = []

    m = mtf.get_variable(
        var.mesh, var.name + "/adam_m", var.shape,
        initializer=tf.zeros_initializer(), trainable=False)

    v = mtf.get_variable(
        var.mesh, var.name + "/adam_v", var.shape,
        initializer=tf.zeros_initializer(), trainable=False)

    # Standard Adam update.
    next_m = self.beta_1 * m + (1.0 - self.beta_1) * grad
    next_v = self.beta_2 * v + (1.0 - self.beta_2) * mtf.square(grad)

    update = next_m / (mtf.sqrt(next_v) + self.epsilon)

    # Just adding the square of the weights to the loss function is *not*
    # the correct way of using L2 regularization/weight decay with Adam,
    # since that will interact with the m and v parameters in strange ways.
    #
    # Instead we want ot decay the weights in a manner that doesn't interact
    # with the m/v parameters. This is equivalent to adding the square
    # of the weights to the loss with plain (non-momentum) SGD.
    if self._do_use_weight_decay(var.name):
      update += self.weight_decay_rate * var.value

    update_with_lr = self.learning_rate * update

    var_update = mtf.assign_sub(var, update_with_lr)

    assignments.extend(
        [var_update,
         mtf.assign(m, next_m),
         mtf.assign(v, next_v)])
    return assignments

  def _do_use_weight_decay(self, param_name):
    """Whether to use L2 weight decay for `param_name`."""
    if not self.weight_decay_rate:
      return False
    if self.exclude_from_weight_decay:
      for r in self.exclude_from_weight_decay:
        if re.search(r, param_name) is not None:
          return False
    return True


@gin.configurable
class AdafactorOptimizer(Optimizer):
  """Adafactor."""

  def __init__(self,
               multiply_by_parameter_scale=True,
               learning_rate=None,
               decay_rate=None,
               beta1=0.0,
               clipping_threshold=1.0,
               factored=True,
               epsilon1=1e-30,
               epsilon2=1e-3,
               min_dim_size_to_factor=128):
    """Construct a new Adafactor optimizer.

    See class comment.

    Args:
      multiply_by_parameter_scale: a boolean
      learning_rate: an optional Scalar.
      decay_rate: an optional Scalar.
      beta1: a float value between 0 and 1
      clipping_threshold: an optional float >= 1
      factored: a boolean - whether to use factored second-moment estimator
        for 2d variables
      epsilon1: Regularization constant for squared gradient.
      epsilon2: Regularization constant for parameter scale.
      min_dim_size_to_factor: only factor accumulator if two tensor dimensions
        are at least this size.

    Raises:
      ValueError: if absolute_update_scale and relative_update_scale_fn are both
        present or both absent.
    """
    self._multiply_by_parameter_scale = multiply_by_parameter_scale
    if learning_rate is None:
      learning_rate = self._learning_rate_default(multiply_by_parameter_scale)
    self._learning_rate = learning_rate
    if decay_rate is None:
      decay_rate = self._decay_rate_default()
    self._decay_rate = decay_rate
    self._beta1 = beta1
    self._clipping_threshold = clipping_threshold
    self._factored = factored
    self._epsilon1 = epsilon1
    self._epsilon2 = epsilon2
    self._min_dim_size_to_factor = min_dim_size_to_factor

  def _factored_dims(self, shape):
    """Should we use a factored second moment estimator.

    Based on the shape of the variable.
    If we factor the accumulator, then this function returns a list of two
    mtf.Dimensions to reduce over.  We always pick the two largest dimensions.
    If there are not two dimensions of size >= min_dim_size_to_factor, then we
    do not factor.

    Args:
      shape: a Shape
    Returns:
      either a list of 2 Dimensions or None
    """
    if not self._factored or shape.ndims < 2:
      return None
    sorted_dims = sorted(shape.dims, key=lambda d: -d.size)
    if sorted_dims[1].size < self._min_dim_size_to_factor:
      return None
    return sorted_dims[:2]

  def _parameter_scale(self, var):
    """Estimate the scale of the parameters from the current values.

    We include a minimum value of 0.001 to give it a chance to escape 0
    if it was zero-initialized.

    Instead of using the value, we could impute the scale from the shape,
    as initializers do.

    Args:
      var: a variable or Tensor.
    Returns:
      a Scalar
    """
    return mtf.maximum(reduce_rms(var), self._epsilon2)

  def apply_grad(self, grad, var):
    if grad is None:
      tf.logging.warning("Gradient is None for variable %s" % var.name)
      return []
    # create slots
    grad = mtf.to_float(grad)
    factored_dims = self._factored_dims(var.shape)
    if factored_dims:
      d0, d1 = factored_dims
      vr_shape = var.shape - d0
      vc_shape = var.shape - d1
      vr = mtf.get_variable(
          var.mesh, var.name + "_slot_vr", vr_shape,
          initializer=tf.zeros_initializer(), trainable=False)
      vc = mtf.get_variable(
          var.mesh, var.name + "_slot_vc", vc_shape,
          initializer=tf.zeros_initializer(), trainable=False)
    else:
      v = mtf.get_variable(
          var.mesh, var.name + "_slot_v", var.shape,
          initializer=tf.zeros_initializer(), trainable=False)
    if self._beta1:
      m = mtf.get_variable(
          var.mesh, var.name + "_slot_m", var.shape,
          initializer=tf.zeros_initializer(), trainable=False)

    with tf.variable_scope(var.name + "/adafactor"):
      grad_squared = mtf.square(grad) + self._epsilon1
      decay_rate = self._decay_rate
      old_val = mtf.to_float(var.value)
      if self._multiply_by_parameter_scale:
        update_scale = self._parameter_scale(old_val) * self._learning_rate
      else:
        update_scale = self._learning_rate
      mixing_rate = 1.0 - decay_rate
      updates = []
      if factored_dims:
        grad_squared_row_mean = mtf.reduce_mean(
            grad_squared, output_shape=vr_shape)
        grad_squared_col_mean = mtf.reduce_mean(
            grad_squared, output_shape=vc_shape)
        new_vr = vr * decay_rate + grad_squared_row_mean * mixing_rate
        new_vc = vc * decay_rate + grad_squared_col_mean * mixing_rate
        vr_update = mtf.assign(vr, new_vr)
        vc_update = mtf.assign(vc, new_vc)
        updates.extend([vr_update, vc_update])
        long_term_mean = mtf.reduce_mean(new_vr, reduced_dim=d1)
        r_factor = mtf.rsqrt(new_vr / long_term_mean)
        c_factor = mtf.rsqrt(new_vc)
        x = grad * r_factor * c_factor
      else:
        new_v = v * decay_rate + grad_squared * mixing_rate
        v_update = mtf.assign(v, new_v)
        updates.append(v_update)
        x = grad * mtf.rsqrt(new_v)
      if self._clipping_threshold is not None:
        clipping_denom = mtf.maximum(
            1.0, reduce_rms(x) / self._clipping_threshold)
        x /= clipping_denom
      subtrahend = x * update_scale
      if self._beta1:
        new_m = (m * tf.constant(self._beta1)
                 + subtrahend * tf.constant(1.0 - self._beta1))
        subtrahend = new_m
        updates.append(mtf.assign(m, new_m))
      # It is critical to use assign_sub instead of mtf.assign(var - subtrahend)
      #  for the case of bfloat16 activations, so as to avoid repeatedly
      #  rounding the slice value, which results in poor quality.
      var_update = mtf.assign_sub(var, subtrahend)
      updates.append(var_update)
      return updates

  def _decay_rate_default(self):
    return adafactor_decay_rate_pow(0.8)

  def _learning_rate_default(self, multiply_by_parameter_scale):
    learning_rate = tf.minimum(tf.math.rsqrt(step_num() + 1.0), 0.01)
    if not multiply_by_parameter_scale:
      learning_rate *= 0.05
    return learning_rate


def adafactor_decay_rate_adam(beta2):
  """Second-moment decay rate like Adam, subsuming the correction factor.

  Args:
    beta2: a float between 0 and 1
  Returns:
    a scalar
  """
  t = tf.cast(tf.train.get_or_create_global_step(), tf.float32) + 1.0
  decay = beta2 * (1.0 - tf.pow(beta2, t - 1.0)) / (1.0 - tf.pow(beta2, t))
  return decay


def adafactor_decay_rate_pow(exponent):
  """Second moment decay rate where memory-length grows as step_num^exponent.

  Args:
    exponent: a float between 0 and 1
  Returns:
    a scalar
  """
  return 1.0 - tf.pow((step_num() + 1.0), -exponent)


def step_num():
  return tf.cast(tf.train.get_or_create_global_step(), tf.float32)


def adafactor_optimizer_from_hparams(hparams, lr):
  """Create an Adafactor optimizer based on model hparams.

  Args:
    hparams: model hyperparameters
    lr: learning rate scalar.
  Returns:
    an AdafactorOptimizer
  Raises:
    ValueError: on illegal values
  """
  if hparams.optimizer_adafactor_decay_type == "Adam":
    decay_rate = adafactor_decay_rate_adam(
        hparams.optimizer_adafactor_beta2)
  elif hparams.optimizer_adafactor_decay_type == "pow":
    decay_rate = adafactor_decay_rate_pow(
        hparams.optimizer_adafactor_memory_exponent)
  else:
    raise ValueError("unknown optimizer_adafactor_decay_type")
  return AdafactorOptimizer(
      multiply_by_parameter_scale=(
          hparams.optimizer_adafactor_multiply_by_parameter_scale),
      learning_rate=lr,
      decay_rate=decay_rate,
      beta1=hparams.optimizer_adafactor_beta1,
      clipping_threshold=hparams.optimizer_adafactor_clipping_threshold,
      factored=hparams.optimizer_adafactor_factored)


def reduce_rms(x):
  return mtf.sqrt(mtf.reduce_mean(mtf.square(x)))
