#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.

"""
This module defines a set of data processing transforms that can be used
to perform common data transformations on a dataset. These transforms are
implemented using the TensorFlow Transform (TFT) library. The transforms
in this module are intended to be used in conjunction with the
MLTransform class, which provides a convenient interface for
applying a sequence of data processing transforms to a dataset.

See the documentation for MLTransform for more details.

Note: The data processing transforms defined in this module don't
perform the transformation immediately. Instead, it returns a
configured operation object, which encapsulates the details of the
transformation. The actual computation takes place later in the Apache Beam
pipeline, after all transformations are set up and the pipeline is run.
"""

# pytype: skip-file

import logging
from typing import Any
from typing import Dict
from typing import Iterable
from typing import List
from typing import Optional
from typing import Tuple
from typing import Union

import tensorflow as tf
import tensorflow_transform as tft
from apache_beam.ml.transforms.base import BaseOperation
from tensorflow_transform import analyzers
from tensorflow_transform import common_types
from tensorflow_transform import tf_utils

__all__ = [
    'ComputeAndApplyVocabulary',
    'ScaleToZScore',
    'ScaleTo01',
    'ApplyBuckets',
    'Bucketize',
    'TFIDF',
    'TFTOperation',
    'ScaleByMinMax',
    'NGrams',
]

# Register the expected input types for each operation
# this will be used to determine schema for the tft.AnalyzeDataset
_EXPECTED_TYPES: Dict[str, Union[int, str, float]] = {}

_LOGGER = logging.getLogger(__name__)


def register_input_dtype(type):
  def wrapper(fn):
    _EXPECTED_TYPES[fn.__name__] = type
    return fn

  return wrapper


class TFTOperation(BaseOperation[common_types.TensorType,
                                 common_types.TensorType]):
  def __init__(self, columns: List[str]) -> None:
    """
    Base Operation class for TFT data processing transformations.
    Processing logic for the transformation is defined in the
    apply_transform() method. If you have a custom transformation that is not
    supported by the existing transforms, you can extend this class
    and implement the apply_transform() method.
    Args:
      columns: List of column names to apply the transformation.
    """
    super().__init__(columns)
    if not columns:
      raise RuntimeError(
          "Columns are not specified. Please specify the column for the "
          " op %s" % self.__class__.__name__)

  def get_artifacts(self, data: common_types.TensorType,
                    col_name: str) -> Dict[str, common_types.TensorType]:
    """
    Returns the artifacts generated by the operation.
    """
    return {}

  @tf.function
  def _split_string_with_delimiter(self, data, delimiter):
    """
    only applicable to string columns.
    """
    data = tf.sparse.to_dense(data)
    # this method acts differently compared to tf.strings.split
    # this will split the string based on multiple delimiters while
    # the latter will split the string based on a single delimiter.
    fn = lambda data: tf.compat.v1.string_split(
        data, delimiter, result_type='RaggedTensor')
    # tf.compat.v1.string_split works on a single string. Use tf.map_fn
    # to apply the function on each element of the input data.
    data = tf.map_fn(
        fn,
        data,
        fn_output_signature=tf.RaggedTensorSpec(
            tf.TensorShape([None, None]), tf.string))
    data = data.values.to_sparse()
    # the columns of the sparse tensor are suffixed with $indices, $values
    # related to sparse tensor. Create a new sparse tensor by extracting
    # the indices, values and dense_shape from the original sparse tensor
    # to preserve the original column name.
    data = tf.sparse.SparseTensor(
        indices=data.indices, values=data.values, dense_shape=data.dense_shape)
    # for list of string, batch dimensions becomes inverted after tf.map_fn,
    #  transpose the data to get the original shape.
    if tf.shape(data)[1] == 1:
      data = tf.sparse.transpose(data)
    return data


