#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Modules (classes) interface."""

from typing import TYPE_CHECKING

from torch.nn.modules import (
    CELU,
    ELU,
    GELU,
    GLU,
    GRU,
    LSTM,
    RNN,
    SELU,
    AdaptiveAvgPool1d,
    AdaptiveAvgPool2d,
    AdaptiveAvgPool3d,
    AdaptiveLogSoftmaxWithLoss,
    AdaptiveMaxPool1d,
    AdaptiveMaxPool2d,
    AdaptiveMaxPool3d,
    AlphaDropout,
    AvgPool1d,
    AvgPool2d,
    AvgPool3d,
    BatchNorm1d,
    BatchNorm2d,
    BatchNorm3d,
    BCELoss,
    BCEWithLogitsLoss,
    Bilinear,
    ChannelShuffle,
    ConstantPad1d,
    ConstantPad2d,
    ConstantPad3d,
    Conv1d,
    Conv2d,
    Conv3d,
    ConvTranspose1d,
    ConvTranspose2d,
    ConvTranspose3d,
    CosineEmbeddingLoss,
    CosineSimilarity,
    CrossEntropyLoss,
    CrossMapLRN2d,
    CTCLoss,
    Dropout,
    Dropout2d,
    Dropout3d,
    Embedding,
    EmbeddingBag,
    FeatureAlphaDropout,
    Flatten,
    Fold,
    FractionalMaxPool2d,
    FractionalMaxPool3d,
    GaussianNLLLoss,
    GroupNorm,
    GRUCell,
    Hardshrink,
    Hardsigmoid,
    Hardswish,
    Hardtanh,
    HingeEmbeddingLoss,
    HuberLoss,
    Identity,
    InstanceNorm1d,
    InstanceNorm2d,
    InstanceNorm3d,
    KLDivLoss,
    L1Loss,
    LayerNorm,
    LazyBatchNorm1d,
    LazyBatchNorm2d,
    LazyBatchNorm3d,
    LazyConv1d,
    LazyConv2d,
    LazyConv3d,
    LazyConvTranspose1d,
    LazyConvTranspose2d,
    LazyConvTranspose3d,
    LazyInstanceNorm1d,
    LazyInstanceNorm2d,
    LazyInstanceNorm3d,
    LazyLinear,
    LeakyReLU,
    Linear,
    LocalResponseNorm,
    LogSigmoid,
    LogSoftmax,
    LPPool1d,
    LPPool2d,
    LSTMCell,
    MarginRankingLoss,
    MaxPool1d,
    MaxPool2d,
    MaxPool3d,
    MaxUnpool1d,
    MaxUnpool2d,
    MaxUnpool3d,
    Mish,
    Module,
    ModuleDict,
    ModuleList,
    MSELoss,
    MultiheadAttention,
    MultiLabelMarginLoss,
    MultiLabelSoftMarginLoss,
    MultiMarginLoss,
    NLLLoss,
    PairwiseDistance,
    ParameterDict,
    ParameterList,
    PixelShuffle,
    PixelUnshuffle,
    PoissonNLLLoss,
    PReLU,
    ReflectionPad1d,
    ReflectionPad2d,
    ReflectionPad3d,
    ReLU,
    ReLU6,
    ReplicationPad1d,
    ReplicationPad2d,
    ReplicationPad3d,
    RNNBase,
    RNNCell,
    RNNCellBase,
    RReLU,
    Sequential,
    Sigmoid,
    SiLU,
    SmoothL1Loss,
    SoftMarginLoss,
    Softmax,
    Softmax2d,
    Softmin,
    Softplus,
    Softshrink,
    Softsign,
    SyncBatchNorm,
    Tanh,
    Tanhshrink,
    Threshold,
    Transformer,
    TransformerDecoder,
    TransformerDecoderLayer,
    TransformerEncoder,
    TransformerEncoderLayer,
    TripletMarginLoss,
    TripletMarginWithDistanceLoss,
    Unflatten,
    Unfold,
    Upsample,
    UpsamplingBilinear2d,
    UpsamplingNearest2d,
    ZeroPad2d,
)

