# Copyright (c) 2020 PaddlePaddle 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.
from __future__ import annotations

from typing import TYPE_CHECKING

import paddle
from paddle import _C_ops
from paddle.utils.decorator_utils import param_one_alias

from ...base.data_feeder import check_variable_and_dtype
from ...base.layer_helper import LayerHelper
from ...common_ops_import import Variable
from ...framework import in_dynamic_or_pir_mode

if TYPE_CHECKING:
    from paddle import Tensor

__all__ = []


@param_one_alias(["x", "input"])
def one_hot(
    x: Tensor,
    num_classes: int = -1,
    name: str | None = None,
) -> Tensor:
    """

    The operator converts each id in the input `x` to an one-hot vector with a
    `num_classes` length. The value in the vector dimension corresponding to the id
    is 1, and the value in the remaining dimension is 0.

    The shape of output Tensor is generated by appending `num_classes` dimension
    behind the last dimension of the `x` shape.

    .. code-block:: text

        Example 1:

        input:
            x.shape = [4]
            x.data = [1, 1, 3, 0]
            num_classes = 4

        output:
            Out.shape = [4, 4]
            Out.data = [[0., 1., 0., 0.],
                        [0., 1., 0., 0.],
                        [0., 0., 0., 1.],
                        [1., 0., 0., 0.]]

        Example 2:

        input:
            x.shape = [4]
            x.data = [1, 1, 5, 0]
            num_classes = 4

        output: Throw an exception for Illegal value
            The second dimension in X is 5, which is greater than num_classes,
            so it throws an exception.


    .. note::
        Alias Support: The parameter name ``input`` can be used as an alias for ``x``.
        For example, ``one_hot(input=tensor_x, ...)`` is equivalent to ``one_hot(x=tensor_x, ...)``.


    Args:
        x (Tensor): Tensor with shape :math:`[N_1, N_2, ..., N_k]` ,
            which contains at least one dimension. The data type is int32 or int64.
            Alias: ``input``.
        num_classes (int): An integer defining the `num_classes` of the one hot dimension. If input `x`
            is word id, `num_classes` is generally the dictionary size. Default value: -1.
        name (str|None, optional): For detailed information, please refer
           to :ref:`api_guide_Name`. Usually name is no need to set and
           None by default.

    Returns:
        Tensor, The one-hot representations of `x`. A Tensor with type float32.

    Examples:
        .. code-block:: pycon

            >>> import paddle
            >>> # Correspond to the first example above, where label.shape is 4 and one_hot_label.shape is [4, 4].
            >>> label = paddle.to_tensor([1, 1, 3, 0], dtype='int64')
            >>> print(label.shape)
            paddle.Size([4])
            >>> one_hot_label = paddle.nn.functional.one_hot(label, num_classes=4)
            >>> print(one_hot_label.shape)
            paddle.Size([4, 4])
            >>> print(one_hot_label)
            Tensor(shape=[4, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
                   [[0., 1., 0., 0.],
                    [0., 1., 0., 0.],
                    [0., 0., 0., 1.],
                    [1., 0., 0., 0.]])

    """
    if not isinstance(num_classes, paddle.pir.Value) and num_classes == -1:
        num_classes = x.max() + 1
    if in_dynamic_or_pir_mode():
        return _C_ops.one_hot(x, num_classes)
    else:
        check_variable_and_dtype(x, 'input', ['int32', 'int64'], 'one_hot_v2')
        helper = LayerHelper("one_hot_v2", **locals())

        one_hot_out = helper.create_variable_for_type_inference(dtype='float32')
        if not isinstance(num_classes, Variable):
            # user attribute
            inputs = {'X': x}
            attrs = {'depth': num_classes, 'allow_out_of_range': False}
        else:
            num_classes.stop_gradient = True
            inputs = {'X': x, 'depth_tensor': num_classes}
            attrs = {'allow_out_of_range': False}
        helper.append_op(
            type="one_hot_v2",
            inputs=inputs,
            attrs=attrs,
            outputs={'Out': one_hot_out},
            stop_gradient=True,
        )
        return one_hot_out


