import typing
from typing import Any, List, Literal, Optional, Set, Tuple, Union

import numpy as np
import torch

from ._base import DType, Shape, TensorLike, TensorLikeBackend, filter_dict


class NumpyBackend(TensorLikeBackend):
    """Backend for NumPy ndarrays."""

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

    _DTYPE_MAP = {
        # String to numpy type
        'float16': np.float16,
        'half': np.float16,
        'float32': np.float32,
        'float': np.float32,
        'float64': np.float64,
        'double': np.float64,
        'int8': np.int8,
        'int16': np.int16,
        'int32': np.int32,
        'int': np.int32,
        'int64': np.int64,
        'bool': np.bool_,
        # Python type to numpy type
        float: np.float32,  # Default float
        int: np.int32,  # Default int
        bool: np.bool_,
        # PyTorch to numpy type
        torch.float16: np.float16,
        torch.float32: np.float32,
        torch.float64: np.float64,
        torch.int8: np.int8,
        torch.int16: np.int16,
        torch.int32: np.int32,
        torch.int64: np.int64,
        torch.bool: np.bool_,
    }

    def type(self, dtype_like: Any) -> np.dtype:
        if self._DTYPE_MAP.get(dtype_like):
            return self._DTYPE_MAP[dtype_like]
        try:
            return np.dtype(dtype_like)
        except TypeError:
            raise ValueError(f'Unsupported or unknown dtype identifier: {dtype_like}')

    def as_tensorlike(self, data: Any, **kwargs) -> np.ndarray:
        # Pop torch-specific kwargs
        kwargs.pop('device', None)
        kwargs.pop('requires_grad', None)

        if 'dtype' in kwargs:
            kwargs['dtype'] = self.type(kwargs['dtype'])
            if kwargs['dtype'] is None:
                del kwargs['dtype']  # 清理掉 None

        # For other data types, create a new array
        return np.asarray(data, **kwargs)

    def is_finite(self, tensor: np.ndarray) -> np.ndarray:
        return np.isfinite(tensor)

    def requires_grad(self, tensor: np.ndarray) -> bool:
        return False

    def detach(self, tensor: np.ndarray) -> np.ndarray:
        return tensor

    def shape(self, tensor: np.ndarray) -> Tuple[int, ...]:
        return tensor.shape

    def device(self, tensor: np.ndarray) -> Optional[str]:
        return 'cpu'

    def to_device(self, tensor: np.ndarray, device: str) -> np.ndarray:
        return tensor

    def sum(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> np.ndarray:
        return np.sum(tensor, axis=axis, keepdims=keepdims)

    def mean(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> np.ndarray:
        return np.mean(tensor, axis=axis, keepdims=keepdims)

    def matmul(self, tensor1: np.ndarray, tensor2: np.ndarray) -> np.ndarray:
        return np.matmul(tensor1, tensor2)

    def reshape(self, tensor: np.ndarray, new_shape: Tuple[int, ...]) -> np.ndarray:
        return tensor.reshape(new_shape)

    def concatenate(self, tensors: typing.Sequence, axis: int = 0) -> np.ndarray:
        return np.concatenate(tensors, axis=axis)

    def ones(self, shape: Tuple[int, ...], **kwargs) -> np.ndarray:
        kwargs = filter_dict(kwargs, allowed=['dtype', 'order', 'like'])
        return np.ones(shape, **kwargs)

    def zeros(self, shape: Tuple[int, ...], **kwargs) -> np.ndarray:
        kwargs = filter_dict(kwargs, allowed=['dtype', 'order', 'like'])
        return np.zeros(shape, **kwargs)

    def numel(self, tensor: np.ndarray) -> int:
        return tensor.size

    def dtype(self, tensor: np.ndarray) -> np.dtype:
        return tensor.dtype

    def astype(self, tensor: np.ndarray, dtype: np.dtype) -> np.ndarray:
        return tensor.astype(dtype)

    def ndim(self, tensor: np.ndarray) -> int:
        return tensor.ndim

    def squeeze(self, tensor: np.ndarray, axis: Optional[int] = None) -> np.ndarray:
        return tensor.squeeze(axis=axis)

    def unsqueeze(self, tensor: np.ndarray, axis: int) -> np.ndarray:
        return np.expand_dims(tensor, axis)

    def max(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> np.ndarray:
        return np.max(tensor, axis=axis, keepdims=keepdims)

    def min(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> np.ndarray:
        return np.min(tensor, axis=axis, keepdims=keepdims)

    def argmax(
        self, tensor: np.ndarray, axis: Optional[Union[int, Tuple[int, ...]]] = None
    ) -> np.ndarray:
        return np.argmax(tensor, axis=axis)  # pyright: ignore[reportArgumentType,reportCallIssue]

    def argmin(
        self, tensor: np.ndarray, axis: Optional[Union[int, Tuple[int, ...]]] = None
    ) -> np.ndarray:
        return np.argmin(tensor, axis=axis)  # pyright: ignore[reportArgumentType,reportCallIssue]

    def abs(self, tensor: np.ndarray) -> np.ndarray:
        return np.abs(tensor)

    def sqrt(self, tensor: np.ndarray) -> np.ndarray:
        return np.sqrt(tensor)

    def exp(self, tensor: np.ndarray) -> np.ndarray:
        return np.exp(tensor)

    def log(self, tensor: np.ndarray) -> np.ndarray:
        return np.log(tensor)

    def pow(
        self, tensor: np.ndarray, exponent: Union[int, float, np.ndarray]
    ) -> np.ndarray:
        return np.power(tensor, exponent)

    def transpose(self, tensor: np.ndarray, axes: Tuple[int, ...]) -> np.ndarray:
        return np.transpose(tensor, axes)

    def repeat(
        self,
        tensor: np.ndarray,
        repeats: Union[int, np.ndarray],
        axis: Optional[int] = None,
    ) -> np.ndarray:
        return np.repeat(tensor, repeats, axis=axis)

    def tile(self, tensor: np.ndarray, reps: Union[int, Tuple[int, ...]]) -> np.ndarray:
        return np.tile(tensor, reps)

    def stack(self, tensors: typing.Sequence[np.ndarray], axis: int = 0) -> np.ndarray:
        return np.stack(tensors, axis=axis)

    def full(self, shape: Shape, fill_value: Union[int, float], **kwargs) -> np.ndarray:
        kwargs = filter_dict(kwargs, allowed=['dtype', 'order', 'like'])
        return np.full(shape, fill_value, **kwargs)

    def rand(self, shape: Shape, **kwargs) -> np.ndarray:
        if 'dtype' in kwargs:
            return np.random.rand(*shape).astype(kwargs['dtype'])
        return np.random.rand(*shape)

    def randn(self, shape: Shape, **kwargs) -> np.ndarray:
        if 'dtype' in kwargs:
            return np.random.randn(*shape).astype(kwargs['dtype'])
        return np.random.randn(*shape)

    def clip(
        self, tensor: np.ndarray, min_val: Union[int, float], max_val: Union[int, float]
    ) -> np.ndarray:
        return np.clip(tensor, min_val, max_val)

    def where(self, condition: np.ndarray, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        return np.where(condition, x, y)

    def allclose(
        self,
        tensor1: np.ndarray,
        tensor2: np.ndarray,
        rtol: float = 1e-5,
        atol: float = 1e-8,
    ) -> bool:
        return np.allclose(tensor1, tensor2, rtol=rtol, atol=atol)

    def all(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> np.ndarray:
        return np.all(tensor, axis=axis, keepdims=keepdims)

    def any(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> np.ndarray:
        return np.any(tensor, axis=axis, keepdims=keepdims)

    def zeros_like(self, tensor: np.ndarray, **kwargs) -> np.ndarray:
        kwargs = filter_dict(kwargs, allowed=['dtype', 'order', 'subok', 'shape'])
        return np.zeros_like(tensor, **kwargs)

    def ones_like(self, tensor: np.ndarray, **kwargs) -> np.ndarray:
        kwargs = filter_dict(kwargs, allowed=['dtype', 'order', 'subok', 'shape'])
        return np.ones_like(tensor, **kwargs)

    def norm(
        self,
        tensor: np.ndarray,
        ord: Optional[Union[int, float, str]] = None,
        axis: Optional[Union[int, tuple]] = None,
        keepdims: bool = False,
        method: str = 'vectorized',
    ) -> np.ndarray:
        """
        计算张量的范数。

        Parameters:
        -----------
        method : str
            - 'vectorized': 将张量展平后计算范数（与torch保持一致）
            - 'matrix': 计算矩阵诱导范数（numpy原生行为）
            - 'auto': 对numpy使用'matrix'方法
        """
        # 参数验证
        if method not in ['vectorized', 'matrix', 'auto']:
            raise ValueError(
                f"method must be one of ['vectorized', 'matrix', 'auto'], got {method}"
            )

        # 自动选择方法
        if method == 'auto':
            method = 'matrix'  # numpy默认使用矩阵范数

        if axis is None:
            if method == 'vectorized':
                # 展平张量以与torch保持一致
                flattened = tensor.flatten()
                result = np.linalg.norm(flattened, ord=ord)  # pyright: ignore[reportArgumentType,reportCallIssue]
                if keepdims:
                    # 保持原始维度结构
                    shape = [1] * tensor.ndim
                    result = result.reshape(shape)
                return result
            elif method == 'matrix':
                # 使用numpy的默认矩阵范数行为
                result = np.linalg.norm(tensor, ord=ord)  # pyright: ignore[reportArgumentType,reportCallIssue]
                if keepdims:
                    shape = [1] * tensor.ndim
                    result = result.reshape(shape)
                return result
        else:
            # 指定了axis时，两种方法行为相同
            return np.linalg.norm(tensor, ord=ord, axis=axis, keepdims=keepdims)  # pyright: ignore[reportArgumentType,reportCallIssue]

        raise AssertionError('unreachable')

    def std(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, tuple]] = None,
        keepdims: bool = False,
        ddof: int = 0,
    ) -> np.ndarray:
        return np.std(tensor, axis=axis, keepdims=keepdims, ddof=ddof)

    def var(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, tuple]] = None,
        keepdims: bool = False,
        ddof: int = 0,
    ) -> np.ndarray:
        return np.var(tensor, axis=axis, keepdims=keepdims, ddof=ddof)

    def prod(
        self,
        tensor: np.ndarray,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> np.ndarray:
        return np.prod(tensor, axis=axis, keepdims=keepdims)

    def tensor_equal(self, tensor1: np.ndarray, tensor2: np.ndarray) -> bool:
        return np.array_equal(tensor1, tensor2)

    def arange(
        self,
        start: Union[int, float],
        stop: Optional[Union[int, float]] = None,
        step: Union[int, float] = 1,
        **kwargs,
    ) -> np.ndarray:
        if stop is None:
            return np.arange(0, start, step, **kwargs)
        return np.arange(start, stop, step, **kwargs)

    def linspace(
        self, start: Union[int, float], stop: Union[int, float], num: int = 50, **kwargs
    ) -> np.ndarray:
        kwargs = filter_dict(kwargs, allowed=['dtype', 'axis', 'endpoint', 'retstep'])
        return np.linspace(start, stop, num=num, **kwargs)

    def eye(self, n: int, m: Optional[int] = None, **kwargs) -> np.ndarray:
        kwargs = filter_dict(kwargs, allowed=['k', 'dtype', 'order', 'like'])
        return np.eye(n, M=m, **kwargs)

    def maximum(
        self, tensor1: np.ndarray, tensor2: Union[np.ndarray, float]
    ) -> np.ndarray:
        return np.maximum(tensor1, tensor2)

    def to_numpy(self, tensor: np.ndarray) -> np.ndarray:
        return tensor

    def round(self, tensor: np.ndarray, decimals: int = 0) -> np.ndarray:
        return np.round(tensor, decimals=decimals)

    def unique(
        self,
        tensor: np.ndarray,
        return_index: bool = False,
        return_inverse: bool = False,
        return_counts: bool = False,
        axis: Optional[int] = None,
    ) -> Union[np.ndarray, Tuple[np.ndarray, ...]]:
        # np.unique 直接支持所有参数
        result = np.unique(
            tensor,
            return_index=return_index,
            return_inverse=return_inverse,
            return_counts=return_counts,
            axis=axis,
        )

        # 如果没有请求额外的返回，np.unique 只返回一个数组
        if not any([return_index, return_inverse, return_counts]):
            return result

        # 否则，它返回一个元组，我们直接返回即可
        return result

    def quantile(
        self, tensor, q, axis=None, keepdims=False, interpolation: str = 'linear'
    ):
        # NumPy >= 1.22 支持 keepdims 和 interpolation
        return np.quantile(  # pyright: ignore[reportArgumentType,reportCallIssue]
            tensor,
            q,
            axis=axis,
            keepdims=keepdims,
            method=typing.cast(typing.Any, interpolation),
        )

    def searchsorted(
        self,
        sorted_sequence: np.ndarray,
        values: np.ndarray | int | float,
        right: bool = False,
        side: Literal['left', 'right'] | None = None,
    ) -> np.ndarray:
        # 统一接口，NumPy 使用 'side'
        if side is None:
            side = 'right' if right else 'left'
        return np.searchsorted(sorted_sequence, values, side=side)
