# Copyright 2015 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.
# ==============================================================================
"""Contains the TFOpLambda layer."""
import tensorflow.compat.v2 as tf

from keras.src import backend
from keras.src.engine import keras_tensor
from keras.src.engine.base_layer import Layer

# isort: off
from tensorflow.python.platform import tf_logging
from tensorflow.python.util.tf_export import (
    get_canonical_name_for_symbol,
)
from tensorflow.python.util.tf_export import (
    get_symbol_from_name,
)


class ClassMethod(Layer):
    """Wraps a TF API Class's class method  in a `Layer` object.

    It is inserted by the Functional API construction whenever users call
    a supported TF Class's class method on KerasTensors.

    This is useful in the case where users do something like:
    x = keras.Input(...)
    y = keras.Input(...)
    out = tf.RaggedTensor.from_row_splits(x, y)
    """

    @tf.__internal__.tracking.no_automatic_dependency_tracking
    def __init__(self, cls_ref, method_name, **kwargs):
        self.cls_ref = cls_ref
        self.method_name = method_name
        self.cls_symbol = get_canonical_name_for_symbol(
            self.cls_ref, add_prefix_to_v1_names=True
        ) or get_canonical_name_for_symbol(
            self.cls_ref, api_name="keras", add_prefix_to_v1_names=True
        )
        if "name" not in kwargs:
            kwargs["name"] = backend.unique_object_name(
                "tf." + self.cls_symbol + "." + self.method_name,
                zero_based=True,
                avoid_observed_names=True,
            )
        kwargs["autocast"] = False

        # Do not individually trace op layers in the SavedModel.
        self._must_restore_from_config = True

        super().__init__(**kwargs)

        # Preserve all argument data structures when saving/loading a config
        # (e.g., don't unnest lists that contain one element)
        self._preserve_input_structure_in_config = True

        self._call_spec.expects_training_arg = False
        self._call_spec.expects_mask_arg = False

    def call(self, args, kwargs):
        return getattr(self.cls_ref, self.method_name)(*args, **kwargs)

    def get_config(self):
        if not self.cls_symbol:
            raise ValueError(
                "This Keras class method conversion tried to convert "
                f"a method belonging to class {self.cls_symbol}, a class "
                "that is not publicly exposed in the TensorFlow API. "
                "To ensure cross-version compatibility of Keras models "
                "that use op layers, only op layers produced from "
                "public TensorFlow API symbols can be serialized."
            )

        config = {
            "cls_symbol": self.cls_symbol,
            "method_name": self.method_name,
        }
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    @classmethod
    def from_config(cls, config, custom_objects=None):
        config = config.copy()
        symbol_name = config.pop("cls_symbol")
        cls_ref = get_symbol_from_name(symbol_name)
        if not cls_ref:
            raise ValueError(
                f"TensorFlow symbol `{symbol_name}` could not be found."
            )

        config["cls_ref"] = cls_ref

        return cls(**config)


class KerasOpDispatcher(tf.__internal__.dispatch.GlobalOpDispatcher):
    """A global dispatcher that allows building a functional model with TF
    Ops."""

    def handle(self, op, args, kwargs):
        """Handle the specified operation with the specified arguments."""
        if any(
            isinstance(x, keras_tensor.KerasTensor)
            for x in tf.nest.flatten([args, kwargs])
        ):
            return TFOpLambda(op)(*args, **kwargs)
        else:
            return self.NOT_SUPPORTED


KerasOpDispatcher().register()


class InstanceProperty(Layer):
    """Wraps an instance property access (e.g.

    `x.foo`) in a Keras Layer.

    This layer takes an attribute name `attr_name` in the constructor and,
    when called on input tensor `obj` returns `obj.attr_name`.

    KerasTensors specialized for specific extension types use it to
    represent instance property accesses on the represented object in the
    case where the property needs to be dynamically accessed as opposed to
    being statically computed from the typespec, e.g.

    x = keras.Input(..., ragged=True)
    out = x.flat_values
    """

    @tf.__internal__.tracking.no_automatic_dependency_tracking
    def __init__(self, attr_name, **kwargs):
        self.attr_name = attr_name

        if "name" not in kwargs:
            kwargs["name"] = backend.unique_object_name(
                "input." + self.attr_name,
                zero_based=True,
                avoid_observed_names=True,
            )
        kwargs["autocast"] = False

        # Do not individually trace op layers in the SavedModel.
        self._must_restore_from_config = True

        super().__init__(**kwargs)

        # Preserve all argument data structures when saving/loading a config
        # (e.g., don't unnest lists that contain one element)
        self._preserve_input_structure_in_config = True

    def call(self, obj):
        return getattr(obj, self.attr_name)

    def get_config(self):
        config = {"attr_name": self.attr_name}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    @classmethod
    def from_config(cls, config, custom_objects=None):
        return cls(**config)