@register_input_dtype(str)
class ComputeAndApplyVocabulary(TFTOperation):
  def __init__(
      self,
      columns: List[str],
      split_string_by_delimiter: Optional[str] = None,
      *,
      default_value: Any = -1,
      top_k: Optional[int] = None,
      frequency_threshold: Optional[int] = None,
      num_oov_buckets: int = 0,
      vocab_filename: Optional[str] = None,
      name: Optional[str] = None):
    """
    This function computes the vocabulary for the given columns of incoming
    data. The transformation converts the input values to indices of the
    vocabulary.

    Args:
      columns: List of column names to apply the transformation.
      split_string_by_delimiter: (Optional) A string that specifies the
        delimiter to split strings.
      default_value: (Optional) The value to use for out-of-vocabulary values.
      top_k: (Optional) The number of most frequent tokens to keep.
      frequency_threshold: (Optional) Limit the generated vocabulary only to
        elements whose absolute frequency is >= to the supplied threshold.
        If set to None, the full vocabulary is generated.
      num_oov_buckets:  Any lookup of an out-of-vocabulary token will return a
        bucket ID based on its hash if `num_oov_buckets` is greater than zero.
        Otherwise it is assigned the `default_value`.
      vocab_filename: The file name for the vocabulary file. If not provided,
        the default name would be `compute_and_apply_vocab'
        NOTE in order to make your pipelines resilient to implementation
        details please set `vocab_filename` when you are using
        the vocab_filename on a downstream component.
    """
    super().__init__(columns)
    self._default_value = default_value
    self._top_k = top_k
    self._frequency_threshold = frequency_threshold
    self._num_oov_buckets = num_oov_buckets
    self._vocab_filename = vocab_filename if vocab_filename else (
        'compute_and_apply_vocab')
    self._name = name
    self.split_string_by_delimiter = split_string_by_delimiter

  def apply_transform(
      self, data: common_types.TensorType,
      output_column_name: str) -> Dict[str, common_types.TensorType]:

    if self.split_string_by_delimiter:
      data = self._split_string_with_delimiter(
          data, self.split_string_by_delimiter)

    return {
        output_column_name: tft.compute_and_apply_vocabulary(
            x=data,
            default_value=self._default_value,
            top_k=self._top_k,
            frequency_threshold=self._frequency_threshold,
            num_oov_buckets=self._num_oov_buckets,
            vocab_filename=self._vocab_filename,
            name=self._name)
    }


@register_input_dtype(float)
class ScaleToZScore(TFTOperation):
  def __init__(
      self,
      columns: List[str],
      *,
      elementwise: bool = False,
      name: Optional[str] = None):
    """
    This function performs a scaling transformation on the specified columns of
    the incoming data. It processes the input data such that it's normalized
    to have a mean of 0 and a variance of 1. The transformation achieves this
    by subtracting the mean from the input data and then dividing it by the
    square root of the variance.

    Args:
      columns: A list of column names to apply the transformation on.
      elementwise: If True, the transformation is applied elementwise.
        Otherwise, the transformation is applied on the entire column.
      name: A name for the operation (optional).

    scale_to_z_score also outputs additional artifacts. The artifacts are
    mean, which is the mean value in the column, and var, which is the
    variance in the column. The artifacts are stored in the column
    named with the suffix <original_col_name>_mean and <original_col_name>_var
    respectively.
    """
    super().__init__(columns)
    self.elementwise = elementwise
    self.name = name

  def apply_transform(
      self, data: common_types.TensorType,
      output_column_name: str) -> Dict[str, common_types.TensorType]:
    output_dict = {
        output_column_name: tft.scale_to_z_score(
            x=data, elementwise=self.elementwise, name=self.name)
    }
    return output_dict

  def get_artifacts(self, data: common_types.TensorType,
                    col_name: str) -> Dict[str, common_types.TensorType]:
    mean_var = tft.analyzers._mean_and_var(data)
    shape = [tf.shape(data)[0], 1]
    return {
        col_name + '_mean': tf.broadcast_to(mean_var[0], shape),
        col_name + '_var': tf.broadcast_to(mean_var[1], shape),
    }


@register_input_dtype(float)
class ScaleTo01(TFTOperation):
  def __init__(
      self,
      columns: List[str],
      elementwise: bool = False,
      name: Optional[str] = None):
    """
    This function applies a scaling transformation on the given columns
    of incoming data. The transformation scales the input values to the
    range [0, 1] by dividing each value by the maximum value in the
    column.

    Args:
      columns: A list of column names to apply the transformation on.
      elementwise: If True, the transformation is applied elementwise.
        Otherwise, the transformation is applied on the entire column.
      name: A name for the operation (optional).

    ScaleTo01 also outputs additional artifacts. The artifacts are
    max, which is the maximum value in the column, and min, which is the
    minimum value in the column. The artifacts are stored in the column
    named with the suffix <original_col_name>_min and <original_col_name>_max
    respectively.

    """
    super().__init__(columns)
    self.elementwise = elementwise
    self.name = name

  def get_artifacts(self, data: common_types.TensorType,
                    col_name: str) -> Dict[str, common_types.TensorType]:
    shape = [tf.shape(data)[0], 1]
    return {
        col_name + '_min': tf.broadcast_to(tft.min(data), shape),
        col_name + '_max': tf.broadcast_to(tft.max(data), shape)
    }

  def apply_transform(
      self, data: common_types.TensorType,
      output_column_name: str) -> Dict[str, common_types.TensorType]:
    output = tft.scale_to_0_1(
        x=data, elementwise=self.elementwise, name=self.name)

    output_dict = {output_column_name: output}
    return output_dict


