# 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.
# ==============================================================================
"""Functions to generate bidirectional feature pyramids based on image features.

Provides bidirectional feature pyramid network (BiFPN) generators that can be
used to build object detection feature extractors, as proposed by Tan et al.
See https://arxiv.org/abs/1911.09070 for more details.
"""
import collections
import functools
from six.moves import range
from six.moves import zip
import tensorflow as tf

from ..utils import bifpn_utils


def _create_bifpn_input_config(fpn_min_level,
                               fpn_max_level,
                               input_max_level,
                               level_scales=None):
  """Creates a BiFPN input config for the input levels from a backbone network.

  Args:
    fpn_min_level: the minimum pyramid level (highest feature map resolution) to
      use in the BiFPN.
    fpn_max_level: the maximum pyramid level (lowest feature map resolution) to
      use in the BiFPN.
    input_max_level: the maximum pyramid level that will be provided as input to
      the BiFPN. Accordingly, the BiFPN will compute additional pyramid levels
      from input_max_level, up to the desired fpn_max_level.
    level_scales: a list of pyramid level scale factors. If 'None', each level's
      scale is set to 2^level by default, which corresponds to each successive
      feature map scaling by a factor of 2.

  Returns:
    A list of dictionaries for each feature map expected as input to the BiFPN,
    where each has entries for the feature map 'name' and 'scale'.
  """
  if not level_scales:
    level_scales = [2**i for i in range(fpn_min_level, fpn_max_level + 1)]

  bifpn_input_params = []
  for i in range(fpn_min_level, min(fpn_max_level, input_max_level) + 1):
    bifpn_input_params.append({
        'name': '0_up_lvl_{}'.format(i),
        'scale': level_scales[i - fpn_min_level]
    })

  return bifpn_input_params


def _get_bifpn_output_node_names(fpn_min_level, fpn_max_level, node_config):
  """Returns a list of BiFPN output node names, given a BiFPN node config.

  Args:
    fpn_min_level: the minimum pyramid level (highest feature map resolution)
      used by the BiFPN.
    fpn_max_level: the maximum pyramid level (lowest feature map resolution)
      used by the BiFPN.
    node_config: the BiFPN node_config, a list of dictionaries corresponding to
      each node in the BiFPN computation graph, where each entry should have an
      associated 'name'.

  Returns:
    A list of strings corresponding to the names of the output BiFPN nodes.
  """
  num_output_nodes = fpn_max_level - fpn_min_level + 1
  return [node['name'] for node in node_config[-num_output_nodes:]]


