# Copyright 2019 The AdaNet 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

#     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.
"""Adanet implementation for weighted Subnetwork and Ensemblers."""

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

import collections

from absl import logging
from adanet import tf_compat
from adanet.ensemble.ensembler import Ensemble
from adanet.ensemble.ensembler import Ensembler
import tensorflow.compat.v2 as tf


def _stringify(key):
  """Flattens tuple and list keys into strings."""

  if isinstance(key, (tuple, list)):
    return "_".join([str(el) for el in key])
  return key


def _lookup_if_dict(target, key):
  if isinstance(target, dict):
    return target[key]
  return target


class WeightedSubnetwork(
    collections.namedtuple(
        "WeightedSubnetwork",
        ["name", "iteration_number", "weight", "logits", "subnetwork"])):
  # pyformat: disable
  """An AdaNet weighted subnetwork.

  A weighted subnetwork is a weight applied to a subnetwork's last layer
  or logits (depending on the mixture weights type).

  Args:
    name: String name of :code:`subnetwork` as defined by its
      :class:`adanet.subnetwork.Builder`.
    iteration_number: Integer iteration when the subnetwork was created.
    weight: The weight :class:`tf.Tensor` or dict of string to weight
      :class:`tf.Tensor` (for multi-head) to apply to this subnetwork. The
      AdaNet paper refers to this weight as :math:`w` in Equations (4), (5),
      and (6).
    logits: The output :class:`tf.Tensor` or dict of string to weight
      :class:`tf.Tensor` (for multi-head) after the matrix multiplication of
      :code:`weight` and the subnetwork's :code:`last_layer`. The output's shape
      is [batch_size, logits_dimension]. It is equivalent to a linear logits
      layer in a neural network.
    subnetwork: The :class:`adanet.subnetwork.Subnetwork` to weight.

  Returns:
    An :class:`adanet.ensemble.WeightedSubnetwork` object.
  """
  # pyformat: enable

  def __new__(cls,
              name="",
              iteration_number=0,
              weight=None,
              logits=None,
              subnetwork=None):
    return super(WeightedSubnetwork, cls).__new__(
        cls,
        name=name,
        iteration_number=iteration_number,
        weight=weight,
        logits=logits,
        subnetwork=subnetwork)


class ComplexityRegularized(
    collections.namedtuple("ComplexityRegularized", [
        "weighted_subnetworks", "bias", "logits", "subnetworks",
        "complexity_regularization"
    ]), Ensemble):
  r"""An AdaNet ensemble where subnetworks are regularized by model complexity.

  Hence an ensemble is a collection of subnetworks which forms a neural network
  through the weighted sum of their outputs:

  .. math::

      F(x) = \sum_{i=1}^{N}w_ih_i(x) + b

  Args:
    weighted_subnetworks: List of :class:`adanet.ensemble.WeightedSubnetwork`
      instances that form this ensemble. Ordered from first to most recent.
    bias: Bias term :class:`tf.Tensor` or dict of string to bias term
      :class:`tf.Tensor` (for multi-head) for the ensemble's logits.
    logits: Logits :class:`tf.Tensor` or dict of string to logits
      :class:`tf.Tensor` (for multi-head). The result of the function *f* as
      defined in Section 5.1 which is the sum of the logits of all
      :class:`adanet.WeightedSubnetwork` instances in ensemble.
    subnetworks: List of :class:`adanet.subnetwork.Subnetwork` instances that
      form this ensemble. This is kept together with weighted_subnetworks for
      legacy reasons.
    complexity_regularization: Regularization to be added in the Adanet loss.

  Returns:
    An :class:`adanet.ensemble.Weighted` instance.
  """

  def __new__(cls,
              weighted_subnetworks,
              bias,
              logits,
              subnetworks=None,
              complexity_regularization=None):
    return super(ComplexityRegularized, cls).__new__(
        cls,
        weighted_subnetworks=list(weighted_subnetworks),
        bias=bias,
        logits=logits,
        subnetworks=list(subnetworks or []),
        complexity_regularization=complexity_regularization)