class InstanceMethod(InstanceProperty):
    """Wraps an instance method access (e.g. `x.foo(arg)` in a Keras Layer.

    This layer takes an attribute name `attr_name` in the constructor and,
    when called on input tensor `obj` with additional arguments `args` and
    `kwargs` returns `obj.attr_name(*args, **kwargs)`.

    KerasTensors specialized for specific extension types use it to
    represent dynamic instance method calls on the represented object, e.g.

    x = keras.Input(..., ragged=True)
    new_values = keras.Input(...)
    out = x.with_values(new_values)
    """

    def call(self, obj, args, kwargs):
        method = getattr(obj, self.attr_name)
        return method(*args, **kwargs)


class TFOpLambda(Layer):
    """Wraps TF API symbols in a `Layer` object.

    It is inserted by the Functional API construction whenever users call
    a supported TF symbol on KerasTensors.

    Like Lambda layers, this layer tries to raise warnings when it detects users
    explicitly use variables in the call. (To let them know
    that the layer will not capture the variables).

    This is useful in the case where users do something like:
    x = keras.Input(...)
    y = tf.Variable(...)
    out = x * tf_variable
    """

    @tf.__internal__.tracking.no_automatic_dependency_tracking
    def __init__(self, function, **kwargs):
        self.function = function
        self.symbol = get_canonical_name_for_symbol(
            self.function, add_prefix_to_v1_names=True
        ) or get_canonical_name_for_symbol(
            self.function, api_name="keras", add_prefix_to_v1_names=True
        )
        if "name" not in kwargs:
            # Generate a name.
            # TFOpLambda layers avoid already-observed names,
            # because users cannot easily control the generated names.
            # Without this avoidance, users would be more likely to run
            # into unavoidable duplicate layer name collisions.
            # (For standard layers users could just set `name` when creating the
            # layer to work around a collision, but they can't do that for
            # auto-generated layers)
            if self.symbol:
                name = "tf." + self.symbol
            else:
                name = self.function.__name__
            kwargs["name"] = backend.unique_object_name(
                name, zero_based=True, avoid_observed_names=True
            )
        kwargs["autocast"] = False

        # Decorate the function to produce this layer's call method
        def _call_wrapper(*args, **kwargs):
            return self._call_wrapper(*args, **kwargs)

        self.call = tf.__internal__.decorator.make_decorator(
            function, _call_wrapper
        )

        # Do not individually trace op layers in the SavedModel.
        self._must_restore_from_config = True

        super().__init__(**kwargs)

        # Preserve all argument data structures when saving/loading a config
        # (e.g., don't unnest lists that contain one element)
        self._preserve_input_structure_in_config = True

        # Warning on every invocation will be quite irksome in Eager mode.
        self._already_warned = False

        self._call_spec.expects_training_arg = False
        self._call_spec.expects_mask_arg = False

    def _call_wrapper(self, *args, **kwargs):
        created_variables = []

        def _variable_creator(next_creator, **creator_kwargs):
            var = next_creator(**creator_kwargs)
            created_variables.append(var)
            return var

        with tf.GradientTape(
            watch_accessed_variables=True
        ) as tape, tf.variable_creator_scope(_variable_creator):
            # We explicitly drop `name` arguments here,
            # to guard against the case where an op explicitly has a
            # `name` passed (which is susceptible to producing
            # multiple ops w/ the same name when the layer is reused)
            kwargs.pop("name", None)
            result = self.function(*args, **kwargs)
        self._check_variables(created_variables, tape.watched_variables())
        return result

    def _check_variables(self, created_variables, accessed_variables):
        if not created_variables and not accessed_variables:
            # In the common case that a Lambda layer does not touch a Variable,
            # we don't want to incur the runtime cost of assembling any state
            # used for checking only to immediately discard it.
            return

        tracked_weights = set(v.ref() for v in self.weights)
        untracked_new_vars = [
            v for v in created_variables if v.ref() not in tracked_weights
        ]
        if untracked_new_vars:
            variable_str = "\n".join(f"  {i}" for i in untracked_new_vars)
            raise ValueError(
                "The following Variables were created within a Lambda layer "
                f"({self.name}) but are not tracked by said layer: "
                f"{variable_str}\n"
                "The layer cannot safely ensure proper Variable reuse "
                "across multiple calls, and consequently this behavior "
                "is disallowed for safety reasons. Lambda layers are "
                "not well suited for stateful computation; instead, "
                "writing a subclassed Layer is the recommend "
                "way to define layers with Variables."
            )

        untracked_used_vars = [
            v for v in accessed_variables if v.ref() not in tracked_weights
        ]
        if untracked_used_vars and not self._already_warned:
            variable_str = "\n".join(f"  {i}" for i in untracked_used_vars)
            self._warn(
                "The following Variables were used in a Lambda layer's call "
                f"({self.name}), but are not present in its tracked objects: "
                f"{variable_str}. This is a strong indication that the Lambda "
                "layer should be rewritten as a subclassed Layer."
            )
            self._already_warned = True

    def _warn(self, msg):
        # This method will be overridden in a unit test to raise an error,
        # because self.assertWarns is not universally implemented.
        return tf_logging.warning(msg)

    def get_config(self):
        if not self.symbol:
            raise ValueError(
                f"This Keras op layer was generated from {self.function}, a "
                "method that is not publicly exposed in the TensorFlow API. "
                "This may have happened if the method was explicitly "
                "decorated to add dispatching support, and it was used "
                "during Functional model construction. "
                "To ensure cross-version compatibility of Keras models "
                "that use op layers, only op layers produced from "
                "public TensorFlow API symbols can be serialized."
            )
        config = {"function": self.symbol}

        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    @classmethod
    def from_config(cls, config, custom_objects=None):
        config = config.copy()
        symbol_name = config["function"]
        function = get_symbol_from_name(symbol_name)
        if not function:
            raise ValueError(f"TF symbol `{symbol_name}` could not be found.")

        config["function"] = function

        return cls(**config)