def embedding_renorm_(
    x: Tensor, weight: Tensor, max_norm: float, norm_type: float = 2.0
) -> Tensor:
    r"""
    Renorm the weight of embedding with respect to the provided :attr:`max_norm` and :attr:`norm_type` .

    Note:
        In the dynamic graph mode, the input weight will be updated in-place, and the return value will be the changed weight.

    Args:
        x (Tensor): A Tensor with type int32/int64, which contains the id information. The value of the input id should
            satisfy :math:`0<= id < weight.shape[0]` .
        weight (Tensor): The weight. A Tensor with shape of lookup table parameter. It should have two elements which
            indicates the size of the dictionary of embeddings and the size of each embedding vector respectively.
        max_norm (float): The maximum norm for each embedding vector.
        norm_type (float, optional): The p of the p-norm to compute for the max_norm option. Default: 2.0.

    Returns:
        Tensor, The updated weight. The data type is the same as :attr:`weight`.
    """
    with paddle.set_grad_enabled(False):
        unique_x = paddle.unique(x)
        selected_rows = paddle.index_select(weight, unique_x)
        norm = paddle.norm(selected_rows, p=norm_type, axis=1, keepdim=True)
        mask = norm > max_norm
        scale = max_norm / (norm + 1e-7)
        scale = paddle.where(mask, scale, paddle.ones_like(scale))
        scale = paddle.expand_as(scale, selected_rows)
        updated_rows = selected_rows * scale
        paddle.scatter_(weight, unique_x, updated_rows, overwrite=True)
        return weight