@register_input_dtype(float)
class ApplyBuckets(TFTOperation):
  def __init__(
      self,
      columns: List[str],
      bucket_boundaries: Iterable[Union[int, float]],
      name: Optional[str] = None):
    """
    This functions is used to map the element to a positive index i for
    which bucket_boundaries[i-1] <= element < bucket_boundaries[i],
    if it exists. If input < bucket_boundaries[0], then element is
    mapped to 0. If element >= bucket_boundaries[-1], then element is
    mapped to len(bucket_boundaries). NaNs are mapped to
    len(bucket_boundaries).

    Args:
      columns: A list of column names to apply the transformation on.
      bucket_boundaries: A rank 2 Tensor or list representing the bucket
        boundaries sorted in ascending order.
      name: (Optional) A string that specifies the name of the operation.
    """
    super().__init__(columns)
    self.bucket_boundaries = [bucket_boundaries]
    self.name = name

  def apply_transform(
      self, data: common_types.TensorType,
      output_column_name: str) -> Dict[str, common_types.TensorType]:
    output = {
        output_column_name: tft.apply_buckets(
            x=data, bucket_boundaries=self.bucket_boundaries, name=self.name)
    }
    return output


@register_input_dtype(float)
class Bucketize(TFTOperation):
  def __init__(
      self,
      columns: List[str],
      num_buckets: int,
      *,
      epsilon: Optional[float] = None,
      elementwise: bool = False,
      name: Optional[str] = None):
    """
    This function applies a bucketizing transformation on the given columns
    of incoming data. The transformation splits the input data range into
    a set of consecutive bins/buckets, and converts the input values to
    bucket IDs (integers) where each ID corresponds to a particular bin.

    Args:
      columns: List of column names to apply the transformation.
      num_buckets: Number of buckets to be created.
      epsilon: (Optional) A float number that specifies the error tolerance
        when computing quantiles, so that we guarantee that any value x will
        have a quantile q such that x is in the interval
        [q - epsilon, q + epsilon] (or the symmetric interval for even
        num_buckets). Must be greater than 0.0.
      elementwise: (Optional) A boolean that specifies whether the quantiles
        should be computed on an element-wise basis. If False, the quantiles
        are computed globally.
      name: (Optional) A string that specifies the name of the operation.
    """
    super().__init__(columns)
    self.num_buckets = num_buckets
    self.epsilon = epsilon
    self.elementwise = elementwise
    self.name = name

  def get_artifacts(self, data: common_types.TensorType,
                    col_name: str) -> Dict[str, common_types.TensorType]:
    num_buckets = self.num_buckets
    epsilon = self.epsilon
    elementwise = self.elementwise

    if num_buckets < 1:
      raise ValueError('Invalid num_buckets %d' % num_buckets)

    if isinstance(data, (tf.SparseTensor, tf.RaggedTensor)) and elementwise:
      raise ValueError(
          'bucketize requires `x` to be dense if `elementwise=True`')

    x_values = tf_utils.get_values(data)

    if epsilon is None:
      # See explanation in args documentation for epsilon.
      epsilon = min(1.0 / num_buckets, 0.01)

    quantiles = analyzers.quantiles(
        x_values, num_buckets, epsilon, reduce_instance_dims=not elementwise)
    shape = [
        tf.shape(data)[0], num_buckets - 1 if num_buckets > 1 else num_buckets
    ]
    # These quantiles are used as the bucket boundaries in the later stages.
    # Should we change the prefix _quantiles to _bucket_boundaries?
    return {col_name + '_quantiles': tf.broadcast_to(quantiles, shape)}

  def apply_transform(
      self, data: common_types.TensorType,
      output_column_name: str) -> Dict[str, common_types.TensorType]:
    output = {
        output_column_name: tft.bucketize(
            x=data,
            num_buckets=self.num_buckets,
            epsilon=self.epsilon,
            elementwise=self.elementwise,
            name=self.name)
    }
    return output


