# Copyright 2020 The TensorFlow 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
#
#    https://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 the PointNet networks.

@inproceedings{qi2017pointnet,
  title={Pointnet: Deep learning on point sets
         for3d classification and segmentation},
  author={Qi, Charles R and Su, Hao and Mo, Kaichun and Guibas, Leonidas J},
  booktitle={Proceedings of the IEEE conference on computer vision and pattern
             recognition},
  pages={652--660},
  year={2017}}

NOTE: scheduling of batchnorm momentum currently not available in keras. However
experimentally, using the batch norm from Keras resulted in better test accuracy
(+1.5%) than the author's [custom batch norm
version](https://github.com/charlesq34/pointnet/blob/master/utils/tf_util.py)
even when coupled with batchnorm momentum decay. Further, note the author's
version is actually performing a "global normalization", as mentioned in the
[tf.nn.moments documentation]
(https://www.tensorflow.org/api_docs/python/tf/nn/moments).

This shorthand notation is used throughout this module:
  `B`: Number of elements in a batch.
  `N`: The number of points in the point set.
  `D`: Number of dimensions (e.g. 2 for 2D, 3 for 3D).
  `C`: The number of feature channels.
"""

from typing import Optional

import tensorflow as tf
from tensorflow_graphics.util import export_api


class PointNetConv2Layer(tf.keras.layers.Layer):
  """The 2D convolution layer used by the feature encoder in PointNet."""

  def __init__(self, channels, momentum):
    """Constructs a Conv2 layer.

    Note:
      Differently from the standard Keras Conv2 layer, the order of ops is:
      1. fully connected layer
      2. batch normalization layer
      3. ReLU activation unit

    Args:
      channels: the number of generated feature.
      momentum: the momentum of the batch normalization layer.
    """
    super(PointNetConv2Layer, self).__init__()
    self.channels = channels