def _create_bifpn_node_config(bifpn_num_iterations,
                              bifpn_num_filters,
                              fpn_min_level,
                              fpn_max_level,
                              input_max_level,
                              bifpn_node_params=None,
                              level_scales=None):
  """Creates a config specifying a bidirectional feature pyramid network.

  Args:
    bifpn_num_iterations: the number of top-down bottom-up feature computations
      to repeat in the BiFPN.
    bifpn_num_filters: the number of filters (channels) for every feature map
      used in the BiFPN.
    fpn_min_level: the minimum pyramid level (highest feature map resolution) to
      use in the BiFPN.
    fpn_max_level: the maximum pyramid level (lowest feature map resolution) to
      use in the BiFPN.
    input_max_level: the maximum pyramid level that will be provided as input to
      the BiFPN. Accordingly, the BiFPN will compute additional pyramid levels
      from input_max_level, up to the desired fpn_max_level.
    bifpn_node_params: If not 'None', a dictionary of additional default BiFPN
      node parameters that will be applied to all BiFPN nodes.
    level_scales: a list of pyramid level scale factors. If 'None', each level's
      scale is set to 2^level by default, which corresponds to each successive
      feature map scaling by a factor of 2.

  Returns:
    A list of dictionaries used to define nodes in the BiFPN computation graph,
    as proposed by EfficientDet, Tan et al (https://arxiv.org/abs/1911.09070).
    Each node's entry has the corresponding keys:
      name: String. The name of this node in the BiFPN. The node name follows
        the format '{bifpn_iteration}_{dn|up}_lvl_{pyramid_level}', where 'dn'
        or 'up' refers to whether the node is in the top-down or bottom-up
        portion of a single BiFPN iteration.
      scale: the scale factor for this node, by default 2^level.
      inputs: A list of names of nodes which are inputs to this node.
      num_channels: The number of channels for this node.
      combine_method: String. Name of the method used to combine input
        node feature maps, 'fast_attention' by default for nodes which have more
        than one input. Otherwise, 'None' for nodes with only one input node.
      input_op: A (partial) function which is called to construct the layers
        that will be applied to this BiFPN node's inputs. This function is
        called with the arguments:
          input_op(name, input_scale, input_num_channels, output_scale,
                   output_num_channels, conv_hyperparams, is_training,
                   freeze_batchnorm)
      post_combine_op: A (partial) function which is called to construct the
        layers that will be applied to the result of the combine operation for
        this BiFPN node. This function will be called with the arguments:
          post_combine_op(name, conv_hyperparams, is_training, freeze_batchnorm)
        If 'None', then no layers will be applied after the combine operation
        for this node.
  """
  if not level_scales:
    level_scales = [2**i for i in range(fpn_min_level, fpn_max_level + 1)]

  default_node_params = {
      'num_channels':
          bifpn_num_filters,
      'combine_method':
          'fast_attention',
      'input_op':
          functools.partial(
              _create_bifpn_resample_block, downsample_method='max_pooling'),
      'post_combine_op':
          functools.partial(
              bifpn_utils.create_conv_block,
              num_filters=bifpn_num_filters,
              kernel_size=3,
              strides=1,
              padding='SAME',
              use_separable=True,
              apply_batchnorm=True,
              apply_activation=True,
              conv_bn_act_pattern=False),
  }
  if bifpn_node_params:
    default_node_params.update(bifpn_node_params)

  bifpn_node_params = []
  # Create additional base pyramid levels not provided as input to the BiFPN.
  # Note, combine_method and post_combine_op are set to None for additional
  # base pyramid levels because they do not combine multiple input BiFPN nodes.
  for i in range(input_max_level + 1, fpn_max_level + 1):
    node_params = dict(default_node_params)
    node_params.update({
        'name': '0_up_lvl_{}'.format(i),
        'scale': level_scales[i - fpn_min_level],
        'inputs': ['0_up_lvl_{}'.format(i - 1)],
        'combine_method': None,
        'post_combine_op': None,
    })
    bifpn_node_params.append(node_params)

  for i in range(bifpn_num_iterations):
    # The first bottom-up feature pyramid (which includes the input pyramid
    # levels from the backbone network and the additional base pyramid levels)
    # is indexed at 0. So, the first top-down bottom-up pass of the BiFPN is
    # indexed from 1, and repeated for bifpn_num_iterations iterations.
    bifpn_i = i + 1

    # Create top-down nodes.
    for level_i in reversed(range(fpn_min_level, fpn_max_level)):
      inputs = []
      # BiFPN nodes in the top-down pass receive input from the corresponding
      # level from the previous BiFPN iteration's bottom-up pass, except for the
      # bottom-most (min) level node, which is computed once in the initial
      # bottom-up pass, and is afterwards only computed in each top-down pass.
      if level_i > fpn_min_level or bifpn_i == 1:
        inputs.append('{}_up_lvl_{}'.format(bifpn_i - 1, level_i))
      else:
        inputs.append('{}_dn_lvl_{}'.format(bifpn_i - 1, level_i))
      inputs.append(bifpn_node_params[-1]['name'])
      node_params = dict(default_node_params)
      node_params.update({
          'name': '{}_dn_lvl_{}'.format(bifpn_i, level_i),
          'scale': level_scales[level_i - fpn_min_level],
          'inputs': inputs
      })
      bifpn_node_params.append(node_params)

    # Create bottom-up nodes.
    for level_i in range(fpn_min_level + 1, fpn_max_level + 1):
      # BiFPN nodes in the bottom-up pass receive input from the corresponding
      # level from the preceding top-down pass, except for the top (max) level
      # which does not have a corresponding node in the top-down pass.
      inputs = ['{}_up_lvl_{}'.format(bifpn_i - 1, level_i)]
      if level_i < fpn_max_level:
        inputs.append('{}_dn_lvl_{}'.format(bifpn_i, level_i))
      inputs.append(bifpn_node_params[-1]['name'])
      node_params = dict(default_node_params)
      node_params.update({
          'name': '{}_up_lvl_{}'.format(bifpn_i, level_i),
          'scale': level_scales[level_i - fpn_min_level],
          'inputs': inputs
      })
      bifpn_node_params.append(node_params)

  return bifpn_node_params


