from typing import Callable, Literal, Union

from torch import nn
from torch.nn import functional as F

XDim = Literal['1d', '2d', '3d', 1, 2, 3]
Conv = type[nn.Conv1d] | type[nn.Conv2d] | type[nn.Conv3d]
BatchNorm = type[nn.BatchNorm1d] | type[nn.BatchNorm2d] | type[nn.BatchNorm3d]
MaxPool = type[nn.MaxPool1d] | type[nn.MaxPool2d] | type[nn.MaxPool3d]
AvgPool = type[nn.AvgPool1d] | type[nn.AvgPool2d] | type[nn.AvgPool3d]
AdaptiveAvgPool = (
    type[nn.AdaptiveAvgPool1d] | type[nn.AdaptiveAvgPool2d] | type[nn.AdaptiveAvgPool3d]
)
AdaptiveMaxPool = (
    type[nn.AdaptiveMaxPool1d] | type[nn.AdaptiveMaxPool2d] | type[nn.AdaptiveMaxPool3d]
)
ConvTranspose = (
    type[nn.ConvTranspose1d] | type[nn.ConvTranspose2d] | type[nn.ConvTranspose3d]
)
Dropout = type[nn.Dropout] | type[nn.Dropout2d] | type[nn.Dropout3d]


def get_activation(
    name: Literal[
        'relu', 'leaky_relu', 'prelu', 'gelu', 'silu', 'tanh', 'sigmoid', 'swish'
    ]
    | None = None,
    return_type: Literal['class', 'instance', 'function'] = 'class',
) -> nn.Module | Callable:
    if name is None:
        return lambda x: x
    activation_class = {
        'relu': (nn.ReLU, F.relu),
        'leaky_relu': (nn.LeakyReLU, F.leaky_relu),
        'prelu': (nn.PReLU, F.prelu),
        'gelu': (nn.GELU, F.gelu),
        'silu': (nn.SiLU, F.silu),
        'tanh': (nn.Tanh, F.tanh),
        'sigmoid': (nn.Sigmoid, F.sigmoid),
        'swish': (nn.SiLU, F.silu),
    }
    activation = activation_class.get(name.lower())
    if activation is None:
        raise ValueError(f'Invalid activation name: {name}')

    if return_type == 'instance':
        return activation[0]()
    elif return_type == 'function':
        return activation[1]
    elif return_type == 'class':
        return activation[0]
    else:
        raise ValueError(f'Invalid return type: {return_type}')


def get_conv(xdim: XDim = 1) -> Conv:
    """获取对应维度的卷积层类

    根据指定的维度返回相应的 PyTorch 卷积层类。

    Parameters
    ----------
    - `xdim` : `XDim`, optional
        - 维度规格，支持字符串 ('1d', '2d', '3d') 或整数 (1, 2, 3), by default 1

    Returns
    ----------
    - `Conv`
        - 对应维度的卷积层类：
          - 1D: nn.Conv1d (用于时序信号，如 EEG、音频)
          - 2D: nn.Conv2d (用于图像、时频图)
          - 3D: nn.Conv3d (用于视频、3D 数据)

    Raises
    ----------
    - `ValueError`
        - 当 xdim 参数无效时抛出

    Examples
    ----------
    ```py
    # 创建 1D 卷积用于 EEG 信号
    Conv1d = conv(1)  # 或 conv('1d')
    eeg_conv = Conv1d(64, 128, kernel_size=3, padding=1)

    # 创建 2D 卷积用于图像
    Conv2d = conv(2)  # 或 conv('2d')
    image_conv = Conv2d(3, 64, kernel_size=3, padding=1)

    # 创建 3D 卷积用于视频
    Conv3d = conv(3)  # 或 conv('3d')
    video_conv = Conv3d(3, 64, kernel_size=3, padding=1)
    ```
    """
    if xdim == '1d':
        return nn.Conv1d
    elif xdim == '2d':
        return nn.Conv2d
    elif xdim == '3d':
        return nn.Conv3d
    elif xdim == 1:
        return nn.Conv1d
    elif xdim == 2:
        return nn.Conv2d
    elif xdim == 3:
        return nn.Conv3d
    else:
        raise ValueError(f'Invalid parameter xdim: {xdim}; 无法转换为对应的 Conv 类')


