# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/nets/05_cait_3d.ipynb.

# %% auto 0
__all__ = ['CaiTAttentionWithMLPConfig', 'CaiTStage1Config', 'CaiTStage2Config', 'CaiTConfig', 'CaiTAttentionWithMLP',
           'CaiTStage1', 'CaiTStage2', 'CaiT1D', 'CaiT3D']

# %% ../../nbs/nets/05_cait_3d.ipynb 2
from functools import wraps

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

from ..blocks.transformer import Attention1DMLP, Attention1DMLPConfig
from ..docstrings import populate_docstring
from ..layers.attention import Attention1D, Attention1DConfig
from ..utils.activation_checkpointing import ActivationCheckpointing
from ..utils.custom_base_model import Field, model_validator
from ..utils.rearrange import rearrange_channels

# %% ../../nbs/nets/05_cait_3d.ipynb 4
class CaiTAttentionWithMLPConfig(Attention1DConfig, Attention1DMLPConfig):
    layer_norm_eps: float = Field(1e-6, description="Epsilon value for the layer normalization.")


class CaiTStage1Config(CaiTAttentionWithMLPConfig):
    stage1_depth: int = Field(..., description="Number of layers in stage 1.", ge=0)


class CaiTStage2Config(CaiTAttentionWithMLPConfig):
    num_class_tokens: int = Field(1, description="Number of class tokens to be added in stage 2.", ge=0)
    stage2_depth: int = Field(..., description="Number of layers in stage 2.", ge=0)


class CaiTConfig(CaiTStage1Config, CaiTStage2Config):
    @model_validator(mode="after")
    def validate(self):
        super().validate()
        assert self.stage1_depth + self.stage2_depth > 0, "There should be atleast one layer in the model."
        return self

