# Copyright (c) 2025 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

if TYPE_CHECKING:
    from collections.abc import Callable
    from typing import TypeVar

    from typing_extensions import ParamSpec

    _InputT = ParamSpec("_InputT")
    _RetT = TypeVar("_RetT")

import math
import warnings

import numpy as np

import paddle

from ..base.framework import in_dygraph_mode, in_pir_mode
from .initializer.constant import Constant
from .initializer.dirac import Dirac
from .initializer.initializer import calculate_gain  # noqa: F401
from .initializer.kaiming import KaimingNormal, KaimingUniform
from .initializer.normal import Normal, TruncatedNormal
from .initializer.orthogonal import Orthogonal
from .initializer.uniform import Uniform
from .initializer.xavier import XavierNormal, XavierUniform


def _calculate_fan_in_and_fan_out(var: paddle.Tensor) -> tuple[int, int]:
    """Compute the fan_in and the fan_out for layers

    This method computes the fan_in and the fan_out
    for neural network layers, if not specified. It is
    not possible to perfectly estimate fan_in and fan_out.
    This method will estimate it correctly for matrix multiply and
    convolutions.

    Args:
        var: variable for which fan_in and fan_out have to be computed.

    Returns:
        tuple of two integers (fan_in, fan_out).
    """
    shape = var.shape
    if not shape or len(shape) == 0:
        fan_in = fan_out = 1
    elif len(shape) == 1:
        fan_in = fan_out = shape[0]
    elif len(shape) == 2:
        # This is the case for simple matrix multiply
        fan_in = shape[0]
        fan_out = shape[1]
    else:
        # Assume this to be a convolutional kernel
        # In PaddlePaddle, the shape of the kernel is like:
        # [num_filters, num_filter_channels, ...] where the remaining
        # dimensions are the filter_size
        receptive_field_size = np.prod(shape[2:])
        fan_in = int(shape[1] * receptive_field_size)
        fan_out = int(shape[0] * receptive_field_size)
    return (fan_in, fan_out)


