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

# %% auto 0
__all__ = ['get_coords_grid', 'ViT3DMHA', 'ViT3DMHSA', 'ViT3DMHCA', 'ViT3DLayerMLP', 'ViT3DEncoderLayer', 'ViT3DDecoderLayer',
           'ViT3DEncoder', 'ViT3DPatchEmbeddings', 'get_3d_position_embeddings',
           'embed_spacings_in_position_embeddings', 'ViT3DEmbeddings', 'ViT3DModel', 'ViT3DMIMDecoder', 'ViT3DMIM',
           'ViT3DSimMIM']

# %% ../../nbs/nets/04_vit_3d.ipynb 2
import numpy as np
import torch
from einops import rearrange, repeat
from huggingface_hub import PyTorchModelHubMixin
from torch import nn

# %% ../../nbs/nets/04_vit_3d.ipynb 5
def get_coords_grid(grid_size):
    d, h, w = grid_size

    grid_d = torch.arange(d, dtype=torch.int32)
    grid_h = torch.arange(h, dtype=torch.int32)
    grid_w = torch.arange(w, dtype=torch.int32)

    grid = torch.meshgrid(grid_w, grid_h, grid_d, indexing="ij")
    grid = torch.stack(grid, axis=0)
    # (3, d, h, w)

    return grid

# %% ../../nbs/nets/04_vit_3d.ipynb 6
class ViT3DMHA(nn.Module):  # Multi-head-attention
    def __init__(
        self,
        dim,
        num_heads,
        attn_drop_prob=0.0,
        proj_drop_prob=0.0,
    ):
        super().__init__()

        assert dim % num_heads == 0, "dimension must be divisible by number of heads"

        self.dim = dim
        self.num_heads = num_heads

        self.per_head_dim = int(dim // num_heads)

        self.W_q = nn.Linear(dim, dim)
        self.W_k = nn.Linear(dim, dim)
        self.W_v = nn.Linear(dim, dim)
        self.attn_drop = nn.Dropout(attn_drop_prob)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop_prob)

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
        # q: (b, num_tokens_of_q, dim)
        # k: (b, num_tokens_of_k, dim)
        # v: (b, num_tokens_of_v, dim)

        query = rearrange(
            self.W_q(q), "b num_tokens (num_heads d) -> b num_heads num_tokens d", num_heads=self.num_heads
        )
        key = rearrange(self.W_k(k), "b num_tokens (num_heads d) -> b num_heads num_tokens d", num_heads=self.num_heads)
        value = rearrange(
            self.W_v(v), "b num_tokens (num_heads d) -> b num_heads num_tokens d", num_heads=self.num_heads
        )
        # Each is (b, num_heads, num_tokens_in_the_tensor, per_head_dim)

        attention_scores = query @ rearrange(key, "b num_heads num_tokens d -> b num_heads d num_tokens")
        attention_scores = attention_scores / (self.per_head_dim**0.5)
        # (b, num_heads, num_tokens_in_q, num_tokens_in_kv)

        attention_probs = nn.functional.softmax(attention_scores, dim=-1)
        attention_probs = self.attn_drop(attention_probs)
        # (b, num_heads, num_tokens_in_q, num_tokens_in_kv)

        context = attention_probs @ value
        # (b, num_heads, num_tokens_in_q, per_head_dim)
        context = rearrange(context, "b num_heads num_tokens d -> b num_tokens (num_heads d)")
        # (b, num_tokens_of_q, dim)

        context = self.proj(context)
        context = self.proj_drop(context)
        # (b, num_tokens_of_q, dim)

        return context

# %% ../../nbs/nets/04_vit_3d.ipynb 7
class ViT3DMHSA(ViT3DMHA):  # Multi head self attention
    def __init__(
        self,
        dim,
        num_heads,
        attn_drop_prob=0.0,
        proj_drop_prob=0.0,
    ):
        super().__init__(
            dim,
            num_heads,
            attn_drop_prob,
            proj_drop_prob,
        )

    def forward(self, qkv: torch.Tensor):
        return super().forward(qkv, qkv, qkv)

# %% ../../nbs/nets/04_vit_3d.ipynb 9
class ViT3DMHCA(ViT3DMHA):  # Multi head cross attention
    def __init__(
        self,
        dim,
        num_heads,
        attn_drop_prob=0.0,
        proj_drop_prob=0.0,
    ):
        super().__init__(
            dim,
            num_heads,
            attn_drop_prob,
            proj_drop_prob,
        )

    def forward(self, q: torch.Tensor, kv: torch.Tensor):
        return super().forward(q, kv, kv)