# %% ../../nbs/nets/05_cait_3d.ipynb 7
@populate_docstring
class CaiTAttentionWithMLP(nn.Module):
    """Attention layer used in the CaiT 3D model. Introduces learnable gamma scaling of hidden states after the self
    attention and MLP layers. {CLASS_DESCRIPTION_1D_DOC}"""

    @populate_docstring
    def __init__(self, config: CaiTAttentionWithMLPConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the CaiT 3D attention layer.

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

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

        self.mhsa = Attention1D(self.config, checkpointing_level=checkpointing_level)
        self.gamma1 = nn.Parameter(torch.empty(1, 1, self.config.dim))
        self.layernorm1 = nn.LayerNorm(self.config.dim, eps=self.config.layer_norm_eps)
        self.mlp = Attention1DMLP(self.config, checkpointing_level=checkpointing_level)
        self.gamma2 = nn.Parameter(torch.empty(1, 1, self.config.dim))
        self.layernorm2 = nn.LayerNorm(self.config.dim, eps=self.config.layer_norm_eps)

        nn.init.uniform_(self.gamma1, a=-1e-4, b=1e-4)
        nn.init.uniform_(self.gamma2, a=-1e-4, b=1e-4)

        self.checkpointing_level3 = ActivationCheckpointing(3, checkpointing_level)

    @populate_docstring
    def _forward(self, q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
        """Pass the input q and kv tensors through the q, k, and v matrices and then pass them through the CaiT
        attention layer.

        Args:
            q: {INPUT_1D_DOC}
            kv: {INPUT_1D_DOC}

        Returns:
            {OUTPUT_1D_DOC}
        """
        # q: (b, num_tokens_in_q, dim)
        # kv: (b, num_tokens_in_kv, dim)

        res_connection1 = q
        # (b, num_tokens, dim)

        hidden_states = self.layernorm1(q)
        hidden_states = self.mhsa(hidden_states, kv, kv)
        hidden_states = self.gamma1 * hidden_states
        # (b, num_tokens, dim)

        res_connection2 = hidden_states + res_connection1
        # (b, num_tokens, dim)

        hidden_states = self.layernorm2(hidden_states)
        hidden_states = self.mlp(res_connection2)
        hidden_states = self.gamma2 * hidden_states
        # (b, num_tokens, dim)

        hidden_states = hidden_states + res_connection2
        # (b, num_tokens, dim)

        return hidden_states

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

# %% ../../nbs/nets/05_cait_3d.ipynb 10
@populate_docstring
class CaiTStage1(nn.Module, PyTorchModelHubMixin):
    """CaiT stage 1. Performs self attention without class tokens focusing on learning features among tokens.
    {CLASS_DESCRIPTION_1D_DOC}"""

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

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

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

        self.layers = nn.ModuleList(
            [
                CaiTAttentionWithMLP(self.config, checkpointing_level=checkpointing_level)
                for _ in range(self.config.stage1_depth)
            ]
        )

        self.checkpointing_level4 = ActivationCheckpointing(4, checkpointing_level)

    @populate_docstring
    def _forward(self, embeddings: torch.Tensor, return_intermediates: bool = False) -> torch.Tensor:
        """Pass the input embeddings through the CaiT stage 1 layers.

        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 layer outputs."""
        # embeddings: (b, num_tokens, dim)

        layer_outputs = []
        for encoder_layer in self.layers:
            embeddings = encoder_layer(embeddings, 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_level4(self._forward, *args, **kwargs)

# %% ../../nbs/nets/05_cait_3d.ipynb 12
@populate_docstring
class CaiTStage2(nn.Module, PyTorchModelHubMixin):
    """CaiT stage 2. Performs cross attention between class tokens and learned features from stage 1.
    {CLASS_DESCRIPTION_1D_DOC}"""

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

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

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

        self.layers = nn.ModuleList(
            [
                CaiTAttentionWithMLP(self.config, checkpointing_level=checkpointing_level)
                for _ in range(self.config.stage2_depth)
            ]
        )

        self.checkpointing_level4 = ActivationCheckpointing(4, checkpointing_level)

    @populate_docstring
    def _forward(
        self, class_tokens: torch.Tensor, embeddings: torch.Tensor, return_intermediates: bool = False
    ) -> torch.Tensor:
        """Pass the input embeddings through the CaiT stage 2 layers.

        Args:

            class_tokens: {INPUT_1D_DOC}
            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 layer outputs."""
        # embeddings: (b, num_tokens, dim)

        class_embeddings = class_tokens

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

            layer_outputs.append(class_embeddings)

        if return_intermediates:
            return class_embeddings, layer_outputs
        return class_embeddings

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

# %% ../../nbs/nets/05_cait_3d.ipynb 15
@populate_docstring
class CaiT1D(nn.Module, PyTorchModelHubMixin):
    """End-to-end CaiT model for classification. {CLASS_DESCRIPTION_1D_DOC}"""

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

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

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

        self.class_tokens = nn.Parameter(torch.randn(1, self.config.num_class_tokens, self.config.dim))

        self.self_attention = CaiTStage1(self.config)
        self.class_attention = CaiTStage2(self.config)
        self.classifiers = nn.ModuleList([nn.Linear(config.dim, 1) for i in range(self.config.num_class_tokens)])

        self.checkpointing_level5 = ActivationCheckpointing(5, checkpointing_level)

    @populate_docstring
    def _forward(self, tokens: torch.Tensor, return_intermediates: bool = False) -> torch.Tensor | tuple:
        """Pass the input embeddings through the CaiT layers. Expects flattened input.

        Args:
            tokens: {INPUT_1D_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

        Returns:
            {OUTPUT_1D_DOC}
        """
        # tokens: (b, num_embedding_tokens, dim)

        embeddings, layer_outputs1 = self.self_attention(tokens, return_intermediates=True)

        class_tokens = repeat(self.class_tokens, "1 n d -> b n d", b=embeddings.shape[0])
        # (b, num_class_tokens, dim)

        class_embeddings, layer_outputs2 = self.class_attention(class_tokens, embeddings, return_intermediates=True)
        # class_embeddings: (b, num_class_tokens, dim)
        # layer_outputs: list of (b, num_embedding_tokens, dim)

        class_logits = torch.cat(
            [self.classifiers[i](class_embeddings[:, i]) for i in range(len(self.classifiers))],
            dim=1,
        )
        # list of (b, num_classes) for each class token

        if return_intermediates:
            return class_logits, class_embeddings, [layer_outputs1, layer_outputs2]
        return class_logits

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

# %% ../../nbs/nets/05_cait_3d.ipynb 17
@populate_docstring
class CaiT3D(CaiT1D):
    """End-to-end CaiT model for classification. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def _forward(
        self, tokens: torch.Tensor, channels_first: bool = True, return_intermediates: bool = False
    ) -> torch.Tensor | tuple:
        """Pass the input embeddings through the CaiT layers. Expects flattened input.

        Args:
            tokens: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

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

        tokens = rearrange_channels(tokens, channels_first, False)
        # tokens: (b, z, y, x, dim)
        tokens = rearrange(tokens, "b z y x dim -> b (z y x) dim").contiguous()
        # (b, T, dim)

        return super()._forward(tokens, return_intermediates)
