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

"""Implements Axial-Blocks proposed in Axial-DeepLab [1].

Axial-Blocks are based on residual bottleneck blocks, but with the 3x3
convolution replaced with two axial-attention layers, one on the height-axis,
followed by the other on the width-axis.

[1] Axial-Deeplab: Stand-Alone Axial-Attention for Panoptic Segmentation,
    ECCV 2020 Spotlight.
      Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille,
      Liang-Chieh Chen.
"""
import tensorflow as tf

from deeplab2.model import utils
from deeplab2.model.layers import activations
from deeplab2.model.layers import axial_layers
from deeplab2.model.layers import convolutions
from deeplab2.model.layers import squeeze_and_excite


class AxialBlock(tf.keras.layers.Layer):
  """An AxialBlock as a building block for an Axial-ResNet model.

  We implement the Axial-Block proposed in [1] in a general way that also
  includes convolutional residual blocks, such as the basic block and the
  bottleneck block (w/ and w/o Switchable Atrous Convolution).

  A basic block consists of two 3x3 convolutions and a residual connection. It
  is the main building block for wide-resnet variants.

  A bottleneck block consists of consecutive 1x1, 3x3, 1x1 convolutions and a
  residual connection. It is the main building block for standard resnet
  variants.

  An axial block consists of a 1x1 input convolution, a self-attention layer
  (either axial-attention or global attention), a 1x1 output convolution, and a
  residual connection. It is the main building block for axial-resnet variants.

  Note: We apply the striding in the first spatial operation (i.e. 3x3
  convolution or self-attention layer).
  """

  def __init__(self,
               filters_list,
               kernel_size=3,
               strides=1,
               atrous_rate=1,
               use_squeeze_and_excite=False,
               use_sac=False,
               bn_layer=tf.keras.layers.BatchNormalization,
               activation='relu',
               name=None,
               conv_kernel_weight_decay=0.0,
               basic_block_second_conv_atrous_rate=None,
               attention_type=None,
               axial_layer_config=None):
    """Initializes an AxialBlock.

    Args:
      filters_list: A list of filter numbers in the residual block. We currently
        support filters_list with two or three elements. Two elements specify
        the filters for two consecutive 3x3 convolutions, while three elements
        specify the filters for three convolutions (1x1, 3x3, and 1x1).
      kernel_size: The size of the convolution kernels (default: 3).
      strides: The strides of the block (default: 1).
      atrous_rate: The atrous rate of the 3x3 convolutions (default: 1). If this
        residual block is a basic block, it is recommendeded to specify correct
        basic_block_second_conv_atrous_rate for the second 3x3 convolution.
        Otherwise, the second conv will also use atrous rate, which might cause
        atrous inconsistency with different output strides, as tested in
        axial_block_groups_test.test_atrous_consistency_basic_block.
      use_squeeze_and_excite: A boolean flag indicating whether
        squeeze-and-excite (SE) is used.
      use_sac: A boolean, using the Switchable Atrous Convolution (SAC) or not.
      bn_layer: A tf.keras.layers.Layer that computes the normalization
        (default: tf.keras.layers.BatchNormalization).
      activation: A string specifying the activation function to apply.
      name: An string specifying the name of the layer (default: None).
      conv_kernel_weight_decay: A float, the weight decay for convolution
        kernels.
      basic_block_second_conv_atrous_rate: An integer, the atrous rate for the
        second convolution of basic block. This is necessary to ensure atrous
        consistency with different output_strides. Defaults to atrous_rate.
      attention_type: A string, type of attention to apply. Support 'axial' and
        'global'.
      axial_layer_config: A dict, an argument dictionary for the axial layer.

    Raises:
      ValueError: If filters_list does not have two or three elements.
      ValueError: If attention_type is not supported.
      ValueError: If double_global_attention is True in axial_layer_config.
    """
    super(AxialBlock, self).__init__(name=name)

    self._filters_list = filters_list
    self._strides = strides
    self._use_squeeze_and_excite = use_squeeze_and_excite
    self._bn_layer = bn_layer
    self._activate_fn = activations.get_activation(activation)
    self._attention_type = attention_type

    if axial_layer_config is None:
      axial_layer_config = {}

    if basic_block_second_conv_atrous_rate is None:
      basic_block_second_conv_atrous_rate = atrous_rate

    if len(filters_list) == 3:
      # Three consecutive convolutions: 1x1, 3x3, and 1x1.
      self._conv1_bn_act = convolutions.Conv2DSame(
          filters_list[0], 1, 'conv1_bn_act',
          use_bias=False,
          use_bn=True,
          bn_layer=bn_layer,
          activation=activation,
          conv_kernel_weight_decay=conv_kernel_weight_decay)

      if attention_type is None or attention_type.lower() == 'none':
        self._conv2_bn_act = convolutions.Conv2DSame(
            filters_list[1], kernel_size, 'conv2_bn_act',
            strides=strides,
            atrous_rate=atrous_rate,
            use_bias=False,
            use_bn=True,
            bn_layer=bn_layer,
            activation=activation,
            use_switchable_atrous_conv=use_sac,
            # We default to use global context in SAC if use_sac is True. This
            # setting is experimentally found effective.
            use_global_context_in_sac=use_sac,
            conv_kernel_weight_decay=conv_kernel_weight_decay)
      elif attention_type == 'axial':
        if 'double_global_attention' in axial_layer_config:
          if axial_layer_config['double_global_attention']:
            raise ValueError('Double_global_attention takes no effect in '
                             'AxialAttention2D.')
          del axial_layer_config['double_global_attention']
        self._attention = axial_layers.AxialAttention2D(
            strides=strides,
            filters=filters_list[1],
            name='attention',
            bn_layer=bn_layer,
            conv_kernel_weight_decay=conv_kernel_weight_decay,
            **axial_layer_config)
      elif attention_type == 'global':
        self._attention = axial_layers.GlobalAttention2D(
            strides=strides,
            filters=filters_list[1],
            name='attention',
            bn_layer=bn_layer,
            conv_kernel_weight_decay=conv_kernel_weight_decay,
            **axial_layer_config)
      else:
        raise ValueError(attention_type + ' is not supported.')

      # Here we apply a batch norm with gamma initialized at zero. This ensures
      # that at random initialization of the model, the skip connections
      # dominate all residual blocks. In this way, all the skip connections
      # construct an identity mapping that passes the gradients (without any
      # distortion from the randomly initialized blocks) to all residual blocks.
      # This trick helps training at early epochs.
      # Reference: "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour".
      # https://arxiv.org/abs/1706.02677
      self._conv3_bn = convolutions.Conv2DSame(
          filters_list[2], 1, 'conv3_bn',
          use_bias=False,
          use_bn=True,
          bn_layer=bn_layer,
          bn_gamma_initializer='zeros',
          activation='none',
          conv_kernel_weight_decay=conv_kernel_weight_decay)
    elif len(filters_list) == 2:
      # Two consecutive convolutions: 3x3 and 3x3.
      self._conv1_bn_act = convolutions.Conv2DSame(
          filters_list[0], kernel_size, 'conv1_bn_act',
          strides=strides,
          atrous_rate=atrous_rate,
          use_bias=False,
          use_bn=True,
          bn_layer=bn_layer,
          activation=activation,
          use_switchable_atrous_conv=use_sac,
          use_global_context_in_sac=use_sac,
          conv_kernel_weight_decay=conv_kernel_weight_decay)
      # Here we apply a batch norm with gamma initialized at zero. This ensures
      # that at random initialization of the model, the skip connections
      # dominate all residual blocks. In this way, all the skip connections
      # construct an identity mapping that passes the gradients (without any
      # distortion from the randomly initialized blocks) to all residual blocks.
      # This trick helps training at early epochs.
      # Reference: "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour".
      # https://arxiv.org/abs/1706.02677
      self._conv2_bn = convolutions.Conv2DSame(
          filters_list[1], kernel_size, 'conv2_bn',
          strides=1,
          atrous_rate=basic_block_second_conv_atrous_rate,
          use_bias=False,
          use_bn=True,
          bn_layer=bn_layer,
          bn_gamma_initializer='zeros',
          activation='none',
          use_switchable_atrous_conv=use_sac,
          use_global_context_in_sac=use_sac,
          conv_kernel_weight_decay=conv_kernel_weight_decay)
    else:
      raise ValueError('Expect filters_list to have length 2 or 3; got %d' %
                       len(filters_list))

    if self._use_squeeze_and_excite:
      self._squeeze_and_excite = squeeze_and_excite.SimplifiedSqueezeAndExcite(
          filters_list[-1])
    self._conv_kernel_weight_decay = conv_kernel_weight_decay

  def build(self, input_shape_list):
    input_tensor_shape = input_shape_list[0]
    self._shortcut = None
    if input_tensor_shape[3] != self._filters_list[-1]:
      self._shortcut = convolutions.Conv2DSame(
          self._filters_list[-1], 1, 'shortcut',
          strides=self._strides,
          use_bias=False,
          use_bn=True,
          bn_layer=self._bn_layer,
          activation='none',
          conv_kernel_weight_decay=self._conv_kernel_weight_decay)

  def call(self, inputs):
    """Performs a forward pass.

    We have to define drop_path_random_mask outside the layer call and pass it
    into the layer, because recompute_grad (gradient checkpointing) does not
    allow any randomness within the function call. In addition, recompute_grad
    only supports float tensors as inputs. For this reason, the training flag
    should be also passed as a float tensor. For the same reason, we cannot
    support passing drop_path_random_mask as None. Instead, we ask the users to
    pass only the first two tensors when drop path is not used.

    Args:
      inputs: A tuple of 2 or 3 tensors, containing
        input_tensor should be an input tensor of type tf.Tensor with shape
          [batch, height, width, channels].
        float_tensor_training should be a float tensor of 0.0 or 1.0, whether
          the model is in training mode.
        (optional) drop_path_random_mask is a drop path random mask of type
          tf.Tensor with shape [batch, 1, 1, 1].

    Returns:
      outputs: two tensors. The first tensor does not use the last activation
        function. The second tensor uses the activation. We return non-activated
        output to support MaX-DeepLab which uses non-activated feature for the
        stacked decoders.

    Raises:
      ValueError: If the length of inputs is not 2 or 3.
    """
    if len(inputs) not in (2, 3):
      raise ValueError('The length of inputs should be either 2 or 3.')

    # Unpack the inputs.
    input_tensor, float_tensor_training, drop_path_random_mask = (
        utils.pad_sequence_with_none(inputs, target_length=3))

    # Recompute_grad takes only float tensors as inputs. It does not allow
    # bools or boolean tensors. For this reason, we cast training to a float
    # tensor outside this call, and now we cast it back to a boolean tensor.
    training = tf.cast(float_tensor_training, tf.bool)

    shortcut = input_tensor
    if self._shortcut is not None:
      shortcut = self._shortcut(shortcut, training=training)
    elif self._strides != 1:
      shortcut = shortcut[:, ::self._strides, ::self._strides, :]

    if len(self._filters_list) == 3:
      x = self._conv1_bn_act(input_tensor, training=training)
      if (self._attention_type is None or
          self._attention_type.lower() == 'none'):
        x = self._conv2_bn_act(x, training=training)
      else:
        x = self._attention(x, training=training)
        x = self._activate_fn(x)
      x = self._conv3_bn(x, training=training)
    if len(self._filters_list) == 2:
      x = self._conv1_bn_act(input_tensor, training=training)
      x = self._conv2_bn(x, training=training)

    if self._use_squeeze_and_excite:
      x = self._squeeze_and_excite(x)

    if drop_path_random_mask is not None:
      x = x * drop_path_random_mask
    x = x + shortcut
    return x, self._activate_fn(x)