# 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.

"""This file contains code to build a DeepLabV3Plus.

Reference:
  - [Encoder-Decoder with Atrous Separable Convolution for Semantic Image
      Segmentation](https://arxiv.org/pdf/1802.02611.pdf)
"""

import tensorflow as tf

from deeplab2 import common
from deeplab2.model import utils
from deeplab2.model.decoder import aspp
from deeplab2.model.layers import convolutions


layers = tf.keras.layers


class DeepLabV3Plus(tf.keras.layers.Layer):
  """A DeepLabV3+ decoder model.

  This model takes in low- and high-level features from an encoder and performs
  multi-scale context aggregation with the help of an ASPP layer on high-level
  features. These are concatenated with the low-level features and used as input
  to the classification head that is used to predict a semantic segmentation.
  """

  def __init__(self,
               decoder_options,
               deeplabv3plus_options,
               bn_layer=tf.keras.layers.BatchNormalization):
    """Creates a DeepLabV3+ decoder of type tf.keras.layers.Layer.

    Args:
      decoder_options: Decoder options as defined in config_pb2.DecoderOptions.
      deeplabv3plus_options: Model options as defined in
        config_pb2.ModelOptions.DeeplabV3PlusOptions.
      bn_layer: An optional tf.keras.layers.Layer that computes the
        normalization (default: tf.keras.layers.BatchNormalization).
    """
    super(DeepLabV3Plus, self).__init__(name='DeepLabv3Plus')

    self._high_level_feature_name = decoder_options.feature_key
    self._low_level_feature_name = deeplabv3plus_options.low_level.feature_key
    self._aspp = aspp.ASPP(decoder_options.aspp_channels,
                           decoder_options.atrous_rates,
                           bn_layer=bn_layer)

    # Layers for low-level feature transformation.
    self._project_conv_bn_act = convolutions.Conv2DSame(
        deeplabv3plus_options.low_level.channels_project,
        kernel_size=1,
        name='project_conv_bn_act',
        use_bias=False,
        use_bn=True,
        bn_layer=bn_layer,
        activation='relu')

    # Layers for fusing low- and high-level features.
    self._fuse = convolutions.StackedConv2DSame(
        conv_type='depthwise_separable_conv',
        num_layers=2,
        output_channels=decoder_options.decoder_channels,
        kernel_size=3,
        name='fuse',
        use_bias=False,
        use_bn=True,
        bn_layer=bn_layer,
        activation='relu')

    self._final_conv = convolutions.Conv2DSame(
        deeplabv3plus_options.num_classes, kernel_size=1, name='final_conv')

  def reset_pooling_layer(self):
    """Resets the ASPP pooling layer to global average pooling."""
    self._aspp.reset_pooling_layer()

  def set_pool_size(self, pool_size):
    """Sets the pooling size of the ASPP pooling layer.

    Args:
      pool_size: A tuple specifying the pooling size of the ASPP pooling layer.
    """
    self._aspp.set_pool_size(pool_size)

  def get_pool_size(self):
    return self._aspp.get_pool_size()

  @property
  def checkpoint_items(self):
    items = {
        common.CKPT_DEEPLABV3PLUS_ASPP: self._aspp,
        common.CKPT_DEEPLABV3PLUS_PROJECT_CONV_BN_ACT:
            self._project_conv_bn_act,
        common.CKPT_DEEPLABV3PLUS_FUSE: self._fuse,
        common.CKPT_SEMANTIC_LAST_LAYER: self._final_conv,
    }
    return items

  def call(self, features, training=False):
    """Performs a forward pass.

    Args:
      features: An input dict of tf.Tensor with shape [batch, height, width,
        channels]. Different keys should point to different features extracted
        by the encoder, e.g. low-level or high-level features.
      training: A boolean flag indicating whether training behavior should be
        used (default: False).

    Returns:
      A dictionary containing the semantic prediction under key
      common.PRED_SEMANTIC_LOGITS_KEY.
    """
    low_level_features = features[self._low_level_feature_name]
    high_level_features = features[self._high_level_feature_name]

    high_level_features = self._aspp(high_level_features, training=training)

    low_level_features = self._project_conv_bn_act(low_level_features,
                                                   training=training)

    target_h = tf.shape(low_level_features)[1]
    target_w = tf.shape(low_level_features)[2]

    high_level_features = utils.resize_align_corners(
        high_level_features, [target_h, target_w])
    x = tf.concat([high_level_features, low_level_features], 3)
    x = self._fuse(x)

    return {common.PRED_SEMANTIC_LOGITS_KEY: self._final_conv(x)}