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 TorchBackend(TensorLikeBackend):
    """Backend for PyTorch tensors."""

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

    _DTYPE_MAP = {
        # Python types
        float: torch.float32,
        int: torch.int64,
        bool: torch.bool,
        # String aliases
        'float': torch.float32,
        'float16': torch.float16,
        'half': torch.float16,
        'float32': torch.float32,
        'float64': torch.float64,
        'double': torch.float64,
        'int': torch.int64,
        'int8': torch.int8,
        'int16': torch.int16,
        'short': torch.int16,
        'int32': torch.int32,
        'int64': torch.int64,
        'long': torch.int64,
        'bool': torch.bool,
        # NumPy dtypes
        np.float16: torch.float16,
        np.float32: torch.float32,
        np.float64: torch.float64,
        np.int8: torch.int8,
        np.int16: torch.int16,
        np.int32: torch.int32,
        np.int64: torch.int64,
        np.bool_: torch.bool,
    }

    def type(self, dtype_like: Any) -> Optional[torch.dtype]:
        if dtype_like is None:
            return None

        if isinstance(dtype_like, torch.dtype):
            return dtype_like

        try:
            return self._DTYPE_MAP[dtype_like]
        except KeyError:
            # Fallback for other potential string representations
            if isinstance(dtype_like, str):
                dtype_str = dtype_like.lower().replace('torch.', '')
                if hasattr(torch, dtype_str):
                    return getattr(torch, dtype_str)

            raise ValueError(f"Cannot convert '{dtype_like}' to a PyTorch dtype.")

    def as_tensorlike(self, data: Any, **kwargs) -> torch.Tensor:
        # 预处理 kwargs，转换 dtype
        if 'dtype' in kwargs:
            kwargs['dtype'] = self.type(kwargs['dtype'])

        if isinstance(data, torch.Tensor):
            tensor = data
            if 'device' in kwargs:
                tensor = tensor.to(device=kwargs['device'])
            if 'dtype' in kwargs:
                tensor = tensor.to(dtype=kwargs['dtype'])
            if 'requires_grad' in kwargs:
                tensor = tensor.requires_grad_(kwargs['requires_grad'])
        elif isinstance(data, np.ndarray):
            # 注意：torch.tensor(numpy_array) 默认会复制数据
            # 如果想保持 from_numpy 的零拷贝特性，需要分步操作
            tensor = torch.from_numpy(data)
            if 'dtype' in kwargs:
                tensor = tensor.to(dtype=kwargs['dtype'])
            if 'device' in kwargs:
                tensor = tensor.to(device=kwargs['device'])
            if kwargs.get('requires_grad', False):
                # 只有浮点数或复数张量可以设置 requires_grad
                if tensor.is_floating_point() or tensor.is_complex():
                    tensor.requires_grad_(True)
                # 如果用户对整数张量设置 requires_grad，PyTorch 会报错，这里让它自然报错即可
        else:
            # torch.tensor 可以直接接受 requires_grad 等参数
            tensor = torch.tensor(data, **kwargs)

        return tensor

    def is_finite(self, tensor: torch.Tensor) -> torch.Tensor:
        return torch.isfinite(tensor)

    def requires_grad(self, tensor: torch.Tensor) -> bool:
        return tensor.requires_grad

    def detach(self, tensor: torch.Tensor) -> torch.Tensor:
        return tensor.detach()

    def shape(self, tensor: torch.Tensor) -> Tuple[int, ...]:
        return tuple(tensor.shape)

    def device(self, tensor: torch.Tensor) -> Optional[str]:
        return str(tensor.device)

    def to_device(self, tensor: torch.Tensor, device: str) -> torch.Tensor:
        return tensor.to(device)

    def sum(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> torch.Tensor:
        if axis is None:
            return torch.sum(tensor)
        return torch.sum(tensor, dim=axis, keepdim=keepdims)

    def mean(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> torch.Tensor:
        if axis is None:
            return torch.mean(tensor)
        return torch.mean(tensor, dim=axis, keepdim=keepdims)

    def matmul(self, tensor1: torch.Tensor, tensor2: torch.Tensor) -> torch.Tensor:
        return torch.matmul(tensor1, tensor2)

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

    def concatenate(self, tensors: typing.Sequence, axis: int = 0) -> torch.Tensor:
        return torch.cat(list(tensors), dim=axis)

    def ones(self, shape: Tuple[int, ...], **kwargs) -> torch.Tensor:
        return torch.ones(shape, **kwargs)

    def zeros(self, shape: Tuple[int, ...], **kwargs) -> torch.Tensor:
        return torch.zeros(shape, **kwargs)

    def numel(self, tensor: torch.Tensor) -> int:
        return tensor.numel()

    def dtype(self, tensor: torch.Tensor) -> torch.dtype:
        return tensor.dtype

    def astype(self, tensor: torch.Tensor, dtype: Any) -> torch.Tensor:
        torch_dtype = self.type(dtype)
        return tensor.to(dtype=torch_dtype)

    def ndim(self, tensor: torch.Tensor) -> int:
        return tensor.ndim

    def squeeze(self, tensor: torch.Tensor, axis: Optional[int] = None) -> torch.Tensor:
        if axis is None:
            return torch.squeeze(tensor)
        return torch.squeeze(tensor, dim=axis)

    def unsqueeze(self, tensor: torch.Tensor, axis: int) -> torch.Tensor:
        return torch.unsqueeze(tensor, dim=axis)

    def max(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> torch.Tensor:
        if axis is None:
            return torch.max(tensor)
        return torch.max(tensor, dim=axis, keepdim=keepdims)[0]  # pyright: ignore[reportArgumentType,reportCallIssue]

    def min(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> torch.Tensor:
        if axis is None:
            return torch.min(tensor)
        return torch.min(tensor, dim=axis, keepdim=keepdims)[0]  # pyright: ignore[reportArgumentType,reportCallIssue]

    def argmax(
        self, tensor: torch.Tensor, axis: Optional[Union[int, Tuple[int, ...]]] = None
    ) -> torch.Tensor:
        if axis is None:
            return torch.argmax(tensor)
        return torch.argmax(tensor, dim=axis)  # pyright: ignore[reportArgumentType]

    def argmin(
        self, tensor: torch.Tensor, axis: Optional[Union[int, Tuple[int, ...]]] = None
    ) -> torch.Tensor:
        if axis is None:
            return torch.argmin(tensor)
        return torch.argmin(tensor, dim=axis)  # pyright: ignore[reportArgumentType]

    def abs(self, tensor: torch.Tensor) -> torch.Tensor:
        return torch.abs(tensor)

    def sqrt(self, tensor: torch.Tensor) -> torch.Tensor:
        return torch.sqrt(tensor)

    def exp(self, tensor: torch.Tensor) -> torch.Tensor:
        return torch.exp(tensor)

    def log(self, tensor: torch.Tensor) -> torch.Tensor:
        return torch.log(tensor)

    def pow(
        self, tensor: torch.Tensor, exponent: Union[int, float, torch.Tensor]
    ) -> torch.Tensor:
        return torch.pow(tensor, exponent)

    def transpose(self, tensor: torch.Tensor, axes: Tuple[int, ...]) -> torch.Tensor:
        return tensor.permute(dims=axes)

    def repeat(
        self,
        tensor: torch.Tensor,
        repeats: Union[int, torch.Tensor],
        axis: Optional[int] = None,
    ) -> torch.Tensor:
        return torch.repeat_interleave(tensor, repeats, dim=axis)

    def tile(
        self, tensor: torch.Tensor, reps: Union[int, Tuple[int, ...]]
    ) -> torch.Tensor:
        if isinstance(reps, int):
            # 如果是单个整数，需要转换为适当的元组
            # PyTorch 的 repeat 需要每个维度都指定重复次数
            reps = (reps,) * tensor.ndim
        return tensor.repeat(*reps)

    def stack(
        self, tensors: typing.Sequence[torch.Tensor], axis: int = 0
    ) -> torch.Tensor:
        return torch.stack(list(tensors), dim=axis)

    def full(self, shape: Shape, fill_value: Union[int, float], **kwargs) -> torch.Tensor:
        return torch.full(shape, fill_value, **kwargs)

    def rand(self, shape: Shape, **kwargs) -> torch.Tensor:
        device = kwargs.pop('device', None)
        tensor = torch.rand(*shape, **kwargs)
        if device is not None:
            tensor = tensor.to(device)
        return tensor

    def randn(self, shape: Shape, **kwargs) -> torch.Tensor:
        device = kwargs.pop('device', None)
        tensor = torch.randn(*shape, **kwargs)
        if device is not None:
            tensor = tensor.to(device)
        return tensor

    def clip(
        self, tensor: torch.Tensor, min_val: Union[int, float], max_val: Union[int, float]
    ) -> torch.Tensor:
        return torch.clamp(tensor, min=min_val, max=max_val)

    def where(
        self, condition: torch.Tensor, x: torch.Tensor, y: torch.Tensor
    ) -> torch.Tensor:
        return torch.where(condition, x, y)

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

    def all(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> torch.Tensor:
        if axis is None:
            return torch.all(tensor)
        return torch.all(tensor, dim=axis, keepdim=keepdims)

    def any(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> torch.Tensor:
        if axis is None:
            return torch.any(tensor)
        return torch.any(tensor, dim=axis, keepdim=keepdims)

    def zeros_like(self, tensor: torch.Tensor, **kwargs) -> torch.Tensor:
        return torch.zeros_like(tensor, **kwargs)

    def ones_like(self, tensor: torch.Tensor, **kwargs) -> torch.Tensor:
        return torch.ones_like(tensor, **kwargs)

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

        Parameters:
        -----------
        method : str
            - 'vectorized': 使用torch的默认范数计算
            - 'matrix': 模拟矩阵诱导范数计算
            - 'auto': 对torch使用'vectorized'方法
        """
        # 参数验证
        if method not in ['vectorized', 'matrix', 'auto']:
            raise ValueError(
                f"method must be one of ['vectorized', 'matrix', 'auto'], got {method}"
            )

        # 自动选择方法
        if method == 'auto':
            method = 'vectorized'  # torch默认使用向量化范数

        if axis is None:
            if method == 'vectorized':
                # 使用torch的默认行为
                return torch.norm(tensor, p=ord)
            elif method == 'matrix':
                # 模拟矩阵诱导范数
                if ord == 1:
                    # L1矩阵范数：最大列和
                    abs_tensor = torch.abs(tensor)
                    col_sums = torch.sum(abs_tensor, dim=0)
                    return torch.max(col_sums)
                elif ord == float('inf'):
                    # L∞矩阵范数：最大行和
                    abs_tensor = torch.abs(tensor)
                    row_sums = torch.sum(abs_tensor, dim=1)
                    return torch.max(row_sums)
                elif ord == 2 or ord is None:
                    # L2矩阵范数：最大奇异值
                    # return torch.norm(tensor, p='fro')  # Frobenius范数对2×2矩阵等于谱范数
                    if tensor.ndim != 2:
                        raise ValueError(
                            'Matrix norm can only be computed for 2D tensors.'
                        )
                    return torch.linalg.norm(tensor, ord=2)
                else:
                    # 其他情况回退到向量化方法
                    return torch.norm(tensor, p=ord)
        else:
            # 指定了axis时，两种方法行为相同
            if isinstance(axis, tuple):
                dim = axis
            else:
                dim = axis
            return torch.norm(tensor, p=ord, dim=dim, keepdim=keepdims)

        raise AssertionError('unreachable')

    def std(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, tuple]] = None,
        keepdims: bool = False,
        ddof: int = 0,
    ) -> torch.Tensor:
        # Handle different ddof values
        if ddof == 1:
            unbiased = True
        elif ddof == 0:
            unbiased = False
        else:
            # For general ddof, calculate from var
            return self.var(tensor, axis=axis, keepdims=keepdims, ddof=ddof).sqrt()

        if axis is None:
            return torch.std(tensor, unbiased=unbiased)

        dim = axis if isinstance(axis, (tuple, list)) else (axis,)
        return torch.std(tensor, dim=dim, keepdim=keepdims, unbiased=unbiased)

    def var(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, tuple]] = None,
        keepdims: bool = False,
        ddof: int = 0,
    ) -> torch.Tensor:
        if ddof == 1:
            unbiased = True
        elif ddof == 0:
            unbiased = False
        else:
            # For general ddof values, we need a manual calculation
            if axis is None:
                numel = tensor.numel()
                if numel - ddof <= 0:
                    return torch.full_like(tensor, float('nan'))
                # Calculate unbiased variance (ddof=0) and then correct it
                return torch.var(tensor, unbiased=False) * (numel / (numel - ddof))

            dim = axis if isinstance(axis, (tuple, list)) else (axis,)
            # Get number of elements along the axis
            n = torch.tensor(tensor.shape)[list(dim)].prod()
            if n - ddof <= 0:
                # Handle case where reduction results in division by zero or negative
                output_shape = list(tensor.shape)
                for d in sorted(dim, reverse=True):
                    if not keepdims:
                        del output_shape[d]
                    else:
                        output_shape[d] = 1
                return torch.full(
                    output_shape, float('nan'), device=tensor.device, dtype=tensor.dtype
                )

            # Calculate unbiased variance (ddof=0) and then correct it
            unbiased_var = torch.var(tensor, dim=dim, keepdim=keepdims, unbiased=False)
            return unbiased_var * (n / (n - ddof))

        if axis is None:
            return torch.var(tensor, unbiased=unbiased)

        dim = axis if isinstance(axis, (tuple, list)) else (axis,)
        return torch.var(tensor, dim=dim, keepdim=keepdims, unbiased=unbiased)

    def prod(
        self,
        tensor: torch.Tensor,
        axis: Optional[Union[int, Tuple[int, ...]]] = None,
        keepdims: bool = False,
    ) -> torch.Tensor:
        if axis is None:
            return torch.prod(tensor)
        return torch.prod(tensor, dim=axis, keepdim=keepdims)  # pyright: ignore[reportArgumentType,reportCallIssue]

    def tensor_equal(self, tensor1: torch.Tensor, tensor2: torch.Tensor) -> bool:
        return torch.equal(tensor1, tensor2)

    def arange(
        self,
        start: Union[int, float],
        stop: Optional[Union[int, float]] = None,
        step: Union[int, float] = 1,
        **kwargs,
    ) -> torch.Tensor:
        kwargs = filter_dict(kwargs, allowed=['dtype', 'device', 'out', 'requires_grad'])
        if 'dtype' in kwargs:
            kwargs['dtype'] = self.type(kwargs['dtype'])

        device = kwargs.pop('device', None)
        if stop is None:
            tensor = torch.arange(start, **kwargs)  # 单参数版本
        else:
            tensor = torch.arange(start, stop, step, **kwargs)
        if device is not None:
            tensor = tensor.to(device)
        return tensor

    def linspace(
        self, start: Union[int, float], stop: Union[int, float], num: int = 50, **kwargs
    ) -> torch.Tensor:
        device = kwargs.pop('device', None)
        tensor = torch.linspace(start, stop, steps=num, **kwargs)
        if device is not None:
            tensor = tensor.to(device)
        return tensor

    def eye(self, n: int, m: Optional[int] = None, **kwargs) -> torch.Tensor:
        kwargs = filter_dict(
            kwargs,
            allowed=['dtype', 'device', 'layout', 'out', 'pin_memory', 'requires_grad'],
        )
        if m is None:
            m = n
        tensor = torch.eye(n, m=m, **kwargs)
        return tensor

    def maximum(
        self, tensor1: torch.Tensor, tensor2: Union[torch.Tensor, float]
    ) -> torch.Tensor:
        if isinstance(tensor2, (float, int)):
            return torch.clamp(tensor1, min=tensor2)
        return torch.maximum(tensor1, tensor2)

    def to_numpy(self, tensor: torch.Tensor) -> np.ndarray:
        return tensor.cpu().numpy()

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

    def unique(
        self,
        tensor: torch.Tensor,
        return_index: bool = False,
        return_inverse: bool = False,
        return_counts: bool = False,
        axis: Optional[int] = None,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
        # 关键点：处理不被支持的参数
        if return_index:
            # 这是一个设计决策：是抛出错误还是尝试模拟？
            # 抛出错误更安全，因为模拟可能效率低下或有边缘情况。
            raise NotImplementedError(
                'torch.unique does not support `return_index` directly.'
            )

        # 参数名转换: axis -> dim
        dim = axis

        # PyTorch 的 unique 默认返回排序后的结果，与 NumPy 行为一致
        result = torch.unique(
            tensor,
            sorted=True,
            return_inverse=return_inverse,
            return_counts=return_counts,
            dim=dim,
        )

        if not any([return_inverse, return_counts]):
            return result

        return result

    def quantile(
        self,
        tensor: torch.Tensor,
        q: Union[float, torch.Tensor],
        axis: Optional[int] = None,
        keepdims: bool = False,
        interpolation: str = 'linear',
    ) -> torch.Tensor:
        # 确保输入张量是浮点数类型
        float_tensor = (
            tensor.to(torch.float32) if not tensor.is_floating_point() else tensor.float()
        )

        # 确保 q 张量的数据类型与输入张量一致
        if not isinstance(q, torch.Tensor):
            q = torch.tensor(q, device=tensor.device, dtype=float_tensor.dtype)
        else:
            q = q.to(dtype=float_tensor.dtype)

        # 调用 torch.quantile
        return torch.quantile(  # pyright: ignore[reportCallIssue]
            float_tensor, q, dim=axis, keepdims=keepdims, interpolation=interpolation
        )

    def searchsorted(
        self,
        sorted_sequence: torch.Tensor,
        values: torch.Tensor,
        right: bool = False,
        side: Literal['left', 'right'] | None = None,
    ) -> torch.Tensor:
        # 统一接口，PyTorch 使用 'right'
        if side is not None:
            right = side == 'right'
        return torch.searchsorted(sorted_sequence, values, right=right)
