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

"""Featurizer layers for Pointpillars."""

from typing import Any, List, Mapping, Optional, Tuple

import numpy as np
import tensorflow as tf

from official.projects.pointpillars.modeling import layers
from official.projects.pointpillars.utils import utils


@tf.keras.utils.register_keras_serializable(package='Vision')
class Featurizer(tf.keras.layers.Layer):
  """The featurizer to convert pillars to a BEV pseudo image.

  The implementation is from the network architecture of PointPillars
  (https://arxiv.org/pdf/1812.05784.pdf). It extract features from pillar
  tensors then scatter them back to bird-eye-view (BEV) image using indices.

  Notations:
    B: batch size
    H: height of the BEV image
    W: width of the BEV image
    P: number of pillars in an example
    N: number of points in a pillar
    D: number of features in a point
    C: channels of the BEV image
  """

  def __init__(
      self,
      image_size: Tuple[int, int],
      pillars_size: Tuple[int, int, int],
      train_batch_size: int,
      eval_batch_size: int,
      num_blocks: int,
      num_channels: int,
      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
      **kwargs):
    """Initialize the featurizer.

    Args:
      image_size: A [int, int] tuple to define the [H, W] of BEV image.
      pillars_size: A [int, int, int] tuple to define the [P, N, D] of pillars.
      train_batch_size: An `int` training batch size per replica.
      eval_batch_size: An `int` evaluation batch size per replica.
      num_blocks: An `int` number of blocks for extracting features.
      num_channels: An `int` number channels of the BEV image.
      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
        block layers. Default to None.
      **kwargs: Additional keyword arguments to be passed.
    """
    super(Featurizer, self).__init__(**kwargs)

    self._config_dict = {
        'image_size': image_size,
        'pillars_size': pillars_size,
        'train_batch_size': train_batch_size,
        'eval_batch_size': eval_batch_size,
        'num_blocks': num_blocks,
        'num_channels': num_channels,
        'kernel_regularizer': kernel_regularizer,
    }
    self._image_shape = [image_size[0], image_size[1], num_channels]

    utils.assert_channels_last()

  def build(self, input_specs: List[tf.TensorShape]):
    """Creates variables for the featurizer."""
    self._blocks = []
    for _ in range(self._config_dict['num_blocks']):
      self._blocks.append(
          layers.ConvBlock(
              filters=self._config_dict['num_channels'],
              kernel_size=1,
              strides=1,
              kernel_regularizer=self._config_dict['kernel_regularizer']))

    # These batch_dims are [B, P, 1] tensors that could be created before
    # call(). They will be used for tf.scatter_nd to convert pillars to BEV
    # images. Because tf.scatter_nd requires a concrete batch size, we need to
    # prepare all possibilities of batch size for train, eval and test mode.
    self._train_batch_dims = self._get_batch_dims(
        self._config_dict['train_batch_size'])
    self._eval_batch_dims = self._get_batch_dims(
        self._config_dict['eval_batch_size'])
    self._test_batch_dims = self._get_batch_dims(1)

    super(Featurizer, self).build(input_specs)

  def _get_batch_dims(self, batch_size: int) -> tf.Tensor:
    p = self._config_dict['pillars_size'][0]
    batch_dims = np.indices([batch_size, p])[0]
    batch_dims = tf.convert_to_tensor(batch_dims, dtype=tf.int32)
    batch_dims = tf.expand_dims(batch_dims, axis=-1)
    return batch_dims

  def _get_batch_size_and_dims(self,
                               training: bool = None) -> Tuple[int, tf.Tensor]:
    # We use training as a ternary indicator, None for test mode.
    # Test mode will be used for saving model and model inference.
    if training is None:
      batch_size = 1
      batch_dims = self._test_batch_dims
    else:
      if training:
        batch_size = self._config_dict['train_batch_size']
        batch_dims = self._train_batch_dims
      else:
        batch_size = self._config_dict['eval_batch_size']
        batch_dims = self._eval_batch_dims
    return batch_size, batch_dims

  def call(self,
           pillars: tf.Tensor,
           indices: tf.Tensor,
           training: bool = None) -> tf.Tensor:
    """Forward pass of the featurizer."""
    # Add batch index to pillar indices.
    # (B, P, 1)
    batch_size, batch_dims = self._get_batch_size_and_dims(training)
    # (B, P, 3)
    batch_indices = tf.concat([batch_dims, indices], axis=-1)

    # Extract features from pillars.
    # (B, P, N, D)
    x = pillars
    # (B, P, N, C)
    for block in self._blocks:
      x = block(x)
    # (B, P, C)
    x = tf.reduce_max(x, axis=2, keepdims=False)

    # Scatter pillars back to form a BEV image.
    # (B, H, W, C)
    image = tf.scatter_nd(
        batch_indices,
        x,
        shape=[batch_size] + self._image_shape)
    self._output_specs = image.get_shape()
    return image

  def get_config(self) -> Mapping[str, Any]:
    return self._config_dict

  @classmethod
  def from_config(cls, config: Mapping[str, Any]) -> tf.keras.Model:
    return cls(**config)

  @property
  def output_specs(self) -> tf.TensorShape:
    return self._output_specs