def kaiming_uniform_(
    tensor: paddle.Tensor,
    a: float = 0,
    mode: str = "fan_in",
    nonlinearity: str = "leaky_relu",
) -> paddle.Tensor:
    """Modify tensor inplace using Kaiming uniform method.

    Args:
        tensor (Tensor):  Paddle Tensor.
        a (float, optional): The negative slope of the rectifier used after this layer.
            Defaults to 0.
        mode (str, optional): Mode to compute the fan. Choose from ["fan_in", "fan_out"].
            When set to 'fan_in', the fan_in parameter is used for initialization.
            When set to 'fan_out', the out_features of trainable Tensor will be used.
            Default is 'fan_in'.
        nonlinearity (str, optional): Nonlinearity method name. Defaults to "leaky_relu".

    Returns:
        Tensor: Initialized tensor.
    """
    init = KaimingUniform(
        negative_slope=a, nonlinearity=nonlinearity, mode=mode
    )

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def kaiming_normal_(
    tensor: paddle.Tensor,
    a: float = 0,
    mode: str = "fan_in",
    nonlinearity: str = "leaky_relu",
) -> paddle.Tensor:
    """Modify tensor inplace using Kaiming normal method.

    Args:
        tensor (Tensor):  Paddle Tensor.
        a (float, optional): The negative slope of the rectifier used after this layer.
            Defaults to 0.
        mode (str, optional): Mode to compute the fan. Choose from ["fan_in", "fan_out"].
            When set to 'fan_in', the fan_in parameter is used for initialization.
            When set to 'fan_out', the out_features of trainable Tensor will be used.
            Default is 'fan_in'.
        nonlinearity (str, optional): Nonlinearity method name. Defaults to "leaky_relu".

    Returns:
        Tensor: Initialized tensor.
    """
    init = KaimingNormal(negative_slope=a, nonlinearity=nonlinearity, mode=mode)

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def xavier_uniform_(
    tensor: paddle.Tensor,
    gain: float = 1.0,
    fan_in: float | None = None,
    fan_out: float | None = None,
) -> paddle.Tensor:
    """Modify tensor inplace using Xavier uniform method.

    Args:
        tensor (Tensor):  Paddle Tensor.
        gain (float, optional): Scaling Tensor. Default is 1.0.
        fan_in (float|None, optional): fan_in for Xavier initialization, which is
                inferred from the Tensor. Default is None.
        fan_out (float|None, optional): fan_out for Xavier initialization, which is
                 inferred from the Tensor. Default is None.

    Returns:
        Tensor: Initialized tensor.
    """
    init = XavierUniform(
        gain=gain,
        fan_in=fan_in,
        fan_out=fan_out,
    )

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def xavier_normal_(
    tensor: paddle.Tensor,
    gain: float = 1.0,
    fan_in: float | None = None,
    fan_out: float | None = None,
) -> paddle.Tensor:
    """Modify tensor inplace using Xavier normal method.

    Args:
        tensor (Tensor):  Paddle Tensor.
        gain (float, optional): Scaling Tensor. Default is 1.0.
        fan_in (float|None, optional): fan_in for Xavier initialization, which is
                inferred from the Tensor. Default is None.
        fan_out (float|None, optional): fan_out for Xavier initialization, which is
                 inferred from the Tensor. Default is None.

    Returns:
        Tensor: Initialized tensor.
    """
    init = XavierNormal(
        gain=gain,
        fan_in=fan_in,
        fan_out=fan_out,
    )

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def uniform_(
    tensor: paddle.Tensor,
    a: float = 0.0,
    b: float = 1.0,
) -> paddle.Tensor:
    """Modify tensor inplace using uniform method.

    Args:
        tensor (Tensor):  Paddle Tensor.
        low (float, optional): Lower boundary of the uniform distribution. Default is :math:`-1.0`.
        high (float, optional): Upper boundary of the uniform distribution. Default is :math:`1.0`.

    Returns:
        Tensor: Initialized tensor.
    """
    init = Uniform(low=a, high=b)

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def normal_(
    tensor: paddle.Tensor,
    mean: float = 0.0,
    std: float = 1.0,
) -> paddle.Tensor:
    """Modify tensor inplace using normal method.

    Args:
        tensor (Tensor):  Paddle Tensor.
        mean (float|complex, optional): mean of the normal distribution. Default is 0.0.
        std (float, optional): standard deviation of the normal distribution. Default is 1.0.

    Returns:
        Tensor: Initialized tensor.
    """
    init = Normal(mean=mean, std=std)

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def trunc_normal_(
    tensor: paddle.Tensor,
    mean: float = 0.0,
    std: float = 1.0,
    a: float = -2.0,
    b: float = 2.0,
) -> paddle.Tensor:
    """Modify tensor inplace using truncated normal method.

    Args:
        tensor (Tensor):  Paddle Tensor.
        mean (float|complex, optional): mean of the normal distribution. Default is 0.0.
        std (float, optional): standard deviation of the normal distribution. Default is 1.0.
        a (float, optional): The minimum cutoff value. Default is -2.0.
        b (float, optional): The maximum cutoff value. Default is 2.0.

    Returns:
        Tensor: Initialized tensor.
    """
    init = TruncatedNormal(mean=mean, std=std, a=a, b=b)

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def constant_(
    tensor: paddle.Tensor,
    val: float,
) -> paddle.Tensor:
    """Modify tensor inplace using constant method.

    Args:
        tensor (Tensor):  Paddle Tensor.
        value (float32|float64, optional): constant value to initialize the parameter.

    Returns:
        Tensor: Initialized tensor.
    """
    init = Constant(value=val)

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def ones_(
    tensor: paddle.Tensor,
) -> paddle.Tensor:
    """Fill the input Tensor with the scalar value 1.

    Args:
        tensor (Tensor):  Paddle Tensor.

    Returns:
        Tensor: Initialized tensor.
    """
    init = Constant(value=1.0)

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def zeros_(
    tensor: paddle.Tensor,
) -> paddle.Tensor:
    """Fill the input Tensor with the scalar value 0.

    Args:
        tensor (Tensor):  Paddle Tensor.

    Returns:
        Tensor: Initialized tensor.
    """
    init = Constant(value=0.0)

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def dirac_(
    tensor: paddle.Tensor,
    groups: int = 1,
) -> paddle.Tensor:
    """Initialize the 3D/4D/5D Tensor with Dirac delta function.

    Args:
        tensor (Tensor):  Paddle Tensor.
        groups (int|None, optional): 0-dimension of the Tensor will be divided by groups,
            each group has the same value. Default: 1.
    Returns:
        Tensor: Initialized tensor.
    """
    init = Dirac(groups=groups)

    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def eye_(
    tensor: paddle.Tensor,
) -> paddle.Tensor:
    """Fill the 2-dimensional input Tensor with the identity matrix.

    Args:
        tensor (Tensor):  Paddle Tensor.
    Returns:
        Tensor: Initialized tensor.
    """

    if len(tensor.shape) != 2:
        raise AssertionError(
            f"Only support 2 dimensional tensor, but got {len(tensor.shape)}."
        )

    if in_dygraph_mode():
        new_tensor = paddle.eye(
            tensor.shape[0], tensor.shape[1], dtype=tensor.dtype
        )
        new_tensor._share_underline_tensor_to(tensor)
        return tensor
    elif in_pir_mode():
        new_tensor = paddle.eye(
            tensor.shape[0], tensor.shape[1], dtype=tensor.dtype
        )
        return new_tensor
    else:
        raise NotImplementedError(
            'Only support run in dygraph mode or PIR mode.'
        )


