# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/layers/02_embeddings.ipynb.

# %% auto 0
__all__ = ['RelativePositionEmbeddings', 'get_coords_grid', 'RelativePositionEmbeddings3D',
           'RelativePositionEmbeddings3DMetaNetwork', 'AbsolutePositionEmbeddings3D', 'PatchEmbeddings3D']

# %% ../../nbs/layers/02_embeddings.ipynb 2
from typing import Union

import numpy as np
import torch
import torch.nn.functional as F
from einops import rearrange, repeat
from torch import nn

from ..utils.normalizations import get_norm_layer

# %% ../../nbs/layers/02_embeddings.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/layers/02_embeddings.ipynb 6
class RelativePositionEmbeddings3D(nn.Module):
    def __init__(
        self,
        num_heads,
        num_patches: tuple[int, int, int],
    ):
        super().__init__()

        self.num_heads = num_heads
        self.num_patches = num_patches

        # TODO: Add embed_spacing_info functionality

        relative_limits = (2 * num_patches[0] - 1, 2 * num_patches[1] - 1, 2 * num_patches[2] - 1)

        self.relative_position_bias_table = nn.Parameter(torch.randn(num_heads, np.prod(relative_limits)))
        # (num_heads, num_patches_z * num_patches_y * num_patches_x)

        # Pair-wise relative position index for each token inside the window
        coords = get_coords_grid(num_patches)
        coords_flatten = rearrange(coords, "three d h w -> three (d h w)", three=3)
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += num_patches[0] - 1
        relative_coords[:, :, 1] += num_patches[1] - 1
        relative_coords[:, :, 2] += num_patches[2] - 1
        relative_position_index: torch.Tensor = (
            relative_coords[:, :, 0] * relative_limits[1] * relative_limits[2]
            + relative_coords[:, :, 1] * relative_limits[2]
            + relative_coords[:, :, 2]
        )
        self.relative_position_index = relative_position_index.flatten()
        # (num_patches, num_patches)

    def forward(self):
        relative_position_embeddings = self.relative_position_bias_table[:, self.relative_position_index]
        # (num_heads, num_patches, num_patches)
        relative_position_embeddings = relative_position_embeddings.reshape(
            1, np.prod(self.num_patches), np.prod(self.num_patches), -1
        )
        # (1, num_patches, num_patches, num_heads)
        relative_position_embeddings = relative_position_embeddings.permute(0, 3, 1, 2).contiguous()
        # (1, num_heads, num_patches, num_patches)
        return relative_position_embeddings

# %% ../../nbs/layers/02_embeddings.ipynb 8
class RelativePositionEmbeddings3DMetaNetwork(nn.Module):
    def __init__(
        self,
        num_heads,
        num_patches: tuple[int, int, int],
    ):
        super().__init__()

        self.num_heads = num_heads
        self.num_patches = num_patches

        # TODO: Add embed_spacing_info functionality
        self.cpb_mlp = nn.Sequential(
            nn.Linear(3, 512, bias=True),
            nn.ReLU(inplace=True),
            nn.Linear(512, num_heads, bias=False),
        )

        relative_limits = (2 * num_patches[0] - 1, 2 * num_patches[1] - 1, 2 * num_patches[2] - 1)

        # Relative coordinates table
        relative_coords_table = get_coords_grid(relative_limits).float()
        relative_coords_table[0] -= num_patches[0] - 1
        relative_coords_table[1] -= num_patches[1] - 1
        relative_coords_table[2] -= num_patches[2] - 1
        relative_coords_table = relative_coords_table.permute(1, 2, 3, 0).contiguous()
        relative_coords_table[:, :, :, 0] /= self.num_patches[0] - 1
        relative_coords_table[:, :, :, 1] /= self.num_patches[1] - 1
        relative_coords_table[:, :, :, 2] /= self.num_patches[2] - 1
        relative_coords_table *= 8  # Normalize to -8, 8
        relative_coords_table = (
            torch.sign(relative_coords_table) * torch.log2(torch.abs(relative_coords_table) + 1.0) / np.log2(8)
        )
        # (num_patches_z, num_patches_y, num_patches_x, 3)
        # Allow moving this to and from cuda whenever required but don't save to state_dict
        self.register_buffer("relative_coords_table", relative_coords_table, persistent=False)

        # Pair-wise relative position index for each token inside the window
        coords = get_coords_grid(num_patches)
        coords_flatten = rearrange(coords, "three d h w -> three (d h w)", three=3)
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += num_patches[0] - 1
        relative_coords[:, :, 1] += num_patches[1] - 1
        relative_coords[:, :, 2] += num_patches[2] - 1
        relative_position_index: torch.Tensor = (
            relative_coords[:, :, 0] * relative_limits[1] * relative_limits[2]
            + relative_coords[:, :, 1] * relative_limits[2]
            + relative_coords[:, :, 2]
        )
        self.relative_position_index = relative_position_index.flatten()
        # (num_patches, num_patches)

    def forward(self):
        # (num_patches_z, num_patches_y, num_patches_x, 3)
        relative_position_embeddings_table = self.cpb_mlp(self.relative_coords_table)
        # (num_patches_z, num_patches_y, num_patches_x, num_heads)
        relative_position_embeddings_table = relative_position_embeddings_table.reshape(-1, self.num_heads)
        # (num_patches, num_heads)
        relative_position_embeddings = relative_position_embeddings_table[self.relative_position_index]
        # (num_patches * num_patches, num_heads)
        relative_position_embeddings = rearrange(
            relative_position_embeddings,
            "(num_patches1 num_patches2) num_heads -> num_heads num_patches1 num_patches2",
            num_patches1=np.prod(self.num_patches),
            num_patches2=np.prod(self.num_patches),
            num_heads=self.num_heads,
        ).contiguous()
        # (num_heads, num_patches, num_patches)
        relative_position_embeddings = 16 * torch.sigmoid(relative_position_embeddings)
        # (num_heads, num_patches, num_patches)
        return relative_position_embeddings

