# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/nets/04_vit_3d.ipynb.

# %% auto 0
__all__ = ['ViTEncoderDecoderConfig', 'ViTEncoderConfig', 'ViT3DEncoderWithPatchEmbeddingsConfig', 'ViTDecoderConfig',
           'ViTEncoder', 'ViTDecoder', 'ViT3DEncoderWithPatchEmbeddings']

# %% ../../nbs/nets/04_vit_3d.ipynb 2
from functools import wraps
from typing import Literal

import torch
from einops import rearrange, repeat
from huggingface_hub import PyTorchModelHubMixin
from torch import nn

from ..blocks.transformer import TransformerDecoderBlock1D, TransformerEncoderBlock1D
from ..docstrings import populate_docstring
from ..layers.embeddings import AbsolutePositionEmbeddings3D, PatchEmbeddings3D
from ..utils.activation_checkpointing import ActivationCheckpointing
from ..utils.custom_base_model import CustomBaseModel, Field
from ..utils.rearrange import rearrange_channels

# %% ../../nbs/nets/04_vit_3d.ipynb 4
class ViTEncoderDecoderConfig(CustomBaseModel):
    dim: int = Field(..., description="Dimension of the patches.")
    num_heads: int = Field(..., description="Number of attention heads.")
    mlp_ratio: int = Field(..., description="Ratio of mlp hidden dim to embedding dim.")
    layer_norm_eps: float = Field(..., description="Epsilon value for the layer normalization.")
    attn_drop_prob: float = Field(0.0, description="Dropout probability for attention weights.")
    proj_drop_prob: float = Field(0.0, description="Dropout probability for the projection layer.")
    mlp_drop_prob: float = Field(0.0, description="Dropout probability for the mlp layer.")
    norm_location: Literal["pre", "post"] = Field(
        "pre", description="Location of the normalization layer with respect to the attention."
    )


class ViTEncoderConfig(ViTEncoderDecoderConfig):
    encoder_depth: int = Field(..., description="Number of encoder blocks.")


class ViT3DEncoderWithPatchEmbeddingsConfig(ViTEncoderConfig):
    patch_size: tuple[int, int, int] = Field(..., description="Size of the patches to be embedded.")
    in_channels: int = Field(..., description="Number of input channels in the input datapoints.")
    num_class_tokens: int = Field(..., description="Number of class tokens to be added.")

    drop_prob: float = Field(0.0, description="Dropout probability for the embedding layer.")


class ViTDecoderConfig(ViTEncoderDecoderConfig):
    decoder_depth: int

