# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/nets/03_swinv2_3d.ipynb.

# %% auto 0
__all__ = ['SwinV23DPatchMergingConfig', 'SwinV23DPatchSplittingConfig', 'SwinV23DBlockConfig', 'SwinV23DStageConfig',
           'SwinV23DEncoderDecoderConfig', 'SwinV23DEncoderWithPatchEmbeddingsConfig', 'SwinV23DLayerLogitScale',
           'SwinV23DLayer', 'SwinV23DBlock', 'SwinV23DPatchMerging', 'SwinV23DPatchSplitting', 'SwinV23DStage',
           'SwinV23DEncoderDecoderBase', 'SwinV23DEncoder', 'SwinV23DDecoder', 'SwinV23DEncoderWithPatchEmbeddings']

# %% ../../nbs/nets/03_swinv2_3d.ipynb 2
import numpy as np
import torch
from huggingface_hub import PyTorchModelHubMixin  # TODO
from torch import nn

from ..blocks.transformer import Attention3DWithMLP, Attention3DWithMLPConfig
from ..docstrings import populate_docstring
from vision_architectures.layers.embeddings import (
    AbsolutePositionEmbeddings3D,
    PatchEmbeddings3D,
    RelativePositionEmbeddings3DConfig,
    RelativePositionEmbeddings3DMetaNetwork,
)
from vision_architectures.nets.swin_3d import (
    Swin3DBlock,
    Swin3DBlockConfig,
    Swin3DEncoderDecoderBase,
    Swin3DEncoderDecoderConfig,
    Swin3DEncoderWithPatchEmbeddings,
    Swin3DEncoderWithPatchEmbeddingsConfig,
    Swin3DLayer,
    Swin3DPatchMerging,
    Swin3DPatchMergingConfig,
    Swin3DPatchSplitting,
    Swin3DPatchSplittingConfig,
    Swin3DStage,
    Swin3DStageConfig,
)
from ..utils.custom_base_model import Field

# %% ../../nbs/nets/03_swinv2_3d.ipynb 4
class SwinV23DPatchMergingConfig(Swin3DPatchMergingConfig):
    pass


class SwinV23DPatchSplittingConfig(Swin3DPatchSplittingConfig):
    pass


class SwinV23DBlockConfig(Swin3DBlockConfig):
    patch_merging: SwinV23DPatchMergingConfig | None = Field(
        None, description="Patch merging config if desired. Patch merging is applied before attention."
    )
    patch_splitting: SwinV23DPatchSplittingConfig | None = Field(
        None, description="Patch splitting config if desired. Patch splitting is applied after attention."
    )


class SwinV23DStageConfig(SwinV23DBlockConfig, Swin3DStageConfig):
    pass


class SwinV23DEncoderDecoderConfig(Swin3DEncoderDecoderConfig):
    stages: list[SwinV23DStageConfig]


class SwinV23DEncoderWithPatchEmbeddingsConfig(SwinV23DEncoderDecoderConfig, Swin3DEncoderWithPatchEmbeddingsConfig):
    pass

# %% ../../nbs/nets/03_swinv2_3d.ipynb 8
class SwinV23DLayerLogitScale(nn.Module):
    def __init__(self, num_heads):
        super().__init__()

        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)

    def forward(self):
        logit_scale = torch.clamp(self.logit_scale, max=np.log(1.0 / 0.01)).exp()
        return logit_scale

# %% ../../nbs/nets/03_swinv2_3d.ipynb 9
@populate_docstring
class SwinV23DLayer(Swin3DLayer):
    """SwinV2 3D Layer applying windowed attention with optional relative position embeddings.
    {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(
        self,
        config: RelativePositionEmbeddings3DConfig | Attention3DWithMLPConfig = {},
        checkpointing_level: int = 0,
        **kwargs
    ):
        """Initializes the SwinV23DLayer.

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

        # Update relative position bias to use the meta network
        relative_position_bias = None
        if self._use_relative_position_bias:
            relative_position_bias = RelativePositionEmbeddings3DMetaNetwork(
                self.embeddings_config, checkpointing_level=checkpointing_level
            )

        # Use SwinV2 logit scale
        logit_scale = SwinV23DLayerLogitScale(self.transformer_config.num_heads)

        # Re-initialize the transformer with the new relative position bias and logit scale
        self.transformer = Attention3DWithMLP(
            self.transformer_config,
            relative_position_bias=relative_position_bias,
            logit_scale=logit_scale,
            checkpointing_level=checkpointing_level,
        )