class MixtureWeightType(object):
  """Mixture weight types available for learning subnetwork contributions.

  The following mixture weight types are defined:

  * `SCALAR`: Produces a rank 0 `Tensor` mixture weight.
  * `VECTOR`: Produces a rank 1 `Tensor` mixture weight.
  * `MATRIX`: Produces a rank 2 `Tensor` mixture weight.
  """

  SCALAR = "scalar"
  VECTOR = "vector"
  MATRIX = "matrix"


class ComplexityRegularizedEnsembler(Ensembler):
  # pyformat: disable
  r"""The AdaNet algorithm implemented as an :class:`adanet.ensemble.Ensembler`.

  The AdaNet algorithm was introduced in the [Cortes et al. ICML 2017] paper:
  https://arxiv.org/abs/1607.01097.

  The AdaNet algorithm uses a weak learning algorithm to iteratively generate a
  set of candidate subnetworks that attempt to minimize the loss function
  defined in Equation (4) as part of an ensemble. At the end of each iteration,
  the best candidate is chosen based on its ensemble's complexity-regularized
  train loss. New subnetworks are allowed to use any subnetwork weights within
  the previous iteration's ensemble in order to improve upon them. If the
  complexity-regularized loss of the new ensemble, as defined in Equation (4),
  is less than that of the previous iteration's ensemble, the AdaNet algorithm
  continues onto the next iteration.

  AdaNet attempts to minimize the following loss function to learn the mixture
  weights :math:`w` of each subnetwork :math:`h` in the ensemble with
  differentiable convex non-increasing surrogate loss function :math:`\Phi`:

  Equation (4):

  .. math::

      F(w) = \frac{1}{m} \sum_{i=1}^{m} \Phi \left(\sum_{j=1}^{N}w_jh_j(x_i),
      y_i \right) + \sum_{j=1}^{N} \left(\lambda r(h_j) + \beta \right) |w_j|

  with :math:`\lambda >= 0` and :math:`\beta >= 0`.

  Args:
    optimizer: String, :class:`tf.train.Optimizer` object, or callable that
      creates the optimizer to use for training the ensemble weights. If left
      as :code:`None`, :meth:`tf.no_op()` is used instead.
    mixture_weight_type: The :class:`adanet.ensemble.MixtureWeightType` defining
      which mixture weight type to learn on top of the subnetworks' logits.
    mixture_weight_initializer: The initializer for mixture_weights. When
      :code:`None`, the default is different according to
      :code:`mixture_weight_type`:

        - :code:`SCALAR` initializes to :math:`1/N` where :math:`N` is the
          number of subnetworks in the ensemble giving a uniform average.
        - :code:`VECTOR` initializes each entry to :math:`1/N` where :math:`N`
          is the number of subnetworks in the ensemble giving a uniform average.
        - :code:`MATRIX` uses :meth:`tf.zeros_initializer`.
    warm_start_mixture_weights: Whether, at the beginning of an iteration, to
      initialize the mixture weights of the subnetworks from the previous
      ensemble to their learned value at the previous iteration, as opposed to
      retraining them from scratch. Takes precedence over the value for
      :code:`mixture_weight_initializer` for subnetworks from previous
      iterations.
    model_dir: The model dir to use for warm-starting mixture weights and bias
      at the logit layer. Ignored if :code:`warm_start_mixture_weights` is
      :code:`False`.
    adanet_lambda: Float multiplier :math:`\lambda` for applying :math:`L1`
      regularization to subnetworks' mixture weights :math:`w` in the ensemble
      proportional to their complexity. See Equation (4) in the AdaNet paper.
    adanet_beta: Float :math:`L1` regularization multiplier :math:`\beta` to apply
      equally to all subnetworks' weights :math:`w` in the ensemble regardless of
      their complexity. See Equation (4) in the AdaNet paper.
    use_bias: Whether to add a bias term to the ensemble's logits.
    name: Optional name for the ensembler. Defaults to 'complexity_regularized'.

  Returns:
    An `adanet.ensemble.ComplexityRegularizedEnsembler` instance.

  Raises:
    ValueError: if :code:`warm_start_mixture_weights` is :code:`True` but
    :code:`model_dir` is :code:`None`.
  """
  # pyformat: enable

  def __init__(self,
               optimizer=None,
               mixture_weight_type=MixtureWeightType.SCALAR,
               mixture_weight_initializer=None,
               warm_start_mixture_weights=False,
               model_dir=None,
               adanet_lambda=0.,
               adanet_beta=0.,
               use_bias=False,
               name=None):
    if warm_start_mixture_weights:
      if model_dir is None:
        raise ValueError("model_dir cannot be None when "
                         "warm_start_mixture_weights is True.")

    self._optimizer = optimizer
    self._mixture_weight_type = mixture_weight_type
    self._mixture_weight_initializer = mixture_weight_initializer
    self._warm_start_mixture_weights = warm_start_mixture_weights
    self._model_dir = model_dir
    self._adanet_lambda = adanet_lambda
    self._adanet_beta = adanet_beta
    self._use_bias = use_bias
    self._name = name

  @property
  def name(self):
    if self._name:
      return self._name
    return "complexity_regularized"

  def build_ensemble(self,
                     subnetworks,
                     previous_ensemble_subnetworks,
                     features,
                     labels,
                     logits_dimension,
                     training,
                     iteration_step,
                     summary,
                     previous_ensemble,
                     previous_iteration_checkpoint=None):
    del features, labels, logits_dimension, training, iteration_step  # unused
    weighted_subnetworks = []
    subnetwork_index = 0
    num_subnetworks = len(subnetworks)

    if previous_ensemble_subnetworks and previous_ensemble:
      num_subnetworks += len(previous_ensemble_subnetworks)
      for weighted_subnetwork in previous_ensemble.weighted_subnetworks:
        if weighted_subnetwork.subnetwork not in previous_ensemble_subnetworks:
          # Pruned.
          continue
        weight_initializer = None
        if self._warm_start_mixture_weights:
          if isinstance(weighted_subnetwork.subnetwork.last_layer, dict):
            weight_initializer = {
                key: self._load_variable(weighted_subnetwork.weight[key],
                                         previous_iteration_checkpoint)
                for key in sorted(weighted_subnetwork.subnetwork.last_layer)
            }
          else:
            weight_initializer = self._load_variable(
                weighted_subnetwork.weight, previous_iteration_checkpoint)
        with tf_compat.v1.variable_scope(
            "weighted_subnetwork_{}".format(subnetwork_index)):
          weighted_subnetworks.append(
              self._build_weighted_subnetwork(
                  weighted_subnetwork.subnetwork,
                  num_subnetworks,
                  weight_initializer=weight_initializer))
        subnetwork_index += 1

    for subnetwork in subnetworks:
      with tf_compat.v1.variable_scope(
          "weighted_subnetwork_{}".format(subnetwork_index)):
        weighted_subnetworks.append(
            self._build_weighted_subnetwork(subnetwork, num_subnetworks))
      subnetwork_index += 1

    if previous_ensemble:
      if len(
          previous_ensemble.subnetworks) == len(previous_ensemble_subnetworks):
        bias = self._create_bias_term(
            weighted_subnetworks,
            prior=previous_ensemble.bias,
            previous_iteration_checkpoint=previous_iteration_checkpoint)
      else:
        bias = self._create_bias_term(
            weighted_subnetworks,
            prior=None,
            previous_iteration_checkpoint=previous_iteration_checkpoint)
        logging.info("Builders using a pruned set of the subnetworks "
                     "from the previous ensemble, so its ensemble's bias "
                     "term will not be warm started with the previous "
                     "ensemble's bias.")
    else:
      bias = self._create_bias_term(weighted_subnetworks)

    logits = self._create_ensemble_logits(weighted_subnetworks, bias, summary)
    complexity_regularization = 0
    if isinstance(logits, dict):
      for key in sorted(logits):
        complexity_regularization += self._compute_complexity_regularization(
            weighted_subnetworks, summary, key)
    else:
      complexity_regularization = self._compute_complexity_regularization(
          weighted_subnetworks, summary)

    return ComplexityRegularized(
        weighted_subnetworks=weighted_subnetworks,
        bias=bias,
        subnetworks=[ws.subnetwork for ws in weighted_subnetworks],
        logits=logits,
        complexity_regularization=complexity_regularization)

  def _load_variable(self, var, previous_iteration_checkpoint):
    latest_checkpoint = tf.train.latest_checkpoint(self._model_dir)
    status = previous_iteration_checkpoint.restore(latest_checkpoint)
    try:
      status.expect_partial().assert_nontrivial_match()
    except AssertionError:
      # Fall back to v1 checkpoint when not using v2 checkpoint.
      return tf.train.load_variable(self._model_dir, tf_compat.tensor_name(var))
    else:
      with tf_compat.v1.Session() as sess:
        status.initialize_or_restore(sess)
        return sess.run(var)

  def _compute_adanet_gamma(self, complexity):
    """For a subnetwork, computes: lambda * r(h) + beta."""

    if self._adanet_lambda == 0.:
      return self._adanet_beta
    return tf.scalar_mul(self._adanet_lambda,
                         tf.cast(complexity,
                                 dtype=tf.float32)) + self._adanet_beta

  def _select_mixture_weight_initializer(self, num_subnetworks):
    if self._mixture_weight_initializer:
      return self._mixture_weight_initializer
    if (self._mixture_weight_type == MixtureWeightType.SCALAR or
        self._mixture_weight_type == MixtureWeightType.VECTOR):
      return tf_compat.v1.constant_initializer(1. / num_subnetworks)
    return tf_compat.v1.zeros_initializer()

  def _build_weighted_subnetwork(self,
                                 subnetwork,
                                 num_subnetworks,
                                 weight_initializer=None):
    """Builds an `adanet.ensemble.WeightedSubnetwork`.

    Args:
      subnetwork: The `Subnetwork` to weight.
      num_subnetworks: The number of subnetworks in the ensemble.
      weight_initializer: Initializer for the weight variable.

    Returns:
      A `WeightedSubnetwork` instance.

    Raises:
      ValueError: When the subnetwork's last layer and logits dimension do
        not match and requiring a SCALAR or VECTOR mixture weight.
    """

    if isinstance(subnetwork.last_layer, dict):
      logits, weight = {}, {}
      for i, key in enumerate(sorted(subnetwork.last_layer)):
        logits[key], weight[key] = self._build_weighted_subnetwork_helper(
            subnetwork, num_subnetworks,
            _lookup_if_dict(weight_initializer, key), key, i)
    else:
      logits, weight = self._build_weighted_subnetwork_helper(
          subnetwork, num_subnetworks, weight_initializer)

    return WeightedSubnetwork(
        subnetwork=subnetwork, logits=logits, weight=weight)

  def _build_weighted_subnetwork_helper(self,
                                        subnetwork,
                                        num_subnetworks,
                                        weight_initializer=None,
                                        key=None,
                                        index=None):
    """Returns the logits and weight of the `WeightedSubnetwork` for key."""

    # Treat subnetworks as if their weights are frozen, and ensure that
    # mixture weight gradients do not propagate through.
    last_layer = _lookup_if_dict(subnetwork.last_layer, key)
    logits = _lookup_if_dict(subnetwork.logits, key)
    weight_shape = None
    last_layer_size = last_layer.get_shape().as_list()[-1]
    logits_size = logits.get_shape().as_list()[-1]
    batch_size = tf.shape(input=last_layer)[0]

    if weight_initializer is None:
      weight_initializer = self._select_mixture_weight_initializer(
          num_subnetworks)
      if self._mixture_weight_type == MixtureWeightType.SCALAR:
        weight_shape = []
      if self._mixture_weight_type == MixtureWeightType.VECTOR:
        weight_shape = [logits_size]
      if self._mixture_weight_type == MixtureWeightType.MATRIX:
        weight_shape = [last_layer_size, logits_size]

    with tf_compat.v1.variable_scope(
        "logits_{}".format(index) if index else "logits"):
      weight = tf_compat.v1.get_variable(
          name="mixture_weight",
          shape=weight_shape,
          initializer=weight_initializer)
      if self._mixture_weight_type == MixtureWeightType.MATRIX:
        # TODO: Add Unit tests for the ndims == 3 path.
        ndims = len(last_layer.get_shape().as_list())
        if ndims > 3:
          raise NotImplementedError(
              "Last Layer with more than 3 dimensions are not supported with "
              "matrix mixture weights.")
        # This is reshaping [batch_size, timesteps, emb_dim ] to
        # [batch_size x timesteps, emb_dim] for matrix multiplication
        # and reshaping back.
        if ndims == 3:
          logging.info("Rank 3 tensors like [batch_size, timesteps, d]  are "
                       "reshaped to rank 2 [ batch_size x timesteps, d] for "
                       "the weight matrix multiplication, and are reshaped "
                       "to their original shape afterwards.")
          last_layer = tf.reshape(last_layer, [-1, last_layer_size])
        logits = tf.matmul(last_layer, weight)
        if ndims == 3:
          logits = tf.reshape(logits, [batch_size, -1, logits_size])
      else:
        logits = tf.multiply(logits, weight)
    return logits, weight

  def _create_bias_term(self,
                        weighted_subnetworks,
                        prior=None,
                        previous_iteration_checkpoint=None):
    """Returns a bias term vector.

    If `use_bias` is set, then it returns a trainable bias variable initialized
    to zero, or warm-started with the given prior. Otherwise it returns
    a non-trainable zero variable.

    Args:
      weighted_subnetworks: List of `WeightedSubnetwork` instances that form
        this ensemble. Ordered from first to most recent.
      prior: Prior bias term `Tensor` of dict of string to `Tensor` (for multi-
        head) for warm-starting the bias term variable.
      previous_iteration_checkpoint: `tf.train.Checkpoint` for iteration t-1.

    Returns:
      A bias term `Tensor` or dict of string to bias term `Tensor` (for multi-
        head).
    """

    if not isinstance(weighted_subnetworks[0].subnetwork.logits, dict):
      return self._create_bias_term_helper(weighted_subnetworks, prior,
                                           previous_iteration_checkpoint)
    bias_terms = {}
    for i, key in enumerate(sorted(weighted_subnetworks[0].subnetwork.logits)):
      bias_terms[key] = self._create_bias_term_helper(
          weighted_subnetworks, prior, previous_iteration_checkpoint, key, i)
    return bias_terms

  def _create_bias_term_helper(self,
                               weighted_subnetworks,
                               prior,
                               previous_iteration_checkpoint,
                               key=None,
                               index=None):
    """Returns a bias term for weights with the given key."""

    shape = None
    if prior is None or not self._warm_start_mixture_weights:
      prior = tf_compat.v1.zeros_initializer()
      logits = _lookup_if_dict(weighted_subnetworks[0].subnetwork.logits, key)
      dims = logits.shape.as_list()

      if len(dims) == 1:
        num_dims = 1
      else:
        assert len(dims) == 2
        num_dims = dims[-1]
        assert num_dims is not None
      shape = num_dims

    else:
      prior = self._load_variable(
          _lookup_if_dict(prior, key), previous_iteration_checkpoint)
    return tf_compat.v1.get_variable(
        name="bias_{}".format(index) if index else "bias",
        shape=shape,
        initializer=prior,
        trainable=self._use_bias)

  def _create_ensemble_logits(self, weighted_subnetworks, bias, summary):
    """Computes the AdaNet weighted ensemble logits.

    Args:
      weighted_subnetworks: List of `WeightedSubnetwork` instances that form
        this ensemble. Ordered from first to most recent.
      bias: Bias term `Tensor` or dict of string to `Tensor` (for multi-head)
        for the AdaNet-weighted ensemble logits.
      summary: A `_ScopedSummary` instance for recording ensemble summaries.

    Returns:
      A two-tuple of:
       1. Ensemble logits `Tensor` or dict of string to logits `Tensor` (for
         multi-head).
       2. Ensemble complexity regularization
    """

    if not isinstance(weighted_subnetworks[0].subnetwork.logits, dict):
      return self._create_ensemble_logits_helper(weighted_subnetworks, bias,
                                                 summary)
    logits_dict = weighted_subnetworks[0].subnetwork.logits
    return {
        key: self._create_ensemble_logits_helper(
            weighted_subnetworks, bias, summary, key=key, index=i)
        for i, key in enumerate(sorted(logits_dict))
    }

  def _create_ensemble_logits_helper(self,
                                     weighted_subnetworks,
                                     bias,
                                     summary,
                                     key=None,
                                     index=None):
    """Returns the AdaNet ensemble logits and regularization term for key."""

    subnetwork_logits = []
    for weighted_subnetwork in weighted_subnetworks:
      subnetwork_logits.append(_lookup_if_dict(weighted_subnetwork.logits, key))
    with tf_compat.v1.variable_scope(
        "logits_{}".format(index) if index else "logits"):
      ensemble_logits = _lookup_if_dict(bias, key)
      for logits in subnetwork_logits:
        ensemble_logits = tf.add(ensemble_logits, logits)
    return ensemble_logits

  def _compute_complexity_regularization(self,
                                         weighted_subnetworks,
                                         summary,
                                         key=None):
    """Returns the AdaNet regularization term contribution for a key."""

    ensemble_complexity_regularization = 0
    total_weight_l1_norms = 0
    weights = []
    for weighted_subnetwork in weighted_subnetworks:
      weight_l1_norm = tf.norm(
          tensor=_lookup_if_dict(weighted_subnetwork.weight, key), ord=1)
      total_weight_l1_norms += weight_l1_norm
      ensemble_complexity_regularization += (
          self._compute_complexity_regularization_helper(
              weight_l1_norm, weighted_subnetwork.subnetwork.complexity))
      weights.append(weight_l1_norm)

    with summary.current_scope():
      # Append a suffix for multi head summaries.
      suffix = "_{}".format(_stringify(key)) if key else ""
      summary.scalar(
          "complexity_regularization/adanet/adanet_weighted_ensemble" + suffix,
          ensemble_complexity_regularization)
      summary.histogram(
          "mixture_weights/adanet/adanet_weighted_ensemble" + suffix, weights)
      for iteration, weight in enumerate(weights):
        scope = "adanet/adanet_weighted_ensemble/subnetwork{}_{}".format(
            suffix, iteration)
        summary.scalar("mixture_weight_norms/{}".format(scope), weight)
        fraction = weight / total_weight_l1_norms
        summary.scalar("mixture_weight_fractions/{}".format(scope), fraction)
    return ensemble_complexity_regularization

  def _compute_complexity_regularization_helper(self, weight_l1_norm,
                                                complexity):
    """For a subnetwork, computes: (lambda * r(h) + beta) * |w|."""

    # Note: Unsafe comparison against float zero.
    if self._adanet_lambda == 0. and self._adanet_beta == 0.:
      return tf.constant(0., name="zero")
    return tf.scalar_mul(self._compute_adanet_gamma(complexity), weight_l1_norm)

  def build_train_op(self, ensemble, loss, var_list, labels, iteration_step,
                     summary, previous_ensemble):
    del labels, iteration_step, summary, previous_ensemble  # unused
    optimizer = self._optimizer
    if callable(optimizer):
      optimizer = optimizer()
    if optimizer is None:
      return tf.no_op()

    # The AdaNet Estimator is responsible for incrementing the global step.
    return optimizer.minimize(
        loss=loss + ensemble.complexity_regularization, var_list=var_list)