# %% ../../nbs/nets/04_vit_3d.ipynb 12
class ViT3DLayerMLP(nn.Module):
    def __init__(self, dim, intermediate_ratio, dropout_prob=0.0):
        super().__init__()
        self.dense1 = nn.Linear(dim, dim * intermediate_ratio)
        self.act = nn.GELU()
        self.dense2 = nn.Linear(dim * intermediate_ratio, dim)
        self.dropout = nn.Dropout(dropout_prob)

    def forward(self, hidden_states: torch.Tensor):
        # hidden_states: (b, num_tokens, dim)
        hidden_states = self.dense1(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.dense2(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states

# %% ../../nbs/nets/04_vit_3d.ipynb 14
class ViT3DEncoderLayer(nn.Module):
    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 = ViT3DMHSA(dim, num_heads, attn_drop_prob, proj_drop_prob)
        self.layernorm1 = nn.LayerNorm(dim, eps=layer_norm_eps)
        self.mlp = ViT3DLayerMLP(dim, intermediate_ratio, mlp_drop_prob)
        self.layernorm2 = nn.LayerNorm(dim, eps=layer_norm_eps)

    def forward(self, qkv: torch.Tensor):
        """Note: This uses post-normalization instead of pre-normalization. This is inspired by SwinV2"""
        # qkv: (b, num_tokens, dim)

        res_connection1 = qkv
        # (b, num_tokens, dim)

        hidden_states = self.mhsa(qkv)
        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.layernorm2(hidden_states)
        # (b, num_tokens, dim)

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

        return hidden_states

# %% ../../nbs/nets/04_vit_3d.ipynb 16
class ViT3DDecoderLayer(nn.Module):
    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 = ViT3DMHSA(dim, num_heads, attn_drop_prob, proj_drop_prob)
        self.layernorm1 = nn.LayerNorm(dim, eps=layer_norm_eps)
        self.mhca = ViT3DMHCA(dim, num_heads, attn_drop_prob, proj_drop_prob)
        self.layernorm2 = nn.LayerNorm(dim, eps=layer_norm_eps)
        self.mlp = ViT3DLayerMLP(dim, intermediate_ratio, mlp_drop_prob)
        self.layernorm3 = nn.LayerNorm(dim, eps=layer_norm_eps)

    def forward(self, q: torch.Tensor, kv: torch.Tensor):  # This uses post-normalization
        # q: (b, num_tokens_in_q, dim)
        # kv: (b, num_tokens_in_kv, dim)

        res_connection1 = q
        # (b, num_tokens_in_q, dim)

        hidden_states_q = self.mhsa(q)
        hidden_states_q = self.layernorm1(hidden_states_q)
        # (b, num_tokens_in_q, dim)

        res_connection2 = hidden_states_q + res_connection1
        # (b, num_tokens_in_q, dim)

        hidden_states = self.mhca(res_connection2, kv)
        hidden_states = self.layernorm2(hidden_states)
        # (b, num_tokens_in_q, dim)

        res_connection3 = hidden_states + res_connection2
        # (b, num_tokens_in_q, dim)

        hidden_states = self.mlp(res_connection3)
        hidden_states = self.layernorm3(hidden_states)
        # (b, num_tokens_in_q, dim)

        hidden_states = hidden_states + res_connection3
        # (b, num_tokens_in_q, dim)

        return hidden_states

# %% ../../nbs/nets/04_vit_3d.ipynb 20
class ViT3DEncoder(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()

        self.layers = nn.ModuleList(
            [
                ViT3DEncoderLayer(
                    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):
        # 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)

        return embeddings, layer_outputs

# %% ../../nbs/nets/04_vit_3d.ipynb 24
class ViT3DPatchEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()

        patch_size = config["patch_size"]
        num_channels = config["in_channels"]
        dim = config["dim"]

        self.patch_embeddings = nn.Conv3d(
            in_channels=num_channels,
            out_channels=dim,
            kernel_size=patch_size,
            stride=patch_size,
        )

    def forward(self, pixel_values: torch.Tensor):
        # pixel_values: (b, c, z, y, x)

        embeddings = self.patch_embeddings(pixel_values)
        # (b, dim, num_tokens_z, num_tokens_y, num_tokens_x)

        return embeddings

# %% ../../nbs/nets/04_vit_3d.ipynb 27
def get_3d_position_embeddings(embedding_size, grid_size, patch_size=(1, 1, 1)):
    if embedding_size % 6 != 0:
        raise ValueError("embed_dim must be divisible by 6")

    grid = get_coords_grid(grid_size)
    # (3, d, h, w)

    grid = rearrange(grid, "x d h w -> x 1 d h w")
    # (3, 1, d, h, w)

    omega = torch.arange(embedding_size // 6, dtype=torch.float32)
    omega /= embedding_size / 6.0
    omega = 1.0 / 10000**omega
    # (d // 6)

    patch_multiplier = torch.Tensor(patch_size) / min(patch_size)

    position_embeddings = []
    for i, grid_subset in enumerate(grid):
        grid_subset = grid_subset.reshape(-1)
        out = torch.einsum("m,d->md", grid_subset, omega)

        emb_sin = torch.sin(out)
        emb_cos = torch.cos(out)

        emb = torch.cat([emb_sin, emb_cos], axis=1) * patch_multiplier[i]
        position_embeddings.append(emb)

    position_embeddings = torch.cat(position_embeddings, axis=1)
    # (embedding_size, d * h * w)
    d, h, w = grid_size
    position_embeddings = rearrange(position_embeddings, "(d h w) e -> 1 e d h w", d=d, h=h, w=w)
    # (1, embedding_size, d, h, w)

    return position_embeddings

# %% ../../nbs/nets/04_vit_3d.ipynb 28
def embed_spacings_in_position_embeddings(embeddings: torch.Tensor, spacings: torch.Tensor):
    assert spacings.ndim == 2, "Please provide spacing information for each batch element"
    _, embedding_size, _, _, _ = embeddings.shape
    assert embedding_size % 3 == 0, "To embed spacing info, the embedding size must be divisible by 3"
    embeddings = embeddings * repeat(spacings, f"B S -> B (S {int(embedding_size / 3)}) 1 1 1", S=3)

    return embeddings

# %% ../../nbs/nets/04_vit_3d.ipynb 29
class ViT3DEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.config = config

        dim = config["dim"]

        self.patch_embeddings = ViT3DPatchEmbeddings(config)
        self.layer_norm = nn.LayerNorm(dim)

        self.absolute_position_embeddings = None
        if config["use_absolute_position_embeddings"]:
            grid_size = (
                config["image_size"][0] // config["patch_size"][0],
                config["image_size"][1] // config["patch_size"][1],
                config["image_size"][2] // config["patch_size"][2],
            )
            if config["learnable_absolute_position_embeddings"]:
                self.absolute_position_embeddings = nn.Parameter(
                    torch.randn(1, dim, grid_size[0], grid_size[1], grid_size[2])
                )
            else:
                self.absolute_position_embeddings = get_3d_position_embeddings(dim, grid_size, config["patch_size"])

    def forward(
        self,
        pixel_values: torch.Tensor,
        spacings: torch.Tensor,
        mask_patches: torch.Tensor = None,
        mask_token: torch.Tensor = None,
    ):
        # pixel_values: (b, c, z, y, x)

        embeddings = self.patch_embeddings(pixel_values)
        # (b, dim, num_tokens_z, num_tokens_y, num_tokens_x)
        embeddings = rearrange(embeddings, "b d nz ny nx -> b nz ny nx d")
        embeddings = self.layer_norm(embeddings)
        embeddings = rearrange(embeddings, "b nz ny nx d -> b d nz ny nx")
        # (b, dim, num_tokens_z, num_tokens_y, num_tokens_x)

        if mask_patches is not None:
            # mask_patches (binary mask): (b, num_tokens_z, num_tokens_y, num_tokens_x)
            # mask_token: (1, dim, 1, 1, 1)
            mask_patches = repeat(mask_patches, "b z y x -> b d z y x", d=embeddings.shape[1])
            embeddings = (embeddings * (1 - mask_patches)) + (mask_patches * mask_token)

        if self.absolute_position_embeddings is not None:
            absolute_position_embeddings = self.absolute_position_embeddings.to(embeddings.device)
            # (1, dim, num_tokens_z, num_tokens_y, num_tokens_x)
            if self.config["embed_spacing_info"]:
                absolute_position_embeddings = embed_spacings_in_position_embeddings(
                    absolute_position_embeddings, spacings
                )
                # (b, dim, num_tokens_z, num_tokens_y, num_tokens_x)

            embeddings = embeddings + absolute_position_embeddings
            # (b, dim, num_tokens_z, num_tokens_y, num_tokens_x)

        embeddings = rearrange(embeddings, "b d nz ny nx -> b (nz ny nx) d")

        return embeddings

# %% ../../nbs/nets/04_vit_3d.ipynb 33
class ViT3DModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()

        self.embeddings = ViT3DEmbeddings(config)
        self.pos_drop = nn.Dropout(config.get("drop_prob", 0.0))
        self.num_class_tokens = config["num_class_tokens"]
        if self.num_class_tokens > 0:
            self.class_tokens = nn.Parameter(torch.randn(1, config["num_class_tokens"], config["dim"]))
        self.encoder = ViT3DEncoder(config)

    def forward(
        self,
        pixel_values: torch.Tensor,
        spacings: torch.Tensor,
        mask_patches: torch.Tensor = None,
        mask_token: torch.Tensor = None,
    ):
        # pixel_values: (b, c, z, y, x)
        # spacings: (b, 3)
        # mask_patches: (num_tokens_z, num_tokens_y, num_tokens_x)

        embeddings = self.embeddings(pixel_values, spacings, mask_patches, mask_token)
        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 + 1, dim)

        encoded, layer_outputs = self.encoder(embeddings)
        # encoded: (b, num_tokens (+ 1), 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 :]

        return class_tokens, encoded, layer_outputs

# %% ../../nbs/nets/04_vit_3d.ipynb 36
class ViT3DMIMDecoder(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.image_size = config["image_size"]
        self.in_channels = config["in_channels"]
        self.patch_size = config["patch_size"]

        dim = config["dim"]
        out_dim = np.prod(self.patch_size) * self.in_channels

        self.decoder = nn.Linear(dim, out_dim)

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

        decoded = self.decoder(encodings)
        # (b, num_tokens, new_dim)

        decoded = rearrange(
            decoded,
            "b (nz ny nx) (c pz py px) -> b c (nz pz) (ny py) (nx px)",
            c=self.in_channels,
            pz=self.patch_size[0],
            py=self.patch_size[1],
            px=self.patch_size[2],
            nz=self.image_size[0] // self.patch_size[0],
            ny=self.image_size[1] // self.patch_size[1],
            nx=self.image_size[2] // self.patch_size[2],
        )
        # (b, c, z, y, x)

        return decoded

# %% ../../nbs/nets/04_vit_3d.ipynb 38
class ViT3DMIM(nn.Module):
    def __init__(self, config):
        super().__init__()

        assert config["num_class_tokens"] == 0, "MIM does not support class tokens"

        self.image_size = config["image_size"]
        self.patch_size = config["patch_size"]
        self.in_channels = config["in_channels"]
        self.mask_ratio = config["mask_ratio"]

        self.vit = ViT3DModel(config)
        self.decoder = ViT3DMIMDecoder(config)

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

    def mask_image(self, pixel_values: torch.Tensor):
        b = pixel_values.shape[0]

        mask_ratio = self.mask_ratio
        grid_size = tuple([size // patch for size, patch in zip(self.image_size, self.patch_size)])
        num_tokens = np.prod(grid_size)
        mask_patches = []
        for _ in range(b):
            _mask_patches = torch.zeros(num_tokens, dtype=torch.int8, device=pixel_values.device)
            _mask_patches[: int(mask_ratio * num_tokens)] = 1
            _mask_patches = _mask_patches[torch.randperm(num_tokens)]
            _mask_patches = rearrange(_mask_patches, "(z y x) -> z y x", z=grid_size[0], y=grid_size[1], x=grid_size[2])
            mask_patches.append(_mask_patches)
        mask_patches: torch.Tensor = torch.stack(mask_patches, dim=0)

        return mask_patches

# %% ../../nbs/nets/04_vit_3d.ipynb 39
class ViT3DSimMIM(ViT3DMIM, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__(config)

    @staticmethod
    def loss_fn(pred: torch.Tensor, target: torch.Tensor, reduction="mean"):
        return nn.functional.l1_loss(pred, target, reduction=reduction)

    def forward(self, pixel_values: torch.Tensor, spacings: torch.Tensor):
        mask_patches = self.mask_image(pixel_values)

        _, encodings, _ = self.vit(pixel_values, spacings, mask_patches, self.mask_token)
        decoded = self.decoder(encodings)

        loss = self.loss_fn(decoded, pixel_values, reduction="none")
        mask = repeat(
            mask_patches,
            "b z y x -> b (z pz) (y py) (x px)",
            pz=self.patch_size[0],
            py=self.patch_size[1],
            px=self.patch_size[2],
        )
        loss = (loss * mask).sum() / ((mask.sum() + 1e-5) * self.in_channels)

        return decoded, loss, mask
