# coding=utf-8
# Copyright 2021 The Deadunits 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.

# Lint as: python2, python3
"""Has some auxiliary layer definitions needed for pruning and mean replacement.

This file implements two layers:
  - MeanReplacer
  - Masked Layer
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from absl import logging
from six.moves import range

import tensorflow.compat.v2 as tf


class MaskedLayer(tf.keras.layers.Wrapper):
  """This layer wraps keras.layers for applying binary masks to parameters.

  Example Usage:
    `ml = MaskedLayer(tf.keras.layers.Dense(32))`

  To ensure masks are saved properly, initiate tf.train.Checkpoint object after
    masks are generated.
  This layer can be used for two tasks.
  1. Pruning
    One can prune a layer by setting the binary mask.
    `ml.set_mask(binary_weight_mask)`

  2. Dynamic_training: using a part of the layer. controlled dropout
     TODO
     This part is NOT yet needed and therefore not supported.
  """

  def __init__(self, layer, mask_initializer=tf.initializers.ones, **kwargs):
    """Initilizes the MaskedLayer.

    MaskedLayer assumes that the `tf.keras.layers.Layer` that it is wrapping has
    their weights at `layer.weights[0]` and bias at `layer.weights[1]`. This
    is the current API for `tf.keras.layers.Conv2D` or `tf.keras.layers.Dense`.
    Args:
      layer: tf.keras.layers.Layer that is going to be wrapped.
      mask_initializer: func, used to initialize the mask.
      **kwargs: any remaining named arguments passed to the super(MaskedLayer).

    Raises:
      ValueError: if the `layer` passed is not a keras layer.
    """
    if not isinstance(layer, tf.keras.layers.Layer):
      raise ValueError('Please initialize `MaskedLayer` layer with a '
                       '`Layer` instance. You passed: %s' % layer)
    super(MaskedLayer, self).__init__(layer, **kwargs)
    self.enabled = True
    self.mask_initializer = mask_initializer

  def build(self, input_shape):
    logging.debug('MaskedLayer generated with shape: %s', input_shape)
    if not self.layer.built:
      self.layer.build(input_shape)
      self.layer.built = True
    # This is the case for now, we only support layers with weights+bias
    assert len(self.layer.weights) == 2
    self.mask_weight = self.add_variable('mask_weight',
                                         shape=self.layer.weights[0].shape,
                                         initializer=self.mask_initializer,
                                         trainable=False,
                                         dtype=self.layer.weights[0].dtype)
    self.mask_bias = self.add_variable('mask_bias',
                                       shape=self.layer.weights[1].shape,
                                       initializer=self.mask_initializer,
                                       trainable=False,
                                       dtype=self.layer.weights[1].dtype)
    super(MaskedLayer, self).build(input_shape)
    self.built = True

  def call(self, inputs, **kwargs):
    # We need to apply the mask to ensure the masked parameters are zero.
    # This is due to the fact that the gradient update might change the masked
    # parameters.

    # Note: Another way to implement this is to mask gradient. This would remove
    # the need for setting masked parameters zero before each forward pass.
    self.apply_masks()
    return self.layer(inputs, **kwargs)

  def set_mask(self, mask, is_bias=False):
    """Verifies the shape of the mask and sets it.

    Args:
      mask: tf.Tensor, same size as the weights (`self.layer.weights[0]`) or
        bias (`self.layer.weights[1]`) if `is_bias=True`.
      is_bias: sets the mask for the bias (layer.weights[1]) parameters if True.
    """
    # One shouldn't set the mask before building the layer /init.
    assert self.built
    current_mask = self.mask_bias if is_bias else self.mask_weight
    assert isinstance(mask, tf.Tensor)
    assert current_mask.shape == mask.shape
    mask = tf.cast(mask, current_mask.dtype)
    current_mask.assign(mask)

  def apply_masks(self):
    assert self.built
    if self.enabled:
      self.layer.weights[0].assign(self.layer.weights[0] * self.mask_weight)
      self.layer.weights[1].assign(self.layer.weights[1] * self.mask_bias)

  # Needs to be implemented (from tf.keras.layers.Layer)
  def compute_output_shape(self, input_shape):
    return input_shape

  def __repr__(self):