#   Copyright 2021 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.
"""Label encoder layer for on-device personalized recommendation model."""
import math
from typing import Dict

import tensorflow as tf

from tensorflow_examples.lite.model_maker.third_party.recommendation.ml.configs import input_config_pb2


class LabelEncoder(tf.keras.layers.Layer):
  """Layer to encode label feature.

  Currently only id-based label encoder is supported. With this encoder, label
  embedding layer will be created based on input config. And label embedding
  will be generated by feeding input label feature to label embedding layer.
  """

  def __init__(self,
               input_config: input_config_pb2.InputConfig):
    """Initialize LabelEncoder layer based on input config.

    Args:
      input_config: The input config containing label feature configuation,
        based on which label encoding could be generated.
    """
    super(LabelEncoder, self).__init__()
    self._input_config = input_config
    self._label_feature = self._input_config.label_feature
    if self._label_feature.feature_type != input_config_pb2.FeatureType.INT:
      raise ValueError('Currently only INT type is supported for label.')
    if not self.label_name:
      raise ValueError('label_feature name is expected.')
    self._label_embedding_layer = tf.keras.layers.Embedding(
        self._label_feature.vocab_size,
        self._label_feature.embedding_dim,
        embeddings_initializer=tf.keras.initializers.truncated_normal(
            mean=0.0,
            stddev=1.0 / math.sqrt(self._label_feature.embedding_dim)),
        mask_zero=True,
        name=self.label_name + 'embedding_layer')

  @property
  def label_name(self):
    return self._label_feature.feature_name

  def encode(self, input_label: tf.Tensor) -> tf.Tensor:
    """Generates label encoding with input label tensor.

    Args:
      input_label: Input label tensor, the shape of tensor should
        [num_label_items, 1].

    Returns:
      Label encoding.
    """
    if isinstance(input_label, tf.SparseTensor):
      input_label = tf.sparse.to_dense(input_label)
    label_embedding = self._label_embedding_layer(input_label)
    if tf.keras.backend.ndim(label_embedding) == 3:
      label_embedding = tf.squeeze(label_embedding)
    return label_embedding

  def call(self, inputs: Dict[str, tf.Tensor]):
    """Generate label encoding with input tentor dictionary.

    Args:
      inputs: Dictionary mapping feature name to input tensors.

    Returns:
      Label encoding.
    """
    input_label = inputs[self.label_name]
    return self.encode(input_label)