# %% ../../nbs/nets/04_vit_3d.ipynb 8
@populate_docstring
class ViTEncoder(nn.Module, PyTorchModelHubMixin):
    """Vision Transformer encoder. {CLASS_DESCRIPTION_1D_DOC}"""

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

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

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

        self.layers = nn.ModuleList(
            [
                TransformerEncoderBlock1D(self.config.model_dump(), checkpointing_level=checkpointing_level)
                for _ in range(self.config.encoder_depth)
            ]
        )

        self.checkpointing_level5 = ActivationCheckpointing(5, checkpointing_level)

    @populate_docstring
    def _forward(
        self, embeddings: torch.Tensor, return_intermediates: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
        """Pass the input embeddings through the ViT encoder (self attention).

        Args:
            embeddings: {INPUT_1D_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

        Returns:
            {OUTPUT_1D_DOC} If `return_intermediates` is True, returns a tuple of the output embeddings and a list of
            intermediate embeddings.
        """
        # hidden_states: (b, num_tokens, dim)

        layer_outputs = []
        for encoder_layer in self.layers:
            embeddings = encoder_layer(embeddings)
            # (b, num_tokens, dim)

            layer_outputs.append(embeddings)

        if return_intermediates:
            return embeddings, layer_outputs
        return embeddings

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

# %% ../../nbs/nets/04_vit_3d.ipynb 11
@populate_docstring
class ViTDecoder(nn.Module, PyTorchModelHubMixin):
    """Vision Transformer decoder. {CLASS_DESCRIPTION_1D_DOC}"""

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

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

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

        self.layers = nn.ModuleList(
            [
                TransformerDecoderBlock1D(self.config.model_dump(), checkpointing_level=checkpointing_level)
                for _ in range(self.config.decoder_depth)
            ]
        )

        self.checkpointing_level5 = ActivationCheckpointing(5, checkpointing_level)

    @populate_docstring
    def _forward(
        self, q: torch.Tensor, kv: torch.Tensor, return_intermediates: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
        """Pass the input embeddings through the ViT decoder (self attention + cross attention).

        Args:
            q: Input to the query matrix. {INPUT_1D_DOC}
            kv: Input tot he key and value matrices. {INPUT_1D_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

        Returns:
            {OUTPUT_1D_DOC} If `return_intermediates` is True, returns a tuple of the output embeddings and a list of
            intermediate embeddings.
        """
        # q: (b, num_q_tokens, dim)
        # kv: (b, num_kv_tokens, dim)

        embeddings = q

        layer_outputs = []
        for decoder_layer in self.layers:
            embeddings = decoder_layer(embeddings, kv)
            # (b, num_q_tokens, dim)

            layer_outputs.append(embeddings)

        if return_intermediates:
            return embeddings, layer_outputs
        return embeddings

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

# %% ../../nbs/nets/04_vit_3d.ipynb 14
@populate_docstring
class ViT3DEncoderWithPatchEmbeddings(nn.Module, PyTorchModelHubMixin):
    """Patchification of input array followd by a ViT encoder. {CLASS_DESCRIPTION_3D_DOC}"""

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

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

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

        self.patchify = PatchEmbeddings3D(
            patch_size=self.config.patch_size,
            in_channels=self.config.in_channels,
            dim=self.config.dim,
            checkpointing_level=checkpointing_level,
        )
        self.absolute_position_embeddings = AbsolutePositionEmbeddings3D(dim=self.config.dim, learnable=False)
        self.pos_drop = nn.Dropout(self.config.drop_prob)
        self.num_class_tokens = self.config.num_class_tokens
        if self.num_class_tokens > 0:
            self.class_tokens = nn.Parameter(torch.randn(1, self.config.num_class_tokens, self.config.dim))
        self.encoder = ViTEncoder(self.config, checkpointing_level=checkpointing_level)

    @populate_docstring
    def forward(
        self,
        pixel_values: torch.Tensor,
        spacings: torch.Tensor,
        channels_first: bool = True,
        return_intermediates: bool = False,
    ) -> tuple[torch.Tensor, list[torch.Tensor]] | tuple[torch.Tensor, list[torch.Tensor], list[torch.Tensor]]:
        """Patchify the input datapoint and then pass through the ViT encoder (self attention).

        Args:
            pixel_values: {INPUT_3D_DOC}
            spacings: {SPACINGS_DOC}
            channels_first: {CHANNELS_FIRST_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

        Returns:
            {OUTPUT_1D_DOC} If `return_intermediates` is True, returns a tuple of the output embeddings and a list of
            intermediate embeddings.
        """
        # pixel_values: (b, [c], z, y, x, [c])
        # spacings: (b, 3)

        pixel_values = rearrange_channels(pixel_values, channels_first, True)
        # (b, c, z, y, x)

        embeddings = self.patchify(pixel_values, channels_first=True)
        # (b, dim, num_patches_z, num_patches_y, num_patches_x)

        embeddings = self.absolute_position_embeddings(embeddings, spacings=spacings, channels_first=True)
        # (b, dim, num_patches_z, num_patches_y, num_patches_x)

        embeddings = rearrange(embeddings, "b e nz ny nx -> b (nz ny nx) e").contiguous()
        # (b, num_tokens, dim)

        embeddings = self.pos_drop(embeddings)
        # (b, num_tokens, dim)

        class_tokens = None
        if self.num_class_tokens > 0:
            class_tokens = repeat(self.class_tokens, "1 n d -> b n d", b=embeddings.shape[0])
            embeddings = torch.cat([class_tokens, embeddings], dim=1)
            # (b, num_tokens + num_class_tokens, dim)

        encoded, layer_outputs = self.encoder(embeddings, return_intermediates=True)
        # encoded: (b, num_tokens (+ num_class_tokens), dim)
        # layer_outputs: list of (b, num_tokens (+ 1), dim)

        if self.num_class_tokens > 0:
            class_tokens = encoded[:, : self.num_class_tokens]
            encoded = encoded[:, self.num_class_tokens :]

        if return_intermediates:
            return encoded, class_tokens, layer_outputs
        return encoded, class_tokens