if TYPE_CHECKING:
    from .activation import LogSoftmaxMultidim, SoftmaxMultidim
    from .container import (
        EModuleDict,
        EModuleList,
        EModulePartial,
        ESequential,
        ModuleDict,
        ModuleList,
        ModulePartial,
        Sequential,
    )
    from .crop import CropDim, CropDims
    from .layer import PositionalEncoding
    from .mask import MaskedMean, MaskedSum
    from .module import EModule, Module
    from .multiclass import (
        IndexToName,
        IndexToOnehot,
        NameToIndex,
        NameToOnehot,
        OnehotToIndex,
        OnehotToName,
        ProbsToIndex,
        ProbsToName,
        ProbsToOnehot,
    )
    from .multilabel import (
        IndicesToMultihot,
        IndicesToMultinames,
        MultihotToIndices,
        MultihotToMultiIndices,
        MultihotToMultinames,
        MultiIndicesToMultihot,
        MultiIndicesToMultinames,
        MultinamesToIndices,
        MultinamesToMultihot,
        MultinamesToMultiIndices,
        ProbsToIndices,
        ProbsToMultihot,
        ProbsToMultiIndices,
        ProbsToMultinames,
    )
    from .numpy import NDArrayToTensor, TensorToNDArray, ToNDArray
    from .padding import PadAndStackRec, PadDim, PadDims
    from .powerset import MultilabelToPowerset, PowersetToMultilabel
    from .tensor import (
        FFT,
        IFFT,
        Abs,
        Angle,
        Exp,
        Exp2,
        Imag,
        Log,
        Log2,
        Log10,
        Max,
        Mean,
        Min,
        Normalize,
        Permute,
        Pow,
        Real,
        Repeat,
        RepeatInterleave,
        Reshape,
        Sort,
        TensorTo,
        ToList,
        Transpose,
        View,
    )
    from .transform import (
        AsTensor,
        Identity,
        MoveToRec,
        PadAndCropDim,
        RepeatInterleaveNd,
        ResampleNearestFreqs,
        ResampleNearestRates,
        ResampleNearestSteps,
        Shuffled,
        Squeeze,
        TFlatten,
        ToItem,
        Topk,
        TopP,
        ToTensor,
        TransformDrop,
        Unsqueeze,
        ViewAsComplex,
        ViewAsReal,
    )
else:
    import lazy_loader as lazy

    __getattr__, __dir__, __all__ = lazy.attach(
        __name__,
        submodules=[
            "activation",
            "container",
            "crop",
            "layer",
            "mask",
            "module",
            "multiclass",
            "multilabel",
            "numpy",
            "padding",
            "powerset",
            "tensor",
            "transform",
        ],
        submod_attrs={
            "activation": ["LogSoftmaxMultidim", "SoftmaxMultidim"],
            "container": [
                "EModuleDict",
                "EModuleList",
                "EModulePartial",
                "ESequential",
                "ModuleDict",
                "ModuleList",
                "ModulePartial",
                "Sequential",
            ],
            "crop": ["CropDim", "CropDims"],
            "layer": ["PositionalEncoding"],
            "mask": ["MaskedMean", "MaskedSum"],
            "module": ["EModule", "Module"],
            "multiclass": [
                "IndexToName",
                "IndexToOnehot",
                "NameToIndex",
                "NameToOnehot",
                "OnehotToIndex",
                "OnehotToName",
                "ProbsToIndex",
                "ProbsToName",
                "ProbsToOnehot",
            ],
            "multilabel": [
                "IndicesToMultihot",
                "IndicesToMultinames",
                "MultihotToIndices",
                "MultihotToMultiIndices",
                "MultihotToMultinames",
                "MultiIndicesToMultihot",
                "MultiIndicesToMultinames",
                "MultinamesToIndices",
                "MultinamesToMultihot",
                "MultinamesToMultiIndices",
                "ProbsToIndices",
                "ProbsToMultihot",
                "ProbsToMultiIndices",
                "ProbsToMultinames",
            ],
            "numpy": ["NDArrayToTensor", "TensorToNDArray", "ToNDArray"],
            "padding": ["PadAndStackRec", "PadDim", "PadDims"],
            "powerset": ["MultilabelToPowerset", "PowersetToMultilabel"],
            "tensor": [
                "FFT",
                "IFFT",
                "Abs",
                "Angle",
                "Exp",
                "Exp2",
                "Imag",
                "Log",
                "Log2",
                "Log10",
                "Max",
                "Mean",
                "Min",
                "Normalize",
                "Permute",
                "Pow",
                "Real",
                "Repeat",
                "RepeatInterleave",
                "Reshape",
                "Sort",
                "TensorTo",
                "ToList",
                "Transpose",
                "View",
            ],
            "transform": [
                "AsTensor",
                "TFlatten",
                "Identity",
                "MoveToRec",
                "PadAndCropDim",
                "RepeatInterleaveNd",
                "ResampleNearestFreqs",
                "ResampleNearestRates",
                "ResampleNearestSteps",
                "Shuffled",
                "Squeeze",
                "ToItem",
                "Topk",
                "TopP",
                "ToTensor",
                "TransformDrop",
                "Unsqueeze",
                "ViewAsComplex",
                "ViewAsReal",
            ],
        },
    )


del TYPE_CHECKING
