from __future__ import annotations

import typing
import abc
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set, Tuple, Union

from zyplib.import_tools import import_lazy

if TYPE_CHECKING:
    import numpy as np
    import torch
    import numpy as NumpyType
    import torch as TorchType
else:
    np = import_lazy('numpy')
    torch = import_lazy('torch')

Shape = Tuple[int, ...]
DType = Union['NumpyType.dtype', 'TorchType.dtype']
TensorLike = Union['NumpyType.ndarray', 'TorchType.Tensor']


class TensorLikeBackend(abc.ABC):
    """
    Abstract base class for a tensor-like data manipulation backend.

    This class defines a common interface for operations on array-like
    data structures, such as NumPy ndarrays and PyTorch tensors.
    """

    backend_type: Literal['numpy', 'torch'] | None = None

    @abc.abstractmethod
    def type(
        self, dtype_like: Union[str, type, 'NumpyType.dtype', 'TorchType.dtype']
    ) -> DType:
        """
        Returns the backend's native data type from various representations.

        Args:
            dtype_like: The data type representation. Can be:
                - A string like 'float32', 'int', 'bool'.
                - A Python type like `float`, `int`, `bool`.
                - A NumPy dtype object (e.g., np.float32).
                - A PyTorch dtype object (e.g., torch.float32).

        Returns:
            The backend's native dtype object (np.dtype or torch.dtype).
        """
        pass

    @abc.abstractmethod
    def as_tensorlike(self, data: Any, **kwargs) -> TensorLike:
        """Converts data to the backend's native tensor format."""
        pass

    @abc.abstractmethod
    def is_finite(self, tensor: TensorLike) -> TensorLike:
        """检查张量元素是否有限"""
        pass

    @abc.abstractmethod
    def requires_grad(self, tensor: TensorLike) -> bool:
        """检查张量是否需要梯度（仅 torch）"""
        pass

    @abc.abstractmethod
    def detach(self, tensor: TensorLike) -> TensorLike:
        """分离张量（仅 torch，numpy 返回自身）"""
        pass

    @abc.abstractmethod
    def shape(self, tensor: TensorLike) -> Tuple[int, ...]:
        """Returns the shape of the tensor."""
        pass

    @abc.abstractmethod
    def device(self, tensor: TensorLike) -> Optional[str]:
        """Returns the device of the tensor."""
        pass

    @abc.abstractmethod
    def to_device(self, tensor: TensorLike, device: str) -> TensorLike:
        """Moves the tensor to the specified device."""
        pass

    @abc.abstractmethod
    def sum(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> TensorLike:
        """Computes the sum of tensor elements."""
        pass

    @abc.abstractmethod
    def mean(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> TensorLike:
        """Computes the mean of tensor elements."""
        pass

    @abc.abstractmethod
    def matmul(self, tensor1: TensorLike, tensor2: TensorLike) -> TensorLike:
        """Performs matrix multiplication."""
        pass

    @abc.abstractmethod
    def reshape(self, tensor: TensorLike, new_shape: Tuple[int, ...]) -> TensorLike:
        """Reshapes the tensor."""
        pass

    @abc.abstractmethod
    def concatenate(
        self, tensors: typing.Sequence[TensorLike], axis: int = 0
    ) -> TensorLike:
        """Concatenates tensors along a given axis."""
        pass

    @abc.abstractmethod
    def ones(self, shape: Tuple[int, ...], **kwargs) -> TensorLike:
        """Creates a tensor of ones."""
        pass

    @abc.abstractmethod
    def zeros(self, shape: Tuple[int, ...], **kwargs) -> TensorLike:
        """Creates a tensor of zeros."""
        pass

    @abc.abstractmethod
    def numel(self, tensor: TensorLike) -> int:
        """Returns the number of elements in the tensor."""
        pass

    @abc.abstractmethod
    def dtype(self, tensor: TensorLike) -> DType:
        """Returns the data type of the tensor."""
        pass

    @abc.abstractmethod
    def astype(self, tensor: TensorLike, dtype: DType) -> TensorLike:
        """Converts the tensor to the specified data type."""
        pass

    @abc.abstractmethod
    def ndim(self, tensor: TensorLike) -> int:
        """Returns the number of dimensions of the tensor."""
        pass

    @abc.abstractmethod
    def squeeze(self, tensor: TensorLike, axis: Optional[int] = None) -> TensorLike:
        """Removes single-dimensional entries from the shape of a tensor."""
        pass

    @abc.abstractmethod
    def unsqueeze(self, tensor: TensorLike, axis: int) -> TensorLike:
        """Inserts a new axis into the tensor."""
        pass

    @abc.abstractmethod
    def max(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> TensorLike:
        """Computes the maximum of tensor elements."""
        pass

    @abc.abstractmethod
    def min(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> TensorLike:
        """Computes the minimum of tensor elements."""
        pass

    @abc.abstractmethod
    def argmax(
        self, tensor: TensorLike, axis: Optional[Union[int, Tuple[int, ...]]] = None
    ) -> TensorLike:
        """Returns the indices of the maximum values along an axis."""
        pass

    @abc.abstractmethod
    def argmin(
        self, tensor: TensorLike, axis: Optional[Union[int, Tuple[int, ...]]] = None
    ) -> TensorLike:
        """Returns the indices of the minimum values along an axis."""
        pass

    @abc.abstractmethod
    def abs(self, tensor: TensorLike) -> TensorLike:
        """Computes the absolute value of each element."""
        pass

    @abc.abstractmethod
    def sqrt(self, tensor: TensorLike) -> TensorLike:
        """Computes the non-negative square-root of each element."""
        pass

    @abc.abstractmethod
    def exp(self, tensor: TensorLike) -> TensorLike:
        """Computes the exponential of each element."""
        pass

    @abc.abstractmethod
    def log(self, tensor: TensorLike) -> TensorLike:
        """Computes the natural logarithm of each element."""
        pass

    @abc.abstractmethod
    def pow(
        self, tensor: TensorLike, exponent: Union[int, float, TensorLike]
    ) -> TensorLike:
        """计算张量的指数"""
        pass

    @abc.abstractmethod
    def transpose(self, tensor: TensorLike, axes: Tuple[int, ...]) -> TensorLike:
        """Permutes the dimensions of a tensor."""
        pass

    @abc.abstractmethod
    def repeat(
        self,
        tensor: TensorLike,
        repeats: Union[int, TensorLike],
        axis: Optional[int] = None,
    ) -> TensorLike:
        """重复张量的元素。

        沿指定轴重复张量中的每个元素。行为等价于 numpy.repeat 和 torch.repeat_interleave。

        Parameters
        ----------
        - `tensor` : `TensorLike`
            - 输入张量
        - `repeats` : `Union[int, TensorLike]`
            - 每个元素的重复次数。如果是整数，则所有元素重复相同次数；
              如果是数组，则会广播到指定轴的形状
        - `axis` : `Optional[int]`, optional
            - 沿哪个轴重复值。如果为 None，则使用扁平化的输入数组并返回扁平输出数组，by default None

        Returns
        ----------
        - `TensorLike`
            - 重复后的张量，除了指定轴（或扁平化）外，形状与输入相同
        """
        pass

    @abc.abstractmethod
    def tile(self, tensor: TensorLike, reps: Union[int, Tuple[int, ...]]) -> TensorLike:
        """瓦片式重复张量。

        通过沿每个维度重复整个张量来构造数组。行为等价于 numpy.tile 和 torch.repeat。

        Parameters
        ----------
        - `tensor` : `TensorLike`
            - 输入张量
        - `reps` : `Union[int, Tuple[int, ...]]`
            - 沿每个轴的重复次数。如果是整数，则在所有维度重复该次数；
              如果是元组，则指定每个维度的重复次数

        Returns
        ----------
        - `TensorLike`
            - 瓦片重复后的张量
        """
        pass

    @abc.abstractmethod
    def stack(self, tensors: typing.Sequence[TensorLike], axis: int = 0) -> TensorLike:
        """Stacks a sequence of tensors along a new axis."""
        pass

    @abc.abstractmethod
    def full(self, shape: Shape, fill_value: Union[int, float], **kwargs) -> TensorLike:
        """Creates a new tensor of a given shape filled with a value."""
        pass

    @abc.abstractmethod
    def rand(self, shape: Shape, **kwargs) -> TensorLike:
        """Creates a tensor with random values uniformly distributed between 0 and 1."""
        pass

    @abc.abstractmethod
    def randn(self, shape: Shape, **kwargs) -> TensorLike:
        """Creates a tensor with random values from a standard normal distribution."""
        pass

    @abc.abstractmethod
    def clip(
        self, tensor: TensorLike, min_val: Union[int, float], max_val: Union[int, float]
    ) -> TensorLike:
        """Clips (limits) the values in a tensor to a specified range."""
        pass

    @abc.abstractmethod
    def where(self, condition: TensorLike, x: TensorLike, y: TensorLike) -> TensorLike:
        """Returns elements chosen from x or y depending on condition."""
        pass

    @abc.abstractmethod
    def allclose(
        self,
        tensor1: TensorLike,
        tensor2: TensorLike,
        rtol: float = 1e-5,
        atol: float = 1e-8,
    ) -> bool:
        """Returns True if two tensors are element-wise equal within a tolerance."""
        pass

    @abc.abstractmethod
    def all(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> TensorLike:
        """Tests if all elements in the tensor evaluate to True."""
        pass

    @abc.abstractmethod
    def any(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> TensorLike:
        """Tests if any element in the tensor evaluates to True."""
        pass

    @abc.abstractmethod
    def zeros_like(self, tensor: TensorLike, **kwargs) -> TensorLike:
        """Creates a tensor of zeros with the same shape as the input tensor."""
        pass

    @abc.abstractmethod
    def ones_like(self, tensor: TensorLike, **kwargs) -> TensorLike:
        """Creates a tensor of ones with the same shape as the input tensor."""
        pass

    @abc.abstractmethod
    def norm(
        self,
        tensor: TensorLike,
        ord: Optional[Union[int, float, str]] = None,
        axis: Optional[Union[int, tuple]] = None,
        keepdims: bool = False,
        method: str = 'vectorized',  # 新增参数
    ) -> TensorLike:
        """计算张量的范数。

        Parameters:
        -----------
        - `tensor` : TensorLike
            输入张量
        - `ord` : Optional[Union[int, float, str]]
            范数的阶数
        - `axis` : Optional[Union[int, tuple]]
            计算范数的轴
        - `keepdims` : bool
            是否保持维度
        - `method` : str
            范数计算方法，支持：
            - 'vectorized': 将张量展平后计算范数（torch默认行为）
            - 'matrix': 计算矩阵诱导范数（numpy默认行为）
            - 'auto': 根据后端类型自动选择

        Returns:
        --------
        - `TensorLike`
            范数结果
        """
        pass

    @abc.abstractmethod
    def std(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, tuple]] = None,
        keepdims: bool = False,
        ddof: int = 0,
    ) -> TensorLike:
        """Computes the standard deviation of tensor elements.

        Parameters
        ----------
        tensor : TensorLike
            Input tensor
        axis : Optional[Union[int, tuple]], optional
            Axis or axes along which the standard deviation is computed
        keepdims : bool, optional
            If True, the reduced dimensions are retained with length 1
        ddof : int, optional
            Delta Degrees of Freedom. The divisor used in calculations is N - ddof,
            where N represents the number of elements. Default is 0 (population standard deviation).

        Returns
        -------
        TensorLike
            Standard deviation of the tensor elements
        """
        pass

    @abc.abstractmethod
    def var(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, tuple]] = None,
        keepdims: bool = False,
        ddof: int = 0,
    ) -> TensorLike:
        """Computes the variance of tensor elements.

        Parameters
        ----------
        tensor : TensorLike
            Input tensor
        axis : Optional[Union[int, tuple]], optional
            Axis or axes along which the variance is computed
        keepdims : bool, optional
            If True, the reduced dimensions are retained with length 1
        ddof : int, optional
            Delta Degrees of Freedom. The divisor used in calculations is N - ddof,
            where N represents the number of elements. Default is 0 (population variance).

        Returns
        -------
        TensorLike
            Variance of the tensor elements
        """
        pass

    @abc.abstractmethod
    def prod(
        self,
        tensor: TensorLike,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> TensorLike:
        """Computes the product of tensor elements."""
        pass

    @abc.abstractmethod
    def tensor_equal(self, tensor1: TensorLike, tensor2: TensorLike) -> bool:
        """Returns True if two tensors have the same shape and elements, False otherwise."""
        pass

    @abc.abstractmethod
    def arange(
        self,
        start: Union[int, float],
        stop: Optional[Union[int, float]] = None,
        step: Union[int, float] = 1,
        **kwargs,
    ) -> TensorLike:
        """Creates a 1-D tensor with evenly spaced values within a given interval."""
        pass

    @abc.abstractmethod
    def linspace(
        self, start: Union[int, float], stop: Union[int, float], num: int = 50, **kwargs
    ) -> TensorLike:
        """Creates a 1-D tensor with evenly spaced values over a specified interval."""
        pass

    @abc.abstractmethod
    def eye(self, n: int, m: Optional[int] = None, **kwargs) -> TensorLike:
        """Creates a 2-D tensor with ones on the diagonal and zeros elsewhere."""
        pass

    @abc.abstractmethod
    def maximum(
        self, tensor1: TensorLike, tensor2: Union[TensorLike, float]
    ) -> TensorLike:
        """Computes the element-wise maximum of two tensors."""
        pass

    @abc.abstractmethod
    def to_numpy(self, tensor: TensorLike) -> 'NumpyType.ndarray':
        """Converts a tensor to a NumPy array."""
        pass

    @abc.abstractmethod
    def round(self, tensor: TensorLike, decimals: int = 0) -> TensorLike:
        """Rounds the elements of a tensor to a specified number of decimal places."""
        pass

    def divide(
        self, tensor1: TensorLike, tensor2: Union[TensorLike, float]
    ) -> TensorLike:
        """Element-wise division of two tensors."""
        return tensor1 / tensor2

    def dot(self, tensor1: TensorLike, tensor2: TensorLike) -> TensorLike:
        """Dot product of two tensors.

        Note: This method is only intended for use with 1-D vectors. For matrix multiplication,
        use the `matmul` method instead. For general tensors, this computes the Frobenius inner product.

        Parameters
        ----------
        tensor1 : TensorLike
            First tensor (preferably 1-D)
        tensor2 : TensorLike
            Second tensor (preferably 1-D)

        Returns
        -------
        TensorLike
            Dot product result
        """
        return self.sum(tensor1 * tensor2)

    def item(self, tensor_like: TensorLike) -> float | complex:
        """Extracts a scalar value from a tensor."""
        if isinstance(tensor_like, (int, float, complex)):
            return tensor_like
        return tensor_like.item()

    @abc.abstractmethod
    def unique(
        self,
        tensor: TensorLike,
        return_index: bool = False,
        return_inverse: bool = False,
        return_counts: bool = False,
        axis: Optional[int] = None,
        **kwargs: Any,
    ) -> Union[TensorLike, Tuple[TensorLike, ...]]:
        pass

    @abc.abstractmethod
    def quantile(
        self,
        tensor: TensorLike,
        q: Union[float, TensorLike],
        axis: Optional[int] = None,
        keepdims: bool = False,
        interpolation: str = 'linear',
    ) -> TensorLike:
        """
        Computes the q-th quantile of a tensor.

        Parameters
        ----------
        tensor : TensorLike
            Input tensor.
        q : Union[float, TensorLike]
            Quantile or sequence of quantiles to compute, which must be between 0 and 1 inclusive.
        axis : Optional[int], optional
            Axis along which the quantiles are computed.
        keepdims : bool, optional
            If True, the reduced dimensions are retained with length 1.

        Returns
        -------
        TensorLike
            The quantile(s) of the tensor.
        """
        pass

    @abc.abstractmethod
    def searchsorted(
        self,
        sorted_sequence: TensorLike,
        values: TensorLike,
        right: bool = False,
        side: Optional[Literal['left', 'right']] = None,
    ) -> TensorLike:
        """
        Finds the indices into a sorted sequence such that, if the corresponding
        values were inserted, the order would be preserved.

        Parameters
        ----------
        sorted_sequence : TensorLike
            The sorted 1-D sequence.
        values : TensorLike
            Values to insert into the sequence.
        right : bool, optional
            If True, the insertion point will be to the right of any existing entries.
            Used for compatibility with PyTorch.
        side : {'left', 'right'}, optional
            If 'left', the insertion point is to the left. If 'right', to the right.
            Used for compatibility with NumPy.

        Returns
        -------
        TensorLike
            Array of insertion points.
        """
        pass


_TMapping = typing.TypeVar('_TMapping', bound=typing.Mapping[str, Any])


@typing.overload
def filter_dict(
    d: _TMapping, *, allowed: None = None, prohibited: None = None
) -> _TMapping: ...


@typing.overload
def filter_dict(
    d: typing.Mapping[str, Any],
    *,
    allowed: typing.Sequence[str] | typing.AbstractSet[str] | None = None,
    prohibited: typing.Sequence[str] | typing.AbstractSet[str] | None = None,
) -> Dict[str, Any]: ...


def filter_dict(
    d: typing.Mapping[str, Any],
    *,
    allowed: typing.Sequence[str] | typing.AbstractSet[str] | None = None,
    prohibited: typing.Sequence[str] | typing.AbstractSet[str] | None = None,
) -> typing.Mapping[str, Any] | Dict[str, Any]:
    """Filters a dictionary to only include allowed keys."""
    if allowed is None and prohibited is None:
        return d
    # return {k: v for k, v in d.items() if k in allowed}
    if allowed is None:
        allowed = set(d.keys())
    if prohibited is None:
        prohibited = set()
    return {k: v for k, v in d.items() if k in allowed and k not in prohibited}