def _create_bifpn_resample_block(name,
                                 input_scale,
                                 input_num_channels,
                                 output_scale,
                                 output_num_channels,
                                 conv_hyperparams,
                                 is_training,
                                 freeze_batchnorm,
                                 downsample_method=None,
                                 use_native_resize_op=False,
                                 maybe_apply_1x1_conv=True,
                                 apply_1x1_pre_sampling=True,
                                 apply_1x1_post_sampling=False):
  """Creates resample block layers for input feature maps to BiFPN nodes.

  Args:
    name: String. Name used for this block of layers.
    input_scale: Scale factor of the input feature map.
    input_num_channels: Number of channels in the input feature map.
    output_scale: Scale factor of the output feature map.
    output_num_channels: Number of channels in the output feature map.
    conv_hyperparams: A `hyperparams_builder.KerasLayerHyperparams` object
      containing hyperparameters for convolution ops.
    is_training: Indicates whether the feature generator is in training mode.
    freeze_batchnorm: Bool. Whether to freeze batch norm parameters during
      training or not. When training with a small batch size (e.g. 1), it is
      desirable to freeze batch norm update and use pretrained batch norm
      params.
    downsample_method: String. Method to use when downsampling feature maps.
    use_native_resize_op: Bool. Whether to use the native resize up when
      upsampling feature maps.
    maybe_apply_1x1_conv: Bool. If 'True', a 1x1 convolution will only be
      applied if the input_num_channels differs from the output_num_channels.
    apply_1x1_pre_sampling: Bool. Whether a 1x1 convolution will be applied to
      the input feature map before the up/down-sampling operation.
    apply_1x1_post_sampling: Bool. Whether a 1x1 convolution will be applied to
      the input feature map after the up/down-sampling operation.

  Returns:
    A list of layers which may be applied to the input feature maps in order to
    compute feature maps with the specified scale and number of channels.
  """
  # By default, 1x1 convolutions are only applied before sampling when the
  # number of input and output channels differ.
  if maybe_apply_1x1_conv and output_num_channels == input_num_channels:
    apply_1x1_pre_sampling = False
    apply_1x1_post_sampling = False

  apply_bn_for_resampling = True
  layers = []
  if apply_1x1_pre_sampling:
    layers.extend(
        bifpn_utils.create_conv_block(
            name=name + '1x1_pre_sample/',
            num_filters=output_num_channels,
            kernel_size=1,
            strides=1,
            padding='SAME',
            use_separable=False,
            apply_batchnorm=apply_bn_for_resampling,
            apply_activation=False,
            conv_hyperparams=conv_hyperparams,
            is_training=is_training,
            freeze_batchnorm=freeze_batchnorm))

  layers.extend(
      bifpn_utils.create_resample_feature_map_ops(input_scale, output_scale,
                                                  downsample_method,
                                                  use_native_resize_op,
                                                  conv_hyperparams, is_training,
                                                  freeze_batchnorm, name))

  if apply_1x1_post_sampling:
    layers.extend(
        bifpn_utils.create_conv_block(
            name=name + '1x1_post_sample/',
            num_filters=output_num_channels,
            kernel_size=1,
            strides=1,
            padding='SAME',
            use_separable=False,
            apply_batchnorm=apply_bn_for_resampling,
            apply_activation=False,
            conv_hyperparams=conv_hyperparams,
            is_training=is_training,
            freeze_batchnorm=freeze_batchnorm))

  return layers


def _create_bifpn_combine_op(num_inputs, name, combine_method):
  """Creates a BiFPN output config, a list of the output BiFPN node names.

  Args:
    num_inputs: The number of inputs to this combine operation.
    name: String. The name of this combine operation.
    combine_method: String. The method used to combine input feature maps.

  Returns:
    A function which may be called with a list of num_inputs feature maps
    and which will return a single feature map.
  """

  combine_op = None
  if num_inputs < 1:
    raise ValueError('Expected at least 1 input for BiFPN combine.')
  elif num_inputs == 1:
    combine_op = lambda x: x[0]
  else:
    combine_op = bifpn_utils.BiFPNCombineLayer(
        combine_method=combine_method, name=name)
  return combine_op


