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

# %% auto 0
__all__ = ['CaiT3DMHSA', 'CaiT3DMHCA', 'CaiT3DLayerMLP', 'CaiT3DStage1Layer', 'CaiT3DStage2Layer', 'CaiT3DStage1', 'CaiT3DStage2',
           'CaiT3DStage2OnlyModel', 'CaiT3DModel']

# %% ../../nbs/nets/05_cait_3d.ipynb 2
import torch
from einops import repeat
from huggingface_hub import PyTorchModelHubMixin
from torch import nn

from .vit_3d import ViT3DLayerMLP, ViT3DMHCA, ViT3DMHSA

# %% ../../nbs/nets/05_cait_3d.ipynb 5
class CaiT3DMHSA(ViT3DMHSA):  # Multi-head self attention
    pass

# %% ../../nbs/nets/05_cait_3d.ipynb 6
class CaiT3DMHCA(ViT3DMHCA):  # Multi-head class attention
    def forward(self, class_tokens: torch.Tensor, embeddings: torch.Tensor):
        return super().forward(class_tokens, torch.cat([class_tokens, embeddings], dim=1))

# %% ../../nbs/nets/05_cait_3d.ipynb 8
class CaiT3DLayerMLP(ViT3DLayerMLP):
    pass

# %% ../../nbs/nets/05_cait_3d.ipynb 10
class CaiT3DStage1Layer(nn.Module):  # Self attention without class tokens
    def __init__(
        self,
        dim,
        num_heads,
        intermediate_ratio,
        layer_norm_eps,
        attn_drop_prob=0.0,
        proj_drop_prob=0.0,
        mlp_drop_prob=0.0,
    ):
        super().__init__()

        self.mhsa = CaiT3DMHSA(dim, num_heads, attn_drop_prob, proj_drop_prob)
        self.gamma1 = nn.Parameter(torch.empty(1, 1, dim))
        self.layernorm1 = nn.LayerNorm(dim, eps=layer_norm_eps)
        self.mlp = CaiT3DLayerMLP(dim, intermediate_ratio, mlp_drop_prob)
        self.gamma2 = nn.Parameter(torch.empty(1, 1, dim))
        self.layernorm2 = nn.LayerNorm(dim, eps=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)

    def forward(self, embeddings: torch.Tensor):  # This uses post-normalization
        # embeddings: (b, num_tokens, dim)

        res_connection1 = embeddings
        # (b, num_tokens, dim)

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

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

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

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

        return hidden_states

# %% ../../nbs/nets/05_cait_3d.ipynb 12
class CaiT3DStage2Layer(nn.Module):  # Attention with class tokens
    def __init__(
        self,
        dim,
        num_heads,
        intermediate_ratio,
        layer_norm_eps,
        attn_drop_prob=0.0,
        proj_drop_prob=0.0,
        mlp_drop_prob=0.0,
    ):
        super().__init__()

        self.mhca = CaiT3DMHCA(dim, num_heads, attn_drop_prob, proj_drop_prob)
        self.gamma1 = nn.Parameter(torch.empty(1, 1, dim))
        self.layernorm1 = nn.LayerNorm(dim, eps=layer_norm_eps)
        self.mlp = CaiT3DLayerMLP(dim, intermediate_ratio, mlp_drop_prob)
        self.gamma2 = nn.Parameter(torch.empty(1, 1, dim))
        self.layernorm2 = nn.LayerNorm(dim, eps=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)

    def forward(self, class_tokens: torch.Tensor, embeddings: torch.Tensor):  # This uses post-normalization
        # class_tokens: (b, num_class_tokens, dim)
        # embeddings: (b, num_embedding_tokens, dim)

        res_connection1 = class_tokens
        # (b, num_class_tokens, dim)

        hidden_states = self.mhca(class_tokens, embeddings)
        hidden_states = self.gamma1 * hidden_states
        hidden_states = self.layernorm1(hidden_states)
        # (b, num_class_tokens, dim)

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

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

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

        return hidden_states

# %% ../../nbs/nets/05_cait_3d.ipynb 15
class CaiT3DStage1(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()

        self.layers = nn.ModuleList(
            [
                CaiT3DStage1Layer(
                    config["dim"],
                    config["num_heads"],
                    config["intermediate_ratio"],
                    config["layer_norm_eps"],
                    config["attn_drop_prob"],
                    config["proj_drop_prob"],
                    config["mlp_drop_prob"],
                )
                for _ in range(config["encoder_depth"])
            ]
        )

    def forward(self, embeddings: torch.Tensor):
        # embeddings: (b, num_tokens, dim)

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

            layer_outputs.append(embeddings)

        return embeddings, layer_outputs

# %% ../../nbs/nets/05_cait_3d.ipynb 17
class CaiT3DStage2(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()

        self.layers = nn.ModuleList(
            [
                CaiT3DStage2Layer(
                    config["dim"],
                    config["num_heads"],
                    config["intermediate_ratio"],
                    config["layer_norm_eps"],
                    config["attn_drop_prob"],
                    config["proj_drop_prob"],
                    config["mlp_drop_prob"],
                )
                for _ in range(config["encoder_depth"])
            ]
        )

    def forward(self, class_tokens: torch.Tensor, embeddings: torch.Tensor):
        # class_tokens: (b, num_class_tokens, dim)
        # embeddings: (b, num_embed_tokens, dim)

        class_embeddings = class_tokens

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

            layer_outputs.append(class_embeddings)

        return class_embeddings, layer_outputs

# %% ../../nbs/nets/05_cait_3d.ipynb 20
class CaiT3DStage2OnlyModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()

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

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

    def forward(self, embeddings: torch.Tensor):
        # embeddings: (b, num_embedding_tokens, dim)

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

        class_embeddings, layer_outputs = self.class_attention(class_tokens, embeddings)
        # 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

        return class_logits, class_embeddings, layer_outputs

# %% ../../nbs/nets/05_cait_3d.ipynb 22
class CaiT3DModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()

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

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

    def forward(self, tokens: torch.Tensor):
        # tokens: (b, num_embedding_tokens, dim)

        embeddings, layer_outputs1 = self.self_attention(tokens)

        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)
        # 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

        return class_logits, class_embeddings, [layer_outputs1, layer_outputs2]