# %% ../../nbs/nets/03_swinv2_3d.ipynb 12
@populate_docstring
class SwinV23DBlock(Swin3DBlock):
    """SwinV2 3D Block consisting of two SwinV23DLayers: one with regular windows and one with shifted windows.
    {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: SwinV23DBlockConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the SwinV23DBlock.

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

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

        self.w_layer = SwinV23DLayer(self.config.model_dump(), checkpointing_level=checkpointing_level)
        self.sw_layer = SwinV23DLayer(self.config.model_dump(), checkpointing_level=checkpointing_level)

# %% ../../nbs/nets/03_swinv2_3d.ipynb 14
@populate_docstring
class SwinV23DPatchMerging(Swin3DPatchMerging):
    """Patch merging layer for SwinV23D. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: SwinV23DPatchMergingConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the SwinV23DPatchMerging layer.

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

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

        in_dim = self.config.in_dim * np.prod(self.config.merge_window_size)
        self.layer_norm = nn.LayerNorm(in_dim)
        self.proj = nn.Linear(in_dim, self.config.out_dim)

# %% ../../nbs/nets/03_swinv2_3d.ipynb 16
@populate_docstring
class SwinV23DPatchSplitting(Swin3DPatchSplitting):
    """Patch splitting layer for SwinV23D. {CLASS_DESCRIPTION_3D_DOC}

    This is a self-implemented class and is not part of the paper."""

    @populate_docstring
    def __init__(self, config: SwinV23DPatchSplittingConfig, checkpointing_level: int = 0, **kwargs):
        """Initialize the SwinV23DPatchSplitting layer.

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

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

        self.layer_norm = nn.LayerNorm(self.config.in_dim)
        self.proj = nn.Linear(self.config.in_dim, self.config.out_dim * np.prod(self.config.final_window_size))

# %% ../../nbs/nets/03_swinv2_3d.ipynb 18
@populate_docstring
class SwinV23DStage(Swin3DStage):
    """SwinV23D stage for SwinV23D. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: SwinV23DStageConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the SwinV23DStage.

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

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

        self.patch_merging = None
        if self.config.patch_merging is not None:
            self.patch_merging = SwinV23DPatchMerging(self.config.patch_merging)

        self.blocks = nn.ModuleList(
            [SwinV23DBlock(self.config) for _ in range(self.config.depth)],
        )

        self.patch_splitting = None
        if self.config.patch_splitting is not None:
            # This has been implemented to create a Swin-based decoder
            self.patch_splitting = SwinV23DPatchSplitting(self.config.patch_splitting)

# %% ../../nbs/nets/03_swinv2_3d.ipynb 22
class SwinV23DEncoderDecoderBase(Swin3DEncoderDecoderBase, PyTorchModelHubMixin):
    @populate_docstring
    def __init__(self, config: SwinV23DEncoderDecoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the SwinV23DEncoder/SwinV23DDecoder.

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

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

        self.stages = nn.ModuleList(
            [SwinV23DStage(stage_config, checkpointing_level) for stage_config in self.config.stages]
        )

# %% ../../nbs/nets/03_swinv2_3d.ipynb 24
@populate_docstring
class SwinV23DEncoder(SwinV23DEncoderDecoderBase):
    """3D Swin Transformer encoder. Assumes input has already been patchified/tokenized. {CLASS_DESCRIPTION_3D_DOC}"""

    def __init__(self, config: SwinV23DEncoderDecoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        super().__init__(config, checkpointing_level, **kwargs)

        for stage_config in self.config.stages:
            if stage_config.patch_splitting is not None:
                assert (
                    stage_config.patch_merging is not None
                ), "SwinV23DEncoder is not for decoding (mid blocks are ok)."

# %% ../../nbs/nets/03_swinv2_3d.ipynb 27
@populate_docstring
class SwinV23DDecoder(SwinV23DEncoderDecoderBase):
    """3D Swin Transformer decoder. Assumes input has already been patchified/tokenized. {CLASS_DESCRIPTION_3D_DOC}"""

    def __init__(self, config: SwinV23DEncoderDecoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        super().__init__(config, checkpointing_level, **kwargs)

        for stage_config in config.stages:
            if stage_config.patch_merging is not None:
                assert (
                    stage_config.patch_splitting is not None
                ), "SwinV23DDecoder is not for encoding (mid blocks are ok)."

# %% ../../nbs/nets/03_swinv2_3d.ipynb 30
@populate_docstring
class SwinV23DEncoderWithPatchEmbeddings(Swin3DEncoderWithPatchEmbeddings, PyTorchModelHubMixin):
    """3D SwinV2 transformer with 3D patch embeddings. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: SwinV23DEncoderWithPatchEmbeddingsConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the SwinV23DEncoderWithPatchEmbeddings.

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

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

        self.patchify = PatchEmbeddings3D(
            patch_size=self.config.patch_size,
            in_channels=self.config.in_channels,
            dim=self.config.stages[0].get_in_dim(),
            checkpointing_level=checkpointing_level,
        )
        self.absolute_position_embeddings = AbsolutePositionEmbeddings3D(
            dim=self.config.stages[0].get_in_dim(), learnable=False
        )
        self.encoder = SwinV23DEncoder(self.config, checkpointing_level=checkpointing_level)
