"""
nn.modules 模块实现了常见的深度学习中的 building block

包含基础模块、卷积层、归一化层、正则化层、注意力机制等。
特别针对脑电信号分析提供了专用的模块实现。
"""

# 基础模块
from ._get import XDim, get_activation, get_batchnorm, get_conv, get_dropout, get_pooling

# 注意力机制
from .attention import (
    CBAM1D,
    ChannelAttention1D,
    EEGChannelAttention,
    SpatialAttention1D,
    SpectralAttention,
    TemporalAttention,
)
from .base import BaseModule, LambdaModule, MaxNormConstraint

# 卷积模块
from .conv import DepthwiseSeperableConv, MaxNormConv, MultiScaleConv1d, SincConv1d

# 归一化模块
from .norm import (
    CategoricalConditionalBatchNorm1d,
    ConditionalBatchNorm1d,
    LayerScale,
    PixelNorm,
    RMSNorm,
    build_norm,
)

# 正则化模块
from .regularization import DropPath, StochasticDepth

# SE 模块
from .senet import SEBlock, SEBlock2D

__all__ = [
    # 基础模块
    'BaseModule',
    'MaxNormConstraint',
    'XDim',
    'get_batchnorm',
    'get_conv',
    'get_activation',
    'get_dropout',
    'get_pooling',
    # 卷积模块
    'DepthwiseSeperableConv',
    'MaxNormConv',
    'MultiScaleConv1d',
    'SincConv1d',
    # 归一化模块
    'build_norm',
    'PixelNorm',
    'ConditionalBatchNorm1d',
    'CategoricalConditionalBatchNorm1d',
    'RMSNorm',
    'LayerScale',
    # 正则化模块
    'DropPath',
    'StochasticDepth',
    # 注意力机制
    'CBAM1D',
    'ChannelAttention1D',
    'SpatialAttention1D',
    'EEGChannelAttention',
    'TemporalAttention',
    'SpectralAttention',
    # SE 模块
    'SEBlock',
    'SEBlock2D',
]
