# Copyright 2020 The TensorFlow 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
#
#     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.
# ==============================================================================
"""Implementation of SpineNet-MBConv model.

SpineNet with MBConv blocks.

X. Du, T-Y. Lin, P. Jin, G. Ghiasi, M. Tan, Y. Cui, Q. V. Le, X. Song
SpineNet: Learning Scale-Permuted Backbone for Recognition and Localization
https://arxiv.org/abs/1912.05027
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
from absl import logging
import tensorflow.compat.v1 as tf

from modeling.architecture import nn_blocks
from modeling.architecture import nn_ops
from ops import spatial_transform_ops

DEFAULT_EXPAND_RATIO = 6

FILTER_SIZE_MAP = {
    0: 8,
    1: 16,
    2: 24,
    3: 40,
    4: 80,
    5: 112,
    6: 112,
    7: 112,
}

# The static SpineNet architecture discovered by NAS.
# Each element represents a specification of a building block:
#   (block_level, block_fn, (input_offset0, input_offset1), is_output).
SPINENET_BLOCK_SPECS = [
    (2, 'mbconv', (0, 1), False),
    (2, 'mbconv', (1, 2), False),
    (4, 'mbconv', (1, 2), False),
    (3, 'mbconv', (3, 4), False),
    (4, 'mbconv', (3, 5), False),
    (6, 'mbconv', (4, 6), False),
    (4, 'mbconv', (4, 6), False),
    (5, 'mbconv', (7, 8), False),
    (7, 'mbconv', (7, 9), False),
    (5, 'mbconv', (9, 10), False),
    (5, 'mbconv', (9, 11), False),
    (4, 'mbconv', (6, 11), True),
    (3, 'mbconv', (5, 11), True),
    (5, 'mbconv', (8, 13), True),
    (7, 'mbconv', (6, 15), True),
    (6, 'mbconv', (13, 15), True),
]


SCALING_MAP = {
    '49': {
        'endpoints_num_filters': 48,
        'filter_size_scale': 1.0,
        'block_repeats': 1,
    },
    '49S': {
        'endpoints_num_filters': 40,
        'filter_size_scale': 0.65,
        'block_repeats': 1,
    },
    '49XS': {
        'endpoints_num_filters': 24,
        'filter_size_scale': 0.6,
        'block_repeats': 1,
    },
}


class BlockSpec(object):
  """A container class that specifies the block configuration for SpineNet."""

  def __init__(self, level, block_fn, input_offsets, is_output):
    self.level = level
    self.block_fn = block_fn
    self.input_offsets = input_offsets
    self.is_output = is_output


def build_block_specs(block_specs=None):
  """Builds the list of BlockSpec objects for SpineNet."""
  if not block_specs:
    block_specs = SPINENET_BLOCK_SPECS
  logging.info('Building SpineNet block specs: %s', block_specs)
  return [BlockSpec(*b) for b in block_specs]


def block_group(inputs,
                in_filters,
                out_filters,
                strides,
                expand_ratio,
                block_repeats,
                se_ratio=0.2,
                batch_norm_activation=nn_ops.BatchNormActivation(),
                dropblock=nn_ops.Dropblock(),
                drop_connect_rate=None,
                data_format='channels_last',
                name=None,
                is_training=False):
  """Creates one group of blocks for Mobile SpineNet."""
  # Apply strides only to the first block in block_group.
  inputs = nn_blocks.mbconv_block(
      inputs,
      in_filters,
      out_filters,
      expand_ratio,
      strides,
      se_ratio=se_ratio,
      batch_norm_activation=batch_norm_activation,
      dropblock=dropblock,
      drop_connect_rate=drop_connect_rate,
      data_format=data_format,
      is_training=is_training)
  for _ in range(1, block_repeats):
    inputs = nn_blocks.mbconv_block(
        inputs,
        out_filters,
        out_filters,
        expand_ratio,
        1,  # strides
        se_ratio=se_ratio,
        batch_norm_activation=batch_norm_activation,
        dropblock=dropblock,
        drop_connect_rate=drop_connect_rate,
        data_format=data_format,
        is_training=is_training)
  return tf.identity(inputs, name)


def resample_with_sepconv(feat,
                          target_width,
                          target_num_filters,
                          use_native_resize_op=False,
                          batch_norm_activation=nn_ops.BatchNormActivation(),
                          data_format='channels_last',
                          name=None,
                          is_training=False):
  """Match resolution and feature dimension to the target block."""
  _, height, width, num_filters = feat.get_shape().as_list()
  if width is None or num_filters is None:
    raise ValueError('Shape of feat is None (shape:{}).'.format(feat.shape))

  with tf.variable_scope('resample_with_sepconv_{}'.format(name)):
    # Down-sample.
    if width > target_width:
      if width % target_width != 0:
        raise ValueError('width ({}) is not divisible by '
                         'target_width ({}).'.format(width, target_width))

      while width > target_width:
        feat = nn_ops.depthwise_conv2d_fixed_padding(
            inputs=feat, kernel_size=3, strides=2, data_format=data_format)
        feat = batch_norm_activation(feat, is_training=is_training)
        width /= 2

    # Up-sample with NN interpolation.
    elif width < target_width:
      if target_width % width != 0:
        raise ValueError('target_wdith ({}) is not divisible by '
                         'width ({}).'.format(target_width, width))
      scale = target_width // width
      if use_native_resize_op:
        feat = tf.image.resize_nearest_neighbor(feat,
                                                [height * scale, width * scale])
      else:
        feat = spatial_transform_ops.nearest_upsampling(feat, scale=scale)

    # Match feature dimension to the target block.
    feat = nn_ops.conv2d_fixed_padding(
        inputs=feat,
        filters=target_num_filters,
        kernel_size=1,
        strides=1,
        data_format=data_format)
    feat = batch_norm_activation(feat, relu=False, is_training=is_training)

  return feat


def get_drop_connect_rate(init_rate, i, n):
  """Get drop connect rate for the ith block."""
  if (init_rate is not None) and (init_rate > 0 and init_rate < 1):
    dc_rate = init_rate * float(i + 1) / n
    logging.info('Drop connect rate %f for block_%d.', dc_rate, i)
  else:
    dc_rate = None
  return dc_rate


class SpineNetMBConv(object):
  """Class to build SpineNet family models with MBConv blocks."""

  def __init__(self,
               min_level=3,
               max_level=7,
               block_specs=build_block_specs(),
               endpoints_num_filters=48,
               use_native_resize_op=False,
               se_ratio=0.2,
               block_repeats=1,
               filter_size_scale=1.0,
               activation='swish',
               batch_norm_activation=nn_ops.BatchNormActivation(
                   activation='swish'),
               init_drop_connect_rate=None,
               data_format='channels_last'):
    """SpineNetMBConv initialization function.

    Args:
      min_level: `int` minimum level in SpineNet endpoints.
      max_level: `int` maximum level in SpineNet endpoints.