@register_input_dtype(float)
class TFIDF(TFTOperation):
  def __init__(
      self,
      columns: List[str],
      vocab_size: Optional[int] = None,
      smooth: bool = True,
      name: Optional[str] = None,
  ):
    """
    This function applies a tf-idf transformation on the given columns
    of incoming data.

    TFIDF outputs two artifacts for each column: the vocabu index and
    the tfidf weight. The vocabu index is a mapping from the original
    vocabulary to the new vocabulary. The tfidf weight is a mapping
    from the original vocabulary to the tfidf score.

    Input passed to the TFIDF is not modified and used to calculate the
    required artifacts.

    Args:
      columns: List of column names to apply the transformation.
      vocab_size: (Optional) An integer that specifies the size of the
        vocabulary. Defaults to None.

        If vocab_size is None, then the size of the vocabulary is
        determined by `tft.get_num_buckets_for_transformed_feature`.
      smooth: (Optional) A boolean that specifies whether to apply
        smoothing to the tf-idf score. Defaults to True.
      name: (Optional) A string that specifies the name of the operation.
    """
    super().__init__(columns)
    self.vocab_size = vocab_size
    self.smooth = smooth
    self.name = name
    self.tfidf_weight = None

  def apply_transform(
      self, data: common_types.TensorType,
      output_column_name: str) -> common_types.TensorType:

    if self.vocab_size is None:
      try:
        _LOGGER.info(
            'vocab_size is not specified. Trying to infer vocab_size '
            'from the input data using '
            'tft.get_num_buckets_for_transformed_feature.')
        vocab_size = tft.get_num_buckets_for_transformed_feature(data)
      except RuntimeError:
        raise RuntimeError(
            'vocab_size is not specified. Tried to infer vocab_size from the '
            'input data using tft.get_num_buckets_for_transformed_feature, but '
            'failed. Please specify vocab_size explicitly.')
    else:
      vocab_size = self.vocab_size

    vocab_index, tfidf_weight = tft.tfidf(
      data,
      vocab_size,
      self.smooth,
      self.name
    )

    output = {
        output_column_name + '_vocab_index': vocab_index,
        output_column_name + '_tfidf_weight': tfidf_weight
    }
    return output


@register_input_dtype(float)
class ScaleByMinMax(TFTOperation):
  def __init__(
      self,
      columns: List[str],
      min_value: float = 0.0,
      max_value: float = 1.0,
      name: Optional[str] = None):
    """
    This function applies a scaling transformation on the given columns
    of incoming data. The transformation scales the input values to the
    range [min_value, max_value].

    Args:
      columns: A list of column names to apply the transformation on.
      min_value: The minimum value of the output range.
      max_value: The maximum value of the output range.
      name: A name for the operation (optional).
    """
    super().__init__(columns)
    self.min_value = min_value
    self.max_value = max_value
    self.name = name

    if self.max_value <= self.min_value:
      raise ValueError('max_value must be greater than min_value')

  def apply_transform(
      self, data: common_types.TensorType,
      output_column_name: str) -> common_types.TensorType:

    output = tft.scale_by_min_max(
        x=data, output_min=self.min_value, output_max=self.max_value)
    return {output_column_name: output}


@register_input_dtype(str)
class NGrams(TFTOperation):
  def __init__(
      self,
      columns: List[str],
      split_string_by_delimiter: Optional[str] = None,
      *,
      ngram_range: Tuple[int, int],
      ngrams_separator: str,
      name: Optional[str] = None):
    """
    An n-gram is a contiguous sequence of n items from a given sample of text
    or speech. This operation applies an n-gram transformation to
    specified columns of incoming data, splitting the input data into a
    set of consecutive n-grams.

    Args:
      columns: A list of column names to apply the transformation on.
      split_string_by_delimiter: (Optional) A string that specifies the
        delimiter to split the input strings before computing ngrams.
      ngram_range: A tuple of integers(inclusive) specifying the range of
        n-gram sizes.
      ngrams_separator: A string that will be inserted between each ngram.
      name: A name for the operation (optional).
    """
    super().__init__(columns)
    self.ngram_range = ngram_range
    self.ngrams_separator = ngrams_separator
    self.name = name
    self.split_string_by_delimiter = split_string_by_delimiter

  def apply_transform(
      self, data: common_types.TensorType,
      output_column_name: str) -> Dict[str, common_types.TensorType]:
    if self.split_string_by_delimiter:
      data = self._split_string_with_delimiter(
          data, self.split_string_by_delimiter)
    output = tft.ngrams(data, self.ngram_range, self.ngrams_separator)
    return {output_column_name: output}