def _delegate_property(keras_tensor_cls, property_name):
    """Register property on a KerasTensor class.

    Calling this multiple times with the same arguments should be a no-op.

    This method exposes a property on the KerasTensor class that will use an
    `InstanceProperty` layer to access the property on the represented
    intermediate values in the model.

    Args:
      keras_tensor_cls: The KerasTensor subclass that should expose the
        property.
      property_name: The name of the property to expose and delegate to the
        represented (Composite)Tensor.
    """
    # We use a lambda because we can't create a Keras layer at import time
    # due to dynamic layer class versioning.
    property_access = property(
        lambda self: InstanceProperty(property_name)(self)
    )
    setattr(keras_tensor_cls, property_name, property_access)


def _delegate_method(keras_tensor_cls, method_name):
    """Register method on a KerasTensor class.

    Calling this function times with the same arguments should be a no-op.

    This method exposes an instance method on the KerasTensor class that will
    use an `InstanceMethod` layer to run the desired method on the represented
    intermediate values in the model.

    Args:
      keras_tensor_cls: The KerasTensor subclass that should expose the
        property.
      method_name: The name of the method to expose and delegate to the
        represented (Composite)Tensor.
    """

    def delegate(self, *args, **kwargs):
        return InstanceMethod(method_name)(self, args, kwargs)

    setattr(keras_tensor_cls, method_name, delegate)


# We do not support the `uniform_row_length` property because it
# returns either `None` or an int tensor, and code that relies on it tends
# to check `is None` directly. Delegating it here would always return a
# `KerasTensor`, regardless of what can be statically inferred. This would
# never equal `None`, breaking code that expects it to be partially-static
# in unpredictable ways.
for ragged_property in [
    "values",
    "flat_values",
    "row_splits",
    "nested_row_splits",
]:
    _delegate_property(keras_tensor.RaggedKerasTensor, ragged_property)

for ragged_method_name in [
    "value_rowids",
    "nested_value_rowids",
    "nrows",
    "row_starts",
    "row_limits",
    "row_lengths",
    "nested_row_lengths",
    "bounding_shape",
    "with_values",
    "with_flat_values",
    "with_row_splits_dtype",
    "merge_dims",
    "to_tensor",
    "to_sparse",
]:
    _delegate_method(keras_tensor.RaggedKerasTensor, ragged_method_name)

for sparse_property in [
    "indices",
    "values",
    "dense_shape",
]:
    _delegate_property(keras_tensor.SparseKerasTensor, sparse_property)