class KerasBiFpnFeatureMaps(tf.keras.Model):
  """Generates Keras based BiFPN feature maps from an input feature map pyramid.

  A Keras model that generates multi-scale feature maps for detection by
  iteratively computing top-down and bottom-up feature pyramids, as in the
  EfficientDet paper by Tan et al, see arxiv.org/abs/1911.09070 for details.
  """

  def __init__(self,
               bifpn_num_iterations,
               bifpn_num_filters,
               fpn_min_level,
               fpn_max_level,
               input_max_level,
               is_training,
               conv_hyperparams,
               freeze_batchnorm,
               bifpn_node_params=None,
               name=None):
    """Constructor.

    Args:
      bifpn_num_iterations: The number of top-down bottom-up iterations.
      bifpn_num_filters: The number of filters (channels) to be used for all
        feature maps in this BiFPN.
      fpn_min_level: The minimum pyramid level (highest feature map resolution)
        to use in the BiFPN.
      fpn_max_level: The maximum pyramid level (lowest feature map resolution)
        to use in the BiFPN.
      input_max_level: The maximum pyramid level that will be provided as input
        to the BiFPN. Accordingly, the BiFPN will compute any additional pyramid
        levels from input_max_level up to the desired fpn_max_level, with each
        successivel level downsampling by a scale factor of 2 by default.
      is_training: Indicates whether the feature generator is in training mode.
      conv_hyperparams: A `hyperparams_builder.KerasLayerHyperparams` object
        containing hyperparameters for convolution ops.
      freeze_batchnorm: Bool. Whether to freeze batch norm parameters during
        training or not. When training with a small batch size (e.g. 1), it is
        desirable to freeze batch norm update and use pretrained batch norm
        params.
      bifpn_node_params: An optional dictionary that may be used to specify
        default parameters for BiFPN nodes, without the need to provide a custom
        bifpn_node_config. For example, if '{ combine_method: 'sum' }', then all
        BiFPN nodes will combine input feature maps by summation, rather than
        by the default fast attention method.
      name: A string name scope to assign to the model. If 'None', Keras
        will auto-generate one from the class name.
    """
    super(KerasBiFpnFeatureMaps, self).__init__(name=name)
    bifpn_node_config = _create_bifpn_node_config(
        bifpn_num_iterations, bifpn_num_filters, fpn_min_level, fpn_max_level,
        input_max_level, bifpn_node_params)
    bifpn_input_config = _create_bifpn_input_config(
        fpn_min_level, fpn_max_level, input_max_level)
    bifpn_output_node_names = _get_bifpn_output_node_names(
        fpn_min_level, fpn_max_level, bifpn_node_config)

    self.bifpn_node_config = bifpn_node_config
    self.bifpn_output_node_names = bifpn_output_node_names
    self.node_input_blocks = []
    self.node_combine_op = []
    self.node_post_combine_block = []

    all_node_params = bifpn_input_config
    all_node_names = [node['name'] for node in all_node_params]
    for node_config in bifpn_node_config:
      # Maybe transform and/or resample input feature maps.
      input_blocks = []
      for input_name in node_config['inputs']:
        if input_name not in all_node_names:
          raise ValueError(
              'Input feature map ({}) does not exist:'.format(input_name))
        input_index = all_node_names.index(input_name)
        input_params = all_node_params[input_index]
        input_block = node_config['input_op'](
            name='{}/input_{}/'.format(node_config['name'], input_name),
            input_scale=input_params['scale'],
            input_num_channels=input_params.get('num_channels', None),
            output_scale=node_config['scale'],
            output_num_channels=node_config['num_channels'],
            conv_hyperparams=conv_hyperparams,
            is_training=is_training,
            freeze_batchnorm=freeze_batchnorm)
        input_blocks.append((input_index, input_block))

      # Combine input feature maps.
      combine_op = _create_bifpn_combine_op(
          num_inputs=len(input_blocks),
          name=(node_config['name'] + '/combine'),
          combine_method=node_config['combine_method'])

      # Post-combine layers.
      post_combine_block = []
      if node_config['post_combine_op']:
        post_combine_block.extend(node_config['post_combine_op'](
            name=node_config['name'] + '/post_combine/',
            conv_hyperparams=conv_hyperparams,
            is_training=is_training,
            freeze_batchnorm=freeze_batchnorm))

      self.node_input_blocks.append(input_blocks)
      self.node_combine_op.append(combine_op)
      self.node_post_combine_block.append(post_combine_block)
      all_node_params.append(node_config)
      all_node_names.append(node_config['name'])

  def call(self, feature_pyramid):
    """Compute BiFPN feature maps from input feature pyramid.

    Executed when calling the `.__call__` method on input.

    Args:
      feature_pyramid: list of tuples of (tensor_name, image_feature_tensor).

    Returns:
      feature_maps: an OrderedDict mapping keys (feature map names) to
        tensors where each tensor has shape [batch, height_i, width_i, depth_i].
    """
    feature_maps = [el[1] for el in feature_pyramid]
    output_feature_maps = [None for node in self.bifpn_output_node_names]

    for index, node in enumerate(self.bifpn_node_config):
      node_scope = 'node_{:02d}'.format(index)
      with tf.name_scope(node_scope):
        # Apply layer blocks to this node's input feature maps.
        input_block_results = []
        for input_index, input_block in self.node_input_blocks[index]:
          block_result = feature_maps[input_index]
          for layer in input_block:
            block_result = layer(block_result)
          input_block_results.append(block_result)

        # Combine the resulting feature maps.
        node_result = self.node_combine_op[index](input_block_results)

        # Apply post-combine layer block if applicable.
        for layer in self.node_post_combine_block[index]:
          node_result = layer(node_result)

        feature_maps.append(node_result)

        if node['name'] in self.bifpn_output_node_names:
          index = self.bifpn_output_node_names.index(node['name'])
          output_feature_maps[index] = node_result

    return collections.OrderedDict(
        zip(self.bifpn_output_node_names, output_feature_maps))
