from __future__ import annotations

import typing
from typing import TYPE_CHECKING, Literal, Optional, TypeAlias, overload

from zyplib.import_tools import import_lazy

if TYPE_CHECKING:
    import numpy as np
    import torch
    from numpy import dtype as np_dtype
    from numpy import ndarray
    from torch import Tensor
    from torch import dtype as torch_dtype
    from torch.types import Device
else:
    np = import_lazy('numpy')
    torch = import_lazy('torch')

Scalar: TypeAlias = int | float | complex
TensorOrArray: TypeAlias = 'Tensor | ndarray'

__all__ = ['ensure_3dims', 'ensure_2dims', 'ensure_npy', 'ensure_tensor', 'ensure_scalar']


def numel(data: TensorOrArray) -> int:
    """计算张量或数组中元素的总数

    Parameters
    ----------
    - `data` : `torch.Tensor | np.ndarray`
        - 输入的张量或数组

    Returns
    ----------
    - `int`
        - 元素总数

    Raises
    ----------
    - `TypeError`
        - 当输入类型不是 torch.Tensor 或 np.ndarray 时抛出
    """
    if isinstance(data, torch.Tensor):
        return data.numel()
    elif isinstance(data, np.ndarray):
        return data.size
    else:
        raise TypeError(
            f'data should be torch.Tensor or np.ndarray, but got {type(data)}'
        )


def ensure_scalar(
    data: Scalar | TensorOrArray | typing.Sequence[typing.Any],
) -> Scalar | None:
    """确保数据是标量

    尝试将输入数据转换为标量值。

    Parameters
    ----------
    - `data` : `int | float | complex | torch.Tensor | np.ndarray | list`
        - 输入数据，可以是标量、张量、数组或列表

    Returns
    ----------
    - `int | float | complex | None`
        - 如果可以转换为标量，返回标量值
        - 如果无法转换为标量，返回 None
    """
    if isinstance(data, (int, float, complex)):
        return data
    elif isinstance(data, (torch.Tensor, np.ndarray)):
        data = data.squeeze()
        if numel(data) == 1:
            if isinstance(data, torch.Tensor):
                return data.cpu().item()
            else:
                return data.item()
    elif isinstance(data, (list, tuple)):
        if len(data) == 1:
            return data[0]
    return None


@overload
def ensure_npy(data: 'ndarray', dtype: Optional['np_dtype'] = None) -> 'ndarray': ...


@overload
def ensure_npy(data: 'Tensor', dtype: Optional['np_dtype'] = None) -> 'ndarray': ...


@overload
def ensure_npy(
    data: list[typing.Any], dtype: Optional['np_dtype'] = None
) -> 'ndarray': ...


def ensure_npy(
    data: 'ndarray | Tensor | list[typing.Any]', dtype: Optional['np_dtype'] = None
) -> 'ndarray':
    """确保数据是 NumPy 数组格式

    将输入数据转换为 NumPy 数组格式。

    Parameters
    ----------
    - `data` : `Union[np.ndarray, torch.Tensor, list]`
        - 输入数据，可以是 NumPy 数组、PyTorch 张量或列表
    - `dtype` : `np.dtype`, optional
        - 目标数据类型, by default None

    Returns
    ----------
    - `np.ndarray`
        - 转换后的 NumPy 数组

    Raises
    ----------
    - `TypeError`
        - 当输入类型不支持转换时抛出
    """
    if isinstance(data, np.ndarray):
        ...
    elif isinstance(data, torch.Tensor):
        data = data.detach().cpu().numpy()
    elif isinstance(data, list):
        data = np.array(data)
    else:
        raise TypeError(
            f'data should be np.ndarray, torch.Tensor, or list, but got {type(data)}'
        )
    if dtype is not None:
        data = data.astype(dtype)
    return data


ensure_ndarray = ensure_npy


@overload
def ensure_tensor(
    data: 'ndarray',
    dtype: Optional['torch_dtype'] = None,
    detach: bool = False,
    device: 'str | Device | None' = None,
) -> 'Tensor': ...


@overload
def ensure_tensor(
    data: 'Tensor',
    dtype: Optional['torch_dtype'] = None,
    detach: bool = False,
    device: 'str | Device | None' = None,
) -> 'Tensor': ...


@overload
def ensure_tensor(
    data: list[typing.Any],
    dtype: Optional['torch_dtype'] = None,
    detach: bool = False,
    device: 'str | Device | None' = None,
) -> 'Tensor': ...


def ensure_tensor(
    data: 'ndarray | Tensor | list[typing.Any]',
    dtype: Optional['torch_dtype'] = None,
    detach: bool = False,
    device: 'str | Device | None' = None,
) -> 'Tensor':
    """确保数据是 PyTorch 张量格式

    将输入数据转换为 PyTorch 张量格式。

    Parameters
    ----------
    - `data` : `Union[np.ndarray, torch.Tensor, list]`
        - 输入数据，可以是 NumPy 数组、PyTorch 张量或列表
    - `dtype` : `torch.dtype`, optional
        - 目标数据类型, by default None
    - `detach` : `bool`, optional
        - 是否分离计算图, by default False
    - `device` : `Union[str, torch.device]`, optional
        - 目标设备, by default None

    Returns
    ----------
    - `torch.Tensor`
        - 转换后的 PyTorch 张量

    Raises
    ----------
    - `TypeError`
        - 当输入类型不支持转换时抛出
    """
    if isinstance(data, np.ndarray):
        # 检查是否有负步长，如果有则创建副本
        if any(stride < 0 for stride in data.strides):
            data = data.copy()
        data = torch.from_numpy(data)
    elif isinstance(data, torch.Tensor):
        ...
    elif isinstance(data, list):
        data = torch.tensor(data)
    else:
        raise TypeError(
            f'data should be np.ndarray, torch.Tensor, or list, but got {type(data)}'
        )
    if dtype is not None:
        data = data.type(dtype)
    if detach:
        data = data.detach()
    if device is not None:
        data = data.to(device)
    return data


