# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/blocks/03_se_3d.ipynb.

# %% auto 0
__all__ = ['SEBlock3DConfig', 'SEBlock3D']

# %% ../../nbs/blocks/03_se_3d.ipynb 2
from functools import wraps

import torch
from torch import nn

from .cnn import CNNBlock3D, CNNBlockConfig
from ..docstrings import populate_docstring
from ..utils.activation_checkpointing import ActivationCheckpointing
from ..utils.custom_base_model import Field
from ..utils.rearrange import rearrange_channels

# %% ../../nbs/blocks/03_se_3d.ipynb 4
class SEBlock3DConfig(CNNBlockConfig):
    dim: int = Field(..., description="Number of input channels.")
    r: float = Field(..., description="Reduction ratio for the number of channels in the SE block.")

    kernel_size: int = Field(1, description="Kernel size for the convolutional layers in the SE block.")
    normalization: str = Field("batchnorm3d", description="Normalization layer to use in the SE block.")
    activation: str = Field("silu", description="Activation function to use in the SE block.")

    in_channels: None = Field(None, description="determined by dim and r")
    out_channels: None = Field(None, description="determined by dim and r")

# %% ../../nbs/blocks/03_se_3d.ipynb 6
@populate_docstring
class SEBlock3D(nn.Module):
    """A Squeeze-and-Excitation (SE) block for 3D data. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: SEBlock3DConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize an SEBlock3D block. Activation checkpointing level 2.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__()

        self.config = SEBlock3DConfig.model_validate(config | kwargs)

        dim = self.config.dim
        r = self.config.r

        excitation_dim = int(dim // r)

        self.squeeze = nn.AdaptiveAvgPool3d((1, 1, 1))
        self.excite = nn.Sequential(
            CNNBlock3D(
                self.config.model_dump()
                | {
                    "in_channels": dim,
                    "out_channels": excitation_dim,
                    "kernel_size": 1,
                    "stride": 1,
                    "padding": 0,
                },
                checkpointing_level,
            ),
            CNNBlock3D(
                self.config.model_dump()
                | {
                    "in_channels": excitation_dim,
                    "out_channels": dim,
                    "kernel_size": 1,
                    "stride": 1,
                    "padding": 0,
                    "activation": "sigmoid",
                },
                checkpointing_level,
            ),
        )

        self.checkpointing_level2 = ActivationCheckpointing(2, checkpointing_level)

    @populate_docstring
    def _forward(self, x: torch.Tensor, channels_first: bool = True) -> torch.Tensor:
        """Forward pass of the SEBlock3D block.

        Args:
            x: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        # x: (b, [dim], z, y, x, [dim])

        x = rearrange_channels(x, channels_first, True)
        # Now x is (b, dim, z, y, x)

        p = self.squeeze(x)
        # (b, dim, 1, 1, 1)
        p = self.excite(p)
        # (b, dim, 1, 1, 1)
        x = x * p
        # (b, dim, z, y, x)

        x = rearrange_channels(x, True, channels_first)
        # (b, [dim], z, y, x, [dim])

        return x

    @wraps(_forward)
    def forward(self, *args, **kwargs):
        return self.checkpointing_level2(self._forward, *args, **kwargs)