def get_batchnorm(xdim: XDim = 1) -> BatchNorm:
    """获取对应维度的批归一化层类

    根据指定的维度返回相应的 PyTorch 批归一化层类。

    Parameters
    ----------
    - `xdim` : `XDim`, optional
        - 维度规格，支持字符串 ('1d', '2d', '3d') 或整数 (1, 2, 3), by default 1

    Returns
    ----------
    - `nn.Module`
        - 对应维度的批归一化层类：
          - 1D: nn.BatchNorm1d (用于时序数据，输入形状 [N, C, T])
          - 2D: nn.BatchNorm2d (用于图像数据，输入形状 [N, C, H, W])
          - 3D: nn.BatchNorm3d (用于3D数据，输入形状 [N, C, D, H, W])

    Raises
    ----------
    - `ValueError`
        - 当 xdim 参数无效时抛出

    Examples
    ----------
    ```py
    # 创建 1D 批归一化用于 EEG 特征
    BatchNorm1d = batchnorm(1)
    eeg_bn = BatchNorm1d(64)  # 64 个通道

    # 创建 2D 批归一化用于图像特征
    BatchNorm2d = batchnorm('2d')
    image_bn = BatchNorm2d(256)  # 256 个特征图
    ```
    """
    if xdim == '1d':
        return nn.BatchNorm1d
    elif xdim == '2d':
        return nn.BatchNorm2d
    elif xdim == '3d':
        return nn.BatchNorm3d
    elif xdim == 1:
        return nn.BatchNorm1d
    elif xdim == 2:
        return nn.BatchNorm2d
    elif xdim == 3:
        return nn.BatchNorm3d
    else:
        raise ValueError(f'Invalid parameter xdim: {xdim}; 无法转换为对应的 BatchNorm 类')


def maxpool(xdim: XDim = 1) -> MaxPool:
    if xdim == '1d':
        return nn.MaxPool1d
    elif xdim == '2d':
        return nn.MaxPool2d
    elif xdim == '3d':
        return nn.MaxPool3d
    elif xdim == 1:
        return nn.MaxPool1d
    elif xdim == 2:
        return nn.MaxPool2d
    elif xdim == 3:
        return nn.MaxPool3d
    else:
        raise ValueError(f'Invalid parameter xdim: {xdim}; 无法转换为对应的 MaxPool 类')


def avgpool(xdim: XDim = 1) -> AvgPool:
    if xdim == '1d':
        return nn.AvgPool1d
    elif xdim == '2d':
        return nn.AvgPool2d
    elif xdim == '3d':
        return nn.AvgPool3d
    elif xdim == 1:
        return nn.AvgPool1d
    elif xdim == 2:
        return nn.AvgPool2d
    elif xdim == 3:
        return nn.AvgPool3d
    else:
        raise ValueError(f'Invalid parameter xdim: {xdim}; 无法转换为对应的 AvgPool 类')


def adaptiveavgpool(xdim: XDim = 1) -> AdaptiveAvgPool:
    if xdim == '1d':
        return nn.AdaptiveAvgPool1d
    elif xdim == '2d':
        return nn.AdaptiveAvgPool2d
    elif xdim == '3d':
        return nn.AdaptiveAvgPool3d
    elif xdim == 1:
        return nn.AdaptiveAvgPool1d
    elif xdim == 2:
        return nn.AdaptiveAvgPool2d
    elif xdim == 3:
        return nn.AdaptiveAvgPool3d
    else:
        raise ValueError(
            f'Invalid parameter xdim: {xdim}; 无法转换为对应的 AdaptiveAvgPool 类'
        )


def adaptivemaxpool(xdim: XDim = 1) -> AdaptiveMaxPool:
    if xdim == '1d':
        return nn.AdaptiveMaxPool1d
    elif xdim == '2d':
        return nn.AdaptiveMaxPool2d
    elif xdim == '3d':
        return nn.AdaptiveMaxPool3d
    elif xdim == 1:
        return nn.AdaptiveMaxPool1d
    elif xdim == 2:
        return nn.AdaptiveMaxPool2d
    elif xdim == 3:
        return nn.AdaptiveMaxPool3d
    else:
        raise ValueError(
            f'Invalid parameter xdim: {xdim}; 无法转换为对应的 AdaptiveMaxPool 类'
        )


def convtranspose(xdim: XDim = 1) -> ConvTranspose:
    if xdim == '1d':
        return nn.ConvTranspose1d
    elif xdim == '2d':
        return nn.ConvTranspose2d
    elif xdim == '3d':
        return nn.ConvTranspose3d
    elif xdim == 1:
        return nn.ConvTranspose1d
    elif xdim == 2:
        return nn.ConvTranspose2d
    elif xdim == 3:
        return nn.ConvTranspose3d
    else:
        raise ValueError(
            f'Invalid parameter xdim: {xdim}; 无法转换为对应的 ConvTranspose 类'
        )


def get_dropout(xdim: XDim = 1) -> Dropout:
    if xdim == '1d':
        return nn.Dropout
    elif xdim == '2d':
        return nn.Dropout2d
    elif xdim == '3d':
        return nn.Dropout3d
    elif xdim == 1:
        return nn.Dropout
    elif xdim == 2:
        return nn.Dropout2d
    elif xdim == 3:
        return nn.Dropout3d
    else:
        raise ValueError(f'Invalid parameter xdim: {xdim}; 无法转换为对应的 Dropout 类')


def get_pooling(
    name: Literal['max', 'avg', 'adaptive_max', 'adaptive_avg'], xdim: XDim = 1
) -> MaxPool | AvgPool | AdaptiveMaxPool | AdaptiveAvgPool:
    if name == 'max':
        return maxpool(xdim)
    elif name == 'avg':
        return avgpool(xdim)
    elif name == 'adaptive_max':
        return adaptivemaxpool(xdim)
    elif name == 'adaptive_avg':
        return adaptiveavgpool(xdim)
    else:
        raise ValueError(f'Invalid pooling name: {name}')
