# Copyright 2021 Google LLC
#
# 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.


r"""Modules for monkey patching to observe tensors created with feature columns.

This file solely focuses on models created with tf.feature_columns. Generation
of tensors are hidden from the user in this setting. Thus, we monkey-patch
certain functions to observe what tensors are created while saving the model.

Hierarchy of feature columns can be seen below:

      Categorical Column                     Dense Column
             ^                <----_  _----->       ^
             |                      \/              |
             |              bucketized_column       |
categorical_column_with_identity           numeric_column
categorical_column_with_vocabulary_file    indicator_column
categorical_column_with_vocabulary_list    embedding_column
categorical_column_with_with_hash_bucket
crossed_column

Categorical columns are densified via indicator or embedding columns before
being fed to a DNN model; but they can be directly fed to linear models. Crossed
column acts as a categorical column even if any of its keys is a dense column.
Hence, we collect input tensors as sparse tensors generated by categorical
columns or dense tensors for numeric columns. We also collect densified
versions of those tensors from indicator and embedding columns for DNNs (Linear
models are also densified, but in a different way). For each input tensor(s), we
collect their encoded tensors from these densifications.

There are many edge and corner cases on how to build a model with feature
columns and how Tensorflow creates tensors for each case. For example,
DNNLinearCombinedClassifier creates two sets of all inputs and encoded tensors
for each feature. We keep track of all input tensors and potentially encoded
tensors with a class called FeatureTensors. For each base feature, we hold an
array of FeatureTensors objects.

feature tensors:
{<feature name 1>: [FeatureTensors(<input tensor>,
                                  [<embedding tensor>, <indicator tensor>])],
 <feature name 2>: [FeatureTensors(<input tensor 2>, [])],
 <feature name 3>: [FeatureTensors(<input tensor 3>, [])]}

Number of FeatureTensors object for each base column will be one most of the
time. In the rare case of wide deep classifier, there will be two sets of input
and encoded tensors, which leads to a list of FeatureTensors object.
"""
import contextlib
from typing import Any, Callable, Dict, Text, List, Set, Union, Optional
from six.moves import zip
import tensorflow.compat.v1 as tf

from tensorflow.python.feature_column import feature_column_v2 as fc2  
from tensorflow_estimator.python.estimator.canned import prediction_keys
from tensorflow_estimator.python.estimator.export import export_lib


_FEATURE_COLUMNS_TO_PATCH_DENSE = [
    fc2.DenseColumn, fc2.NumericColumn, fc2.EmbeddingColumn,
    fc2.IndicatorColumn, fc2.BucketizedColumn  # Bucketized is dense and sparse.
]

_FEATURE_COLUMNS_TO_PATCH_SPARSE = [
    fc2.CrossedColumn, fc2.BucketizedColumn,  # Bucketized is dense and sparse.
    fc2.WeightedCategoricalColumn
]

_FEATURE_COLUMNS_TO_PATCH_TRANSFORM = [
    fc2.VocabularyListCategoricalColumn, fc2.VocabularyFileCategoricalColumn,
    fc2.IdentityCategoricalColumn, fc2.HashedCategoricalColumn,
]


class FeatureTensors(object):
  """Groups of tensors representing a feature."""

  def __init__(
      self,
      input_tensor: Union[tf.Tensor, fc2.CategoricalColumn.IdWeightPair],
      encoded_tensors: List[tf.Tensor] = ()):  # pytype: disable=annotation-type-mismatch
    """Initialize a FeatureTensors object.

    Args:
      input_tensor: Input tensor representing this feature. Can be sparse.
      encoded_tensors: A set of tensors encoding input_tensor.
    """
    self._input_tensor = input_tensor
    self._encoded_tensors = encoded_tensors or []

  @property
  def input_tensor(self):
    return self._input_tensor

  @input_tensor.setter
  def input_tensor(
      self, tensor: Union[tf.Tensor, fc2.CategoricalColumn.IdWeightPair]):
    self._input_tensor = tensor

  @property
  def encoded_tensors(self):
    return self._encoded_tensors

  def __contains__(
      self,
      tensor: Union[tf.Tensor, fc2.CategoricalColumn.IdWeightPair]) -> bool:
    """Checks if tensors for the feature contains the given tensor."""
    if isinstance(tensor, tf.Tensor):
      return (self._input_tensor == tensor or
              any([tensor == encoded for encoded in self._encoded_tensors]))
    elif isinstance(tensor, fc2.CategoricalColumn.IdWeightPair):
      if not isinstance(self._input_tensor, fc2.CategoricalColumn.IdWeightPair):
        return False
      # Check equivalence of any tensor of the sparse tensor tuple.
      return tensor.id_tensor.indices == self._input_tensor.id_tensor.indices
    return False


class EstimatorMonkeyPatchHelper(object):
  """Monkey patches functions to observe important tensors in Estimators."""

  def __init__(self):
    """Create an EstimatorMonkeyPatchHelper instance.

    EstimatorMonkeyPatchHelper is used for observing tensors created during an
    estimator export. It provides a context manager with exporting_context. Once
    initialized, estimators can be exported within that context. Later, observed
    parameters can be fetched from feature_tensors_dict, output_tensors_dict,
    crossed_columns properties.
    """
    self._feature_tensors_dict = {}
    self._output_tensors_dict = {}
    self._crossed_columns = set()

  @property
  def feature_tensors_dict(self) -> Dict[Text, List[FeatureTensors]]:
    return self._feature_tensors_dict

  @property
  def output_tensors_dict(self) -> Dict[Text, tf.Tensor]:
    return self._output_tensors_dict

  @property
  def crossed_columns(self) -> Set[Text]:
    return self._crossed_columns

  def _make_observing_export_outputs_for_mode(
      self, old_fn: Callable[..., Any],
      observing_dict: Dict[Text, tf.Tensor],
      output_key: Optional[Text] = None) -> Callable[..., Any]:
    """Make an observing function that records output tensors.

    When a model is exported, export_outputs_for_mode() function is called. We
    observe output tensors in predictions parameter given output key. If no
    output key is provided, we resort to common keys which are 'predictions' for
    regression and 'probabilities' for classification.

    Args:
      old_fn: Original export_outputs_for_mode() function.
      observing_dict: Dictionary to record output names and corresponding tensor
        names.
      output_key: Key in prediction dictionary.