"""
cukks.nn - Encrypted neural network modules.

This module provides PyTorch-like layers for encrypted inference.
Each layer mirrors its torch.nn counterpart but operates on encrypted data.

Example:
    >>> import torch.nn as nn
    >>> import cukks
    >>> 
    >>> # Train a plaintext model
    >>> model = nn.Sequential(
    ...     nn.Linear(784, 128),
    ...     nn.ReLU(),
    ...     nn.Linear(128, 10)
    ... )
    >>> train(model, data)
    >>> 
    >>> # Convert to encrypted model
    >>> enc_model = cukks.convert(model, ctx)
    >>> 
    >>> # Run encrypted inference
    >>> enc_output = enc_model(enc_input)
"""

from .module import EncryptedModule, EncryptedIdentity
from .linear import EncryptedLinear
from .block_diagonal import BlockDiagonalLinear
from .block_diagonal_low_rank import BlockDiagLowRankLinear
from .encrypted_block_diag_lr import EncryptedBlockDiagLowRank
from .activations import (
    EncryptedSquare,
    EncryptedReLU,
    EncryptedGELU,
    EncryptedSiLU,
    EncryptedSigmoid,
    EncryptedTanh,
)
from .conv import EncryptedConv2d
from .pooling import EncryptedAvgPool2d, EncryptedMaxPool2d
from .flatten import EncryptedFlatten
from .sequential import EncryptedSequential
from .batchnorm import EncryptedBatchNorm1d, EncryptedBatchNorm2d
from .layernorm import EncryptedLayerNorm
from .dropout import EncryptedDropout
from .residual import EncryptedResidualBlock
from .attention import EncryptedApproxAttention

__all__ = [
    "EncryptedModule",
    "EncryptedIdentity",
    "EncryptedLinear",
    "BlockDiagonalLinear",
    "BlockDiagLowRankLinear",
    "EncryptedBlockDiagLowRank",
    "EncryptedSquare",
    "EncryptedReLU",
    "EncryptedGELU",
    "EncryptedSiLU",
    "EncryptedSigmoid",
    "EncryptedTanh",
    "EncryptedConv2d",
    "EncryptedAvgPool2d",
    "EncryptedMaxPool2d",
    "EncryptedFlatten",
    "EncryptedSequential",
    "EncryptedBatchNorm1d",
    "EncryptedBatchNorm2d",
    "EncryptedLayerNorm",
    "EncryptedDropout",
    "EncryptedResidualBlock",
    "EncryptedApproxAttention",
]