@overload
def ensure_3dims(
    signals: 'ndarray', newaxis: Literal['batch', 'channel'] = 'batch'
) -> 'ndarray': ...


@overload
def ensure_3dims(
    signals: 'Tensor', newaxis: Literal['batch', 'channel'] = 'batch'
) -> 'Tensor': ...


def ensure_3dims(
    signals: 'ndarray | Tensor', newaxis: Literal['batch', 'channel'] = 'batch'
) -> 'ndarray | Tensor':
    """把数据转换为标准的 [Batch, Channel, T] 格式

    Parameters
    ----------
    - `signals` : `Union[np.ndarray, torch.Tensor]`
        - 信号张量，形状可以为 [T] 或者 [N, T]
    - `newaxis` : `Literal['batch', 'channel']`, optional
        - 指定在哪个维度上添加新的维度, by default 'batch'
        - 'batch': 在 batch 维度上添加新的维度
        - 'channel': 在 channel 维度上添加新的维度

    Returns
    ----------
    - `Union[np.ndarray, torch.Tensor]`
        - 转换后的三维数据，形状为 [Batch, Channel, T]

    Examples
    ----------
    ```python
    X = np.random.randn(2, 3)
    X1 = ensure_3dims(X, newaxis='channel')
    print(X1.shape)  # [2, 1, 3]
    X2 = ensure_3dims(X, newaxis='batch')
    print(X2.shape)  # [1, 2, 3]
    ```
    """
    # NumPy array handling
    if isinstance(signals, np.ndarray):
        if signals.ndim != 3:
            if signals.ndim == 1:
                signals = signals[np.newaxis, np.newaxis, :]
            elif signals.ndim == 2:
                if newaxis == 'channel':
                    signals = signals[:, np.newaxis, :]
                else:
                    signals = signals[np.newaxis, :, :]
            else:
                raise ValueError(f'signals.ndim should be 1 or 2, but got {signals.ndim}')
        return signals

    # PyTorch tensor handling
    elif isinstance(signals, torch.Tensor):
        if signals.ndim != 3:
            if signals.ndim == 1:
                signals = signals.unsqueeze(0).unsqueeze(0)
            elif signals.ndim == 2:
                if newaxis == 'channel':
                    signals = signals.unsqueeze(1)
                else:
                    signals = signals.unsqueeze(0)
            else:
                raise ValueError(f'signals.ndim should be 1 or 2, but got {signals.ndim}')
        return signals

    else:
        raise TypeError(
            f'signals should be np.ndarray or torch.Tensor, but got {type(signals)}'
        )


@overload
def ensure_2dims(signals: 'ndarray') -> 'ndarray': ...


@overload
def ensure_2dims(signals: 'Tensor') -> 'Tensor': ...


def ensure_2dims(signals: 'ndarray | Tensor') -> 'ndarray | Tensor':
    """把数据转换为标准的 [Batch, T] 格式

    Parameters
    ----------
    - `signals` : `Union[np.ndarray, torch.Tensor]`
        - 信号张量，形状可以为 [T] 或者 [N, T]

    Returns
    ----------
    - `Union[np.ndarray, torch.Tensor]`
        - 转换后的二维数据，形状为 [Batch, T]
    """
    # NumPy array handling
    if isinstance(signals, np.ndarray):
        if signals.ndim == 2:
            return signals
        elif signals.ndim == 1:
            return signals[np.newaxis, :]
        elif signals.ndim == 3:
            squeezed_signals = signals.squeeze()
            if squeezed_signals.ndim == 2:
                return squeezed_signals
            elif squeezed_signals.ndim == 1:
                return squeezed_signals[np.newaxis, :]
            else:
                # For (N, C, T) where N>1 and C>1, the user must decide how to flatten.
                raise ValueError(
                    f'Cannot automatically convert array of shape {signals.shape} to 2D. '
                    'Please reshape manually.'
                )
        else:
            raise ValueError(f'Unsupported array dimension: {signals.ndim}')

    # PyTorch tensor handling
    elif isinstance(signals, torch.Tensor):
        if signals.ndim == 2:
            return signals
        elif signals.ndim == 1:
            return signals.unsqueeze(0)
        elif signals.ndim == 3:
            squeezed_signals = signals.squeeze()
            if squeezed_signals.ndim == 2:
                return squeezed_signals
            elif squeezed_signals.ndim == 1:
                return squeezed_signals.unsqueeze(0)
            else:
                # For (N, C, T) where N>1 and C>1, the user must decide how to flatten.
                raise ValueError(
                    f'Cannot automatically convert tensor of shape {signals.shape} to 2D. '
                    'Please reshape manually.'
                )
        else:
            raise ValueError(f'Unsupported tensor dimension: {signals.ndim}')

    else:
        raise TypeError(
            f'signals should be np.ndarray or torch.Tensor, but got {type(signals)}'
        )