for sparse_method in [
    "with_values",
]:
    _delegate_method(keras_tensor.SparseKerasTensor, sparse_method)


class TFClassMethodDispatcher(tf.__internal__.dispatch.OpDispatcher):
    """A class method dispatcher that allows building a functional model with TF
    class methods."""

    def __init__(self, cls, method_name):
        self.cls = cls
        self.method_name = method_name

    def handle(self, args, kwargs):
        """Handle the specified operation with the specified arguments."""
        if any(
            isinstance(x, keras_tensor.KerasTensor)
            for x in tf.nest.flatten([args, kwargs])
        ):
            return ClassMethod(self.cls, self.method_name)(args[1:], kwargs)
        else:
            return self.NOT_SUPPORTED


for ragged_class_method in [
    "from_value_rowids",
    "from_row_splits",
    "from_row_lengths",
    "from_row_starts",
    "from_row_limits",
    "from_uniform_row_length",
    "from_nested_value_rowids",
    "from_nested_row_splits",
    "from_nested_row_lengths",
    "from_tensor",
    "from_sparse",
]:
    TFClassMethodDispatcher(tf.RaggedTensor, ragged_class_method).register(
        getattr(tf.RaggedTensor, ragged_class_method)
    )


class SlicingOpLambda(TFOpLambda):
    """Wraps TF API symbols in a `Layer` object.

    It is inserted by the Functional API construction whenever users call
    a supported TF symbol on KerasTensors.

    Like Lambda layers, this layer tries to raise warnings when it detects users
    explicitly use variables in the call. (To let them know
    that the layer will not capture the variables).

    This is useful in the case where users do something like:
    x = keras.Input(...)
    y = tf.Variable(...)
    out = x * tf_variable
    """

    @tf.__internal__.tracking.no_automatic_dependency_tracking
    def __init__(self, function, **kwargs):
        super().__init__(function, **kwargs)

        original_call = self.call

        # Decorate the function to produce this layer's call method
        def _call_wrapper(*args, **kwargs):
            # Turn any slice dicts in the args back into `slice` objects.
            # This conversion cannot use nest.flatten/map_structure,
            # because dicts are flattened by nest while slices aren't.
            # So, map_structure would only see the individual elements in the
            # dict.
            # This can't use map_structure_up_to either because the
            # 'shallowness' of the shallow tree would have to vary depending on
            # if only one dim or multiple are being sliced.
            new_args = []
            for arg in args:
                arg = _dict_to_slice(arg)
                if isinstance(arg, (list, tuple)):
                    new_arg = []
                    for sub_arg in arg:
                        new_arg.append(_dict_to_slice(sub_arg))
                    arg = new_arg
                new_args.append(arg)

            # Handle the kwargs too.
            new_kwargs = {}
            for key, value in kwargs.items():
                value = _dict_to_slice(value)
                if isinstance(value, (list, tuple)):
                    new_value = []
                    for v in value:
                        new_value.append(_dict_to_slice(v))
                    value = new_value
                new_kwargs[key] = value

            return original_call(*new_args, **new_kwargs)

        self.call = tf.__internal__.decorator.make_decorator(
            original_call, _call_wrapper
        )


def _slice_to_dict(x):
    if isinstance(x, slice):
        return {"start": x.start, "stop": x.stop, "step": x.step}
    return x


def _dict_to_slice(x):
    if isinstance(x, dict):
        return slice(x["start"], x["stop"], x["step"])
    return x


class TFSlicingOpDispatcher(tf.__internal__.dispatch.OpDispatcher):
    """A global dispatcher that allows building a functional model with TF
    Ops."""

    def __init__(self, op):
        self.op = op

    def handle(self, args, kwargs):
        """Handle the specified operation with the specified arguments."""
        args = tf.nest.map_structure(_slice_to_dict, args)
        kwargs = tf.nest.map_structure(_slice_to_dict, kwargs)
        if any(
            isinstance(x, keras_tensor.KerasTensor)
            for x in tf.nest.flatten([args, kwargs])
        ):
            return SlicingOpLambda(self.op)(*args, **kwargs)
        else:
            return self.NOT_SUPPORTED


for slicing_op in [
    tf.__operators__.getitem,
    tf.compat.v1.boolean_mask,
    tf.boolean_mask,
    tf.__operators__.ragged_getitem,
]:
    TFSlicingOpDispatcher(slicing_op).register(slicing_op)