@param_one_alias(["x", "input"])
def embedding(
    x: Tensor,
    weight: Tensor,
    padding_idx: int | None = None,
    max_norm: float | None = None,
    norm_type: float = 2.0,
    sparse: bool = False,
    scale_grad_by_freq: bool = False,
    name: str | None = None,
) -> Tensor:
    r"""
    Used to lookup embeddings vector of ids provided by :attr:`x` .

    The shape of output Tensor is generated by appending the last dimension of the input Tensor shape
    with embedding size.

    Note:
        The id in :attr:`x` must satisfy :math:`0 <= id < weight.shape[0]` ,
        otherwise the program will throw an exception and exit.

    .. code-block:: text

            x is a Tensor.
                padding_idx = -1
                x.data = [[1, 3], [2, 4], [4, 127]]
                x.shape = [3, 2]
                weight.shape = [128, 16]
            output is a Tensor:
                out.shape = [3, 2, 16]
                out.data = [[[0.129435295, 0.244512452, ..., 0.436322452],
                            [0.345421456, 0.524563927, ..., 0.144534654]],
                            [[0.345249859, 0.124939536, ..., 0.194353745],
                            [0.945345345, 0.435394634, ..., 0.435345365]],
                            [[0.945345345, 0.435394634, ..., 0.435345365],
                            [0.0,         0.0,         ..., 0.0        ]]]  # padding data

            The input padding_idx is less than 0, it is automatically converted to padding_idx = -1 + 128 = 127
            It will pad all-zero data when id is 127.

    .. note::
        Alias Support: The parameter name ``input`` can be used as an alias for ``x``.
        For example, ``embedding(input=tensor_x, ...)`` is equivalent to ``embedding(x=tensor_x, ...)``.

    Args:
        x (Tensor): A Tensor with type int32/int64, which contains the id information. The value of the input id should
            satisfy :math:`0 <= id < weight.shape[0]` .
            Alias: ``input``.
        weight (Tensor): The weight. A Tensor with shape of lookup table parameter. It should have two elements which
            indicates the size of the dictionary of embeddings and the size of each embedding vector respectively.
        sparse (bool, optional): The flag indicating whether to use sparse update. This parameter only
            affects the performance of the backwards gradient update. It is recommended to set
            True because sparse update is faster. But some optimizers does not support sparse update,
            such as :ref:`api_paddle_optimizer_adadelta_Adadelta` , :ref:`api_paddle_optimizer_adamax_Adamax` , :ref:`api_paddle_optimizer_lamb_Lamb`.
            In these cases, sparse must be False. Default: False.
        padding_idx (int|None, optional): padding_idx needs to be in the interval [-weight.shape[0], weight.shape[0]).
            If :math:`padding\_idx < 0`, the :math:`padding\_idx` will automatically be converted
            to :math:`weight.shape[0] + padding\_idx` . It will output all-zero padding data whenever lookup
            encounters :math:`padding\_idx` in id. And the padding data will not be updated while training.
            If set None, it makes no effect to output. Default: None.
        max_norm (float, optional): If provided, will renormalize the embedding vectors to have a norm larger than
            :attr:`max\_norm` . It will inplace update the input embedding weight in dynamic graph mode. Default: None.
        norm_type (float, optional): The p of the p-norm to compute for the max_norm option. Default: 2.0.
        scale_grad_by_freq (bool, optional): Indicating whether to scale the gradients by the inverse frequency of the
            word ids in input `x`. Default: False.
        name (str|None, optional): For detailed information, please refer
           to :ref:`api_guide_Name`. Usually name is no need to set and
           None by default.

    Returns:
        Tensor, Embedding Tensor mapped by x. The data type is the same as :attr:`weight`.

    Examples:

        .. code-block:: pycon

            >>> import paddle
            >>> import paddle.nn as nn

            >>> x0 = paddle.arange(3, 6).reshape((3, 1)).astype(paddle.int64)
            >>> w0 = paddle.full(shape=(10, 3), fill_value=2).astype(paddle.float32)

            >>> x = paddle.to_tensor(x0, stop_gradient=False)
            >>> print(x.numpy())
            [[3]
             [4]
             [5]]
            >>> print(x.shape)
            paddle.Size([3, 1])

            >>> w = paddle.to_tensor(w0, stop_gradient=False)
            >>> print(w.numpy())
            [[2. 2. 2.]
             [2. 2. 2.]
             [2. 2. 2.]
             [2. 2. 2.]
             [2. 2. 2.]
             [2. 2. 2.]
             [2. 2. 2.]
             [2. 2. 2.]
             [2. 2. 2.]
             [2. 2. 2.]]
            >>> print(w.shape)
            paddle.Size([10, 3])

            >>> emb = nn.functional.embedding(x=x, weight=w, sparse=True, name="embedding")
            >>> print(emb.numpy())
            [[[2. 2. 2.]]
             [[2. 2. 2.]]
             [[2. 2. 2.]]]
            >>> print(emb.shape)
            paddle.Size([3, 1, 3])

    """
    padding_idx = (
        -1
        if padding_idx is None
        else (
            padding_idx if padding_idx >= 0 else (weight.shape[0] + padding_idx)
        )
    )

    if weight.shape[0] != 0 and (
        padding_idx >= weight.shape[0] or padding_idx < -weight.shape[0]
    ):
        raise ValueError(
            f"padding_idx must be within [-{weight.shape[0]}, {weight.shape[0]})"
        )

    if max_norm and weight.size != 0:
        weight = embedding_renorm_(
            x, weight, max_norm=max_norm, norm_type=norm_type
        )

    if scale_grad_by_freq:
        if sparse:
            raise AttributeError(
                "scale_grad_by_freq = True is not supported with sparse update."
            )
        if in_dynamic_or_pir_mode():
            return _C_ops.embedding_with_scaled_gradient(x, weight, padding_idx)
        else:
            helper = LayerHelper('embedding_with_scaled_gradient', **locals())
            dtype = helper.input_dtype(input_param_name='weight')

            check_variable_and_dtype(
                x,
                'input',
                ['uint8', 'int8', 'int16', 'int32', 'int64'],
                'embedding_with_scaled_gradient',
            )
            tmp = helper.create_variable_for_type_inference(dtype)

            helper.append_op(
                type='embedding_with_scaled_gradient',
                inputs={'x': x, 'weight': weight},
                outputs={'out': tmp},
                attrs={'padding_idx': padding_idx},
            )
            return tmp
    else:
        if in_dynamic_or_pir_mode():
            return _C_ops.embedding(x, weight, padding_idx, sparse)
        else:
            helper = LayerHelper('embedding', **locals())
            dtype = helper.input_dtype(input_param_name='weight')

            check_variable_and_dtype(
                x,
                'input',
                ['uint8', 'int8', 'int16', 'int32', 'int64'],
                'embedding',
            )

            is_distributed = False
            remote_prefetch = sparse and (not is_distributed)

            tmp = helper.create_variable_for_type_inference(dtype)

            helper.append_op(
                type='lookup_table_v2',
                inputs={'Ids': x, 'W': weight},
                outputs={'Out': tmp},
                attrs={
                    'is_sparse': sparse,
                    'is_distributed': is_distributed,
                    'remote_prefetch': remote_prefetch,
                    'padding_idx': padding_idx,
                },
            )
            return tmp