def orthogonal_(
    tensor: paddle.Tensor,
    gain: float = 1,
) -> paddle.Tensor:
    """Fill the input Tensor with a (semi) orthogonal matrix.

    Args:
        tensor (Tensor):  Paddle Tensor.
        gain(float, optional): The multiplication coefficient for initialized tensor. Default: 1.0.
    Returns:
        Tensor: Initialized tensor.
    """
    init = Orthogonal(gain=gain)
    if in_dygraph_mode():
        init(tensor)
        return tensor
    return init(tensor)


def sparse_(
    tensor: paddle.Tensor, sparsity: float, std: float = 0.01
) -> paddle.Tensor:
    """Fill the 2D input Tensor as a sparse matrix.

    The non-zero elements will be drawn from the normal distribution.

    Args:
        tensor (Tensor): Paddle Tensor with 2 dimensions.
        sparsity (float): The fraction of elements in each column to be set to zero.
        std (float): the standard deviation of the normal distribution used to generate
            the non-zero values. Default is 0.01.

    Examples:
        >>> tensor = paddle.empty(3, 5)
        >>> result = paddle.nn.init.sparse_(tensor, sparsity=0.1)
    """
    if tensor.ndimension() != 2:
        raise ValueError("Only tensors with 2 dimensions are supported")
    rows, cols = tensor.shape
    num_zeros = math.ceil(sparsity * rows)

    with paddle.no_grad():
        tensor = normal_(tensor, mean=0, std=std)
        for col_idx in range(cols):
            row_indices = paddle.randperm(rows)
            zero_indices = row_indices[:num_zeros]
            tensor[zero_indices, col_idx] = 0
    return tensor


def _make_deprecate(func: Callable[_InputT, _RetT]) -> Callable[_InputT, _RetT]:
    new_name = func.__name__
    old_name = new_name[:-1]

    def deprecated_init(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
        warnings.warn(
            f"`nn.init.{old_name}` is now deprecated in favor of `nn.init.{new_name}`.",
            FutureWarning,
            stacklevel=2,
        )
        return func(*args, **kwargs)

    deprecated_init.__doc__ = rf"""
    {old_name}(...)

    .. warning::
        This method is now deprecated in favor of :func:`paddle.nn.init.{new_name}`.

    See :func:`~paddle.nn.init.{new_name}` for details."""
    deprecated_init.__name__ = old_name
    return deprecated_init


uniform = _make_deprecate(uniform_)
normal = _make_deprecate(normal_)
constant = _make_deprecate(constant_)
eye = _make_deprecate(eye_)
dirac = _make_deprecate(dirac_)
xavier_uniform = _make_deprecate(xavier_uniform_)
xavier_normal = _make_deprecate(xavier_normal_)
kaiming_uniform = _make_deprecate(kaiming_uniform_)
kaiming_normal = _make_deprecate(kaiming_normal_)
orthogonal = _make_deprecate(orthogonal_)
sparse = _make_deprecate(sparse_)