# %% ../../nbs/layers/02_embeddings.ipynb 10
RelativePositionEmbeddings = Union[RelativePositionEmbeddings3D, RelativePositionEmbeddings3DMetaNetwork]

# %% ../../nbs/layers/02_embeddings.ipynb 11
class AbsolutePositionEmbeddings3D(nn.Module):
    def __init__(self, dim, grid_size, learnable=False, spacing=(1, 1, 1)):
        super().__init__()

        if dim % 6 != 0:
            raise ValueError("embed_dim must be divisible by 6")

        self.dim = dim
        self.grid_size = grid_size
        self.spacing = spacing
        grid = get_coords_grid(self.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(self.dim // 6, dtype=torch.float32)
        omega /= self.dim / 6.0
        omega = 1.0 / 10000**omega
        # (dim // 6)

        patch_multiplier = torch.Tensor(self.spacing) / min(self.spacing)

        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)
        # (dim, d * h * w)
        d, h, w = self.grid_size
        position_embeddings = rearrange(position_embeddings, "(d h w) e -> 1 e d h w", d=d, h=h, w=w)
        # (1, dim, d, h, w)

        self.position_embeddings = nn.Parameter(position_embeddings, requires_grad=learnable)

    def forward(self, batch_size=None, spacings=None):
        assert batch_size is not None or spacings is not None, "Either batch_size or spacings must be provided"

        position_embeddings = self.position_embeddings
        # (1, dim, d, h, w)

        if batch_size is not None:
            b = batch_size
        else:
            assert spacings.ndim == 2 and spacings.shape[1] == 3, "spacings must be of shape (batch_size, 3)"

            dim = position_embeddings.shape[1]
            assert dim % 3 == 0, "embed_dim must be divisible by 3"

            b = spacings.shape[0]

        position_embeddings = repeat(position_embeddings, "1 e d h w -> b e d h w", b=b)

        if spacings is not None:
            # (b, 3)
            spacings = repeat(spacings, "b three -> b (three dim_by_three) 1 1 1", three=3, dim_by_three=dim // 3)
            # (b, dim, 1, 1, 1)

            position_embeddings = position_embeddings * spacings
            # (b, dim, d, h, w)

        return position_embeddings

# %% ../../nbs/layers/02_embeddings.ipynb 14
class PatchEmbeddings3D(nn.Module):
    def __init__(self, patch_size: tuple[int, int, int], in_channels: int, dim: int, norm_layer= 'layernorm'):
        super().__init__()

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

        if norm_layer is None:
            self.normalization = nn.Identity()
        elif isinstance(norm_layer, nn.Module):
            self.normalization = norm_layer(dim)
        else:
            self.normalization = get_norm_layer(norm_layer, dim)

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

        embeddings = self.patch_embeddings(pixel_values)
        # (b, dim, num_patches_z, num_patches_y, num_patches_x)
        embeddings = rearrange(embeddings, "b d z y x -> b z y x d")
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)
        embeddings = self.normalization(embeddings)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        return embeddings
