# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/nets/06_detr_3d.ipynb.

# %% auto 0
__all__ = ['DETR3DEncoderConfig', 'DETR3DDecoderConfig', 'DETRBBoxMLPConfig', 'DETR3DConfig', 'DETR3DEncoderLayer',
           'DETR3DDecoderLayer', 'DETR3DEncoder', 'DETR3DDecoder', 'DETRBBoxMLP', 'DETR3D']

# %% ../../nbs/nets/06_detr_3d.ipynb 2
from collections.abc import Callable
from functools import partial, wraps
from typing import Literal

import torch
from einops import rearrange, repeat
from huggingface_hub import PyTorchModelHubMixin
from monai.data.box_utils import box_giou, box_iou, box_pair_giou, convert_box_mode
from scipy.optimize import linear_sum_assignment
from torch import nn
from torch.nn import functional as F
from torch.nn.utils.rnn import pad_sequence

from vision_architectures.blocks.transformer import (
    Attention1D,
    Attention1DMLP,
    Attention3D,
    Attention3DMLP,
    TransformerDecoderBlock1DConfig,
    TransformerEncoderBlock3DConfig,
)
from ..docstrings import populate_docstring
from ..layers.embeddings import AbsolutePositionEmbeddings3D, AbsolutePositionEmbeddings3DConfig
from vision_architectures.losses.class_balanced_cross_entropy_loss import (
    ClassBalancedCrossEntropyLoss,
    ClassBalancedCrossEntropyLossConfig,
)
from ..utils.activation_checkpointing import ActivationCheckpointing
from ..utils.activations import get_act_layer
from ..utils.custom_base_model import CustomBaseModel, Field, model_validator
from ..utils.rearrange import rearrange_channels
from ..utils.residuals import Residual

# %% ../../nbs/nets/06_detr_3d.ipynb 4
class DETR3DEncoderConfig(TransformerEncoderBlock3DConfig, AbsolutePositionEmbeddings3DConfig):
    num_encoder_layers: int = Field(..., description="Number of transformer encoder layers.")


class DETR3DDecoderConfig(TransformerDecoderBlock1DConfig, AbsolutePositionEmbeddings3DConfig):
    num_decoder_layers: int = Field(..., description="Number of transformer decoder layers.")


class DETRBBoxMLPConfig(CustomBaseModel):
    dim: int = Field(..., description="Dimension of the input features.")
    num_classes: int = Field(..., description="Number of classes for the bounding box predictions.")
    activation: str = Field("relu", description="Activation function for the MLP.")
    bbox_size_activation: Literal["sigmoid", "softplus"] = Field(
        "sigmoid",
        description=(
            'Activation function for the bounding box size. "sigmoid" for normalized coordinates, "softplus" '
            "for absolute coordinates."
        ),
    )

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        assert self.num_classes > 0, "Number of classes must be greater than zero. If not predicting classes, set to 1."
        return self


class DETR3DConfig(DETR3DEncoderConfig, DETR3DDecoderConfig, DETRBBoxMLPConfig):
    num_objects: int = Field(..., description="Maximum number of objects to detect.")
    classification_cost_fn: Literal["softmax", "log_softmax"] = Field(
        "softmax",
        description=(
            "Method used to compute classification cost in Hungarian matching. While softmax is the default (as per "
            "the paper), log_softmax is encouraged as the final bipartite matching loss calculation uses cross entropy "
            "loss which performs log_softmax and NLLLoss internally making log_softmax more apt."
        ),
    )
    classification_loss_fn: (
        Literal["cross_entropy", "class_balanced_cross_entropy"] | dict | ClassBalancedCrossEntropyLossConfig
    ) = Field("cross_entropy", description="Loss function for bbox classification.")

    @model_validator(mode="after")
    def validate(self):
        super().validate()

        # Create an object of ClassBalancedCrossEntropyLossConfig if applicable
        if self.classification_loss_fn == "class_balanced_cross_entropy":
            self.classification_loss_fn = ClassBalancedCrossEntropyLossConfig(num_classes=self.num_classes + 1)
        elif isinstance(self.classification_loss_fn, dict):
            self.classification_loss_fn = ClassBalancedCrossEntropyLossConfig.model_validate(
                self.classification_loss_fn
            )

        return self

# %% ../../nbs/nets/06_detr_3d.ipynb 7
@populate_docstring
class DETR3DEncoderLayer(nn.Module):
    """DETR3D Encoder layer. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(
        self,
        config: TransformerEncoderBlock3DConfig = {},
        checkpointing_level: int = 0,
        **kwargs,
    ):
        """Initialize a DETR3DDecoderLayer block. Activation checkpointing level 3.

        This transformer encoder layer allows modifying query and key values independent of the residual connection.

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

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

        dim_qk = self.config.dim_qk
        layer_norm_eps = self.config.layer_norm_eps

        self.attn = Attention3D(
            self.config,
            checkpointing_level=checkpointing_level,
        )
        self.layernorm1 = nn.LayerNorm(dim_qk, eps=layer_norm_eps)
        self.mlp = Attention3DMLP(self.config, checkpointing_level=checkpointing_level)
        self.layernorm2 = nn.LayerNorm(dim_qk, eps=layer_norm_eps)

        self.residual = Residual()

        self.checkpointing_level3 = ActivationCheckpointing(3, checkpointing_level)

    @populate_docstring
    def _forward(
        self,
        qkv: torch.Tensor,
        q_modifier: Callable | None = None,
        k_modifier: Callable | None = None,
        channels_first: bool = True,
    ) -> torch.Tensor:
        """Forward pass of the Attention3DWithMLP block.

        Args:
            qkv: {INPUT_3D_DOC}
            q_modifier: If provided, the query will be passed to the callable before the attention operation.
            k_modifier: If provided, the key will be passed to the callable before the attention operation.
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        # Each is (b, [dim], tokens_z, tokens_y, tokens_x, [dim])

        qkv = rearrange_channels(qkv, channels_first, False)
        # (b, tokens_z, tokens_y, tokens_x, dim)

        q, k, v = qkv, qkv, qkv

        res_connection1 = q
        # (b, tokens_z, tokens_y, tokens_x, dim)

        if q_modifier is not None:
            q = q_modifier(q)
        if k_modifier is not None:
            k = k_modifier(k)
        # (b, tokens_z, tokens_y, tokens_x, dim)

        if self.config.norm_location == "pre":
            q = self.layernorm1(q)
            k = self.layernorm1(k)
            v = self.layernorm1(v)
            # (b, tokens_z, tokens_y, tokens_x, dim)

        hidden_states = self.attn(q, k, v, channels_first=False)
        # (b, tokens_z, tokens_y, tokens_x, dim)

        if self.config.norm_location == "post":
            hidden_states = self.layernorm1(hidden_states)
            # (b, tokens_z, tokens_y, tokens_x, dim)

        hidden_states = self.residual(hidden_states, res_connection1)
        res_connection2 = hidden_states
        # (b, tokens_z, tokens_y, tokens_x, dim)

        if self.config.norm_location == "pre":
            hidden_states = self.layernorm2(hidden_states)
            # (b, tokens_z, tokens_y, tokens_x, dim)

        hidden_states = self.mlp(hidden_states, channels_first=False)
        # (b, tokens_z, tokens_y, tokens_x, dim)

        if self.config.norm_location == "post":
            hidden_states = self.layernorm2(hidden_states)
            # (b, tokens_z, tokens_y, tokens_x, dim)

        hidden_states = self.residual(hidden_states, res_connection2)
        # (b, tokens_z, tokens_y, tokens_x, dim)

        hidden_states = rearrange_channels(hidden_states, False, channels_first)
        # (b, [dim], tokens_z, tokens_y, tokens_x, [dim])

        return hidden_states

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

# %% ../../nbs/nets/06_detr_3d.ipynb 9
@populate_docstring
class DETR3DDecoderLayer(nn.Module):
    """A cross attention transformer block. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: TransformerDecoderBlock1DConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize a DETR3DDecoderLayer block. Activation checkpointing level 3.

        This transformer decoder layer allows modifying query and key values independent of the residual connection.

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

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

        dim = self.config.dim
        num_heads = self.config.num_heads
        mlp_ratio = self.config.mlp_ratio
        layer_norm_eps = self.config.layer_norm_eps
        attn_drop_prob = self.config.attn_drop_prob
        proj_drop_prob = self.config.proj_drop_prob
        mlp_drop_prob = self.config.mlp_drop_prob

        self.attn1 = Attention1D(
            dim=dim,
            num_heads=num_heads,
            attn_drop_prob=attn_drop_prob,
            proj_drop_prob=proj_drop_prob,
        )
        self.layernorm1 = nn.LayerNorm(dim, eps=layer_norm_eps)
        self.attn2 = Attention1D(
            dim=dim,
            num_heads=num_heads,
            attn_drop_prob=attn_drop_prob,
            proj_drop_prob=proj_drop_prob,
        )
        self.layernorm2 = nn.LayerNorm(dim, eps=layer_norm_eps)
        self.mlp = Attention1DMLP(dim=dim, mlp_ratio=mlp_ratio, mlp_drop_prob=mlp_drop_prob)
        self.layernorm3 = nn.LayerNorm(dim, eps=layer_norm_eps)

        self.residual = Residual()

        self.checkpointing_level3 = ActivationCheckpointing(3, checkpointing_level)

    @populate_docstring
    def _forward(
        self,
        q: torch.Tensor,
        kv: torch.Tensor,
        q1_modifier: Callable | None = None,
        k1_modifier: Callable | None = None,
        q2_modifier: Callable | None = None,
        k2_modifier: Callable | None = None,
        channels_first: bool = True,
    ) -> torch.Tensor:
        """Forward pass of the TransformerDecoderBlock1D block.

        Args:
            q: {INPUT_1D_DOC} The query, key, and value tensor used for self-attention.
            kv: {INPUT_3D_DOC} The key and value tensors used for cross-attention.
            q1_modifier: If provided, the self-attention query tensor will be passed through the callable before the
                attention operation.
            k1_modifier: If provided, the self-attention key tensor will be passed through the callable before the
                attention operation.
            q2_modifier: If provided, the cross-attention query tensor will be passed through the callable before the
                attention operation.
            k2_modifier: If provided, the cross-attention key tensor will be passed through the callable before the
                attention operation.
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        # q: (b, num_possible_objects, dim)
        # kv: (b, [dim], tokens_z, tokens_y, tokens_x, [dim])

        q1, k1, v1 = q, q, q
        k2, v2 = kv, kv

        k2 = rearrange_channels(k2, channels_first, False)
        v2 = rearrange_channels(v2, channels_first, False)
        # (b, tokens_z, tokens_y, tokens_x, dim)

        res_connection1 = q1
        # (b, num_possible_objects, dim)

        if q1_modifier is not None:
            q1 = q1_modifier(q1)
        if k1_modifier is not None:
            k1 = k1_modifier(k1)
        # (b, num_possible_objects, dim)

        if self.config.norm_location == "pre":
            q1 = self.layernorm1(q1)
            k1 = self.layernorm1(k1)
            v1 = self.layernorm1(v1)
            # (b, num_possible_objects, dim)

        q2: torch.Tensor = self.attn1(q1, k1, v1)
        # (b, num_possible_objects, dim)

        if self.config.norm_location == "post":
            q2 = self.layernorm1(q2)
            # (b, num_possible_objects, dim)

        q2 = self.residual(q2, res_connection1)
        res_connection2 = q2
        # (b, num_possible_objects, dim)

        if q2_modifier is not None:
            q2 = q2_modifier(q2)
            # (b, num_possible_objects, dim)
        if k2_modifier is not None:
            k2 = k2_modifier(k2)
            # (b, tokens_z, tokens_y, tokens_x, dim)

        k2 = rearrange(k2, "b z y x d -> b (z y x) d").contiguous()
        v2 = rearrange(v2, "b z y x d -> b (z y x) d").contiguous()
        # (b, num_possible_objects, dim)

        if self.config.norm_location == "pre":
            q2 = self.layernorm2(q2)
            # (b, num_possible_objects, dim)

        hidden_states = self.attn2(q2, k2, v2)
        # (b, num_possible_objects, dim)

        if self.config.norm_location == "post":
            hidden_states = self.layernorm2(hidden_states)
            # (b, num_possible_objects, dim)

        hidden_states = self.residual(hidden_states, res_connection2)
        res_connection3 = hidden_states
        # (b, num_possible_objects, dim)

        if self.config.norm_location == "pre":
            hidden_states = self.layernorm3(hidden_states)
            # (b, num_possible_objects, dim)

        hidden_states = self.mlp(hidden_states)
        # (b, num_possible_objects, dim)

        if self.config.norm_location == "post":
            hidden_states = self.layernorm3(hidden_states)
            # (b, num_possible_objects, dim)

        hidden_states = self.residual(hidden_states, res_connection3)
        # (b, num_possible_objects, dim)

        return hidden_states

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

# %% ../../nbs/nets/06_detr_3d.ipynb 12
class DETR3DEncoder(nn.Module, PyTorchModelHubMixin):
    """DETR Transformer encoder."""

    @populate_docstring
    def __init__(self, config: DETR3DEncoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the DETR3DEncoder. Activation checkpointing level 4.

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

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

        self.position_embeddings = AbsolutePositionEmbeddings3D(config)
        self.layers = nn.ModuleList(
            [DETR3DEncoderLayer(config, checkpointing_level) for _ in range(self.config.num_encoder_layers)]
        )

        self.checkpointing_level4 = ActivationCheckpointing(4, checkpointing_level)

    @populate_docstring
    def _forward(
        self,
        embeddings: torch.Tensor,
        spacings: torch.Tensor | None = None,
        return_intermediates: bool = False,
        channels_first: bool = True,
    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
        """Forward pass of the DETR3D encoder.

        Args:
            embeddings: Embeddings from the backbone. {INPUT_3D_DOC}
            spacings: {SPACINGS_DOC}
            return_intermediates: If True, also returns the outputs of all layers. Defaults to False.
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            If return_intermediates is True, returns the final object embeddings and a list of outputs from all layers.
            Otherwise, returns only the final object embeddings.
        """
        # (b, [dim], tokens_z, tokens_y, tokens_x, [dim])

        embeddings = rearrange_channels(embeddings, channels_first, False)
        # (b, tokens_z, tokens_y, tokens_x, dim)

        position_embeddings_modifier = partial(self.position_embeddings, spacings=spacings, channels_first=False)

        layer_outputs = []
        for layer in self.layers:
            embeddings = layer(
                qkv=embeddings,
                q_modifier=position_embeddings_modifier,
                k_modifier=position_embeddings_modifier,
                channels_first=False,
            )
            layer_outputs.append(embeddings)

        embeddings = rearrange_channels(embeddings, False, channels_first)

        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/06_detr_3d.ipynb 15
class DETR3DDecoder(nn.Module, PyTorchModelHubMixin):
    """DETR Transformer decoder."""

    @populate_docstring
    def __init__(self, config: DETR3DDecoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the DETR3DDecoder. Activation checkpointing level 4.

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

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

        self.position_embeddings = AbsolutePositionEmbeddings3D(config)
        self.layers = nn.ModuleList(
            [DETR3DDecoderLayer(config, checkpointing_level) for _ in range(self.config.num_decoder_layers)]
        )

        self.checkpointing_level4 = ActivationCheckpointing(4, checkpointing_level)

    @populate_docstring
    def _forward(
        self,
        object_queries: torch.Tensor,
        embeddings: torch.Tensor,
        spacings: torch.Tensor | None = None,
        return_intermediates: bool = False,
        channels_first: bool = True,
    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
        """Forward pass of the DETR3D decoder.

        Args:
            object_queries: Tokens that represent object queries. {INPUT_1D_DOC}
            embeddings: Actual embeddings of the input. {INPUT_3D_DOC}
            spacings: {SPACINGS_DOC}
            return_intermediates: If True, also returns the outputs of all layers. Defaults to False.
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            If return_intermediates is True, returns the final object embeddings and a list of outputs from all layers.
            Otherwise, returns only the final object embeddings.
        """
        # object_queries: (b, num_possible_objects, dim)
        # embeddings: (b, [dim], tokens_z, tokens_y, tokens_x, [dim])

        embeddings = rearrange_channels(embeddings, channels_first, False)
        # (b, tokens_z, tokens_y, tokens_x, dim)

        position_embeddings_modifier = partial(self.position_embeddings, spacings=spacings, channels_first=False)

        def add_object_queries(ten: torch.Tensor, iteration: int) -> torch.Tensor:
            if iteration == 0:
                return ten  # don't add object queries for the first decoder layer
            return ten + object_queries

        object_embeddings = object_queries
        layer_outputs = []
        for i, layer in enumerate(self.layers):
            object_queries_modifier = partial(add_object_queries, iteration=i)
            object_embeddings = layer(
                q=object_embeddings,
                kv=embeddings,
                q1_modifier=object_queries_modifier,
                k1_modifier=object_queries_modifier,
                q2_modifier=object_queries_modifier,
                k2_modifier=position_embeddings_modifier,
                channels_first=False,
            )
            layer_outputs.append(object_embeddings)

        if return_intermediates:
            return object_embeddings, layer_outputs
        return object_embeddings

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

# %% ../../nbs/nets/06_detr_3d.ipynb 18
class DETRBBoxMLP(nn.Module):
    """DETR Bounding Box MLP. This module predicts bounding boxes and class scores from object query embeddings."""

    @populate_docstring
    def __init__(self, config: DETRBBoxMLPConfig = {}, **kwargs):
        """Initialize the DETRBBoxMLP.

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

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

        self.bbox_head = nn.Sequential(
            nn.Linear(self.config.dim, self.config.dim),
            get_act_layer(self.config.activation),
            nn.Linear(self.config.dim, self.config.dim),
            get_act_layer(self.config.activation),
            nn.Linear(self.config.dim, 6),
        )
        self.class_head = nn.Linear(self.config.dim, 1 + self.config.num_classes)

    @populate_docstring
    def forward(
        self,
        object_embeddings: torch.Tensor,
    ) -> torch.Tensor:
        """Forward pass of the DETRBBoxMLP.

        Args:
            object_embeddings: Object embeddings from the DETR decoder. {INPUT_1D_DOC}

        Returns:
            A tensor of shape (b, num_possible_objects, 1 objectness class + 6 bounding box parameters + num_classes)
            containing the predicted bounding boxes and class scores. {BOUNDING_BOXES_FORMAT_DOC}
        """
        # object_embeddings: (b, num_possible_objects, dim)

        bbox_parameters: torch.Tensor = self.bbox_head(object_embeddings)
        class_logits: torch.Tensor = self.class_head(object_embeddings)
        # (b, num_possible_objects, 6 + 1 + num_classes)

        # Bound the bounding box parameters
        centers, sizes = bbox_parameters[:, :, :3], bbox_parameters[:, :, 3:6]
        centers = centers.sigmoid()  # center coordinates should be in the bbox
        if self.config.bbox_size_activation == "sigmoid":
            sizes = sizes.sigmoid()  # sizes should be between 0 and 1
        elif self.config.bbox_size_activation == "softplus":
            sizes = F.softplus(sizes)  # sizes should be positive but unbounded
        bboxes = torch.cat([centers, sizes, class_logits], dim=-1)

        return bboxes

# %% ../../nbs/nets/06_detr_3d.ipynb 21
class DETR3D(nn.Module, PyTorchModelHubMixin):
    """DETR 3D model. Also implements bipartite matching loss which is essential for DETR training."""

    @populate_docstring
    def __init__(self, config: DETR3DConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the DETR3D. Activation checkpointing level 4.

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

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

        self.num_possible_objects = self.config.num_objects
        self.encoder = DETR3DEncoder(config, checkpointing_level)
        self.object_queries = nn.Parameter(torch.randn(1, self.num_possible_objects, self.config.dim) * 0.01)
        self.decoder = DETR3DDecoder(config, checkpointing_level)
        self.bbox_mlp = DETRBBoxMLP(config)

        self.class_balanced_cross_entropy_loss = None
        if isinstance(self.config.classification_loss_fn, ClassBalancedCrossEntropyLossConfig):
            self.class_balanced_cross_entropy_loss = ClassBalancedCrossEntropyLoss(
                config=self.config.classification_loss_fn
            )

        self.checkpointing_level4 = ActivationCheckpointing(4, checkpointing_level)

    @populate_docstring
    def _forward(
        self,
        embeddings: torch.Tensor,
        spacings: torch.Tensor | None = None,
        channels_first: bool = True,
        return_intermediates: bool = False,
        process_intermediates: bool = True,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor], list[torch.Tensor]]:
        """Forward pass of the DETR3D.

        Args:
            embeddings: Input features from backbone. {INPUT_3D_DOC}
            spacings: {SPACINGS_DOC}
            channels_first: {CHANNELS_FIRST_DOC}
            return_intermediates: If True, also returns the outputs of all layers. Defaults to False.
            process_intermediates: If True, passes the layer outputs through the bbox_mlp too. Requires
                `return_intermediates` to be True too.

        Returns:
            A tuple containing bounding boxes, object embeddings, decoder layer outputs, and encoder layer outputs if
            return_intermediates is True. Else, returns only the bounding boxes. {BOUNDING_BOXES_FORMAT_DOC}
        """
        # embeddings: (b, [dim], num_tokens_z, num_tokens_y, num_tokens_x, [dim])
        # spacings: (b, 3)

        embeddings = rearrange_channels(embeddings, channels_first, False)
        # (b, num_tokens_z, num_tokens_y, num_tokens_x, dim)

        embeddings, encoder_layer_outputs = self.encoder(
            embeddings, spacings=spacings, return_intermediates=True, channels_first=False
        )
        # embeddings: (b, num_tokens_z, num_tokens_y, num_tokens_x, dim)
        # encoder_layer_outputs: list of (b, num_tokens_z, num_tokens_y, num_tokens_x, dim)

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

        object_embeddings, decoder_layer_outputs = self.decoder(
            object_queries, embeddings, spacings=spacings, return_intermediates=True, channels_first=False
        )
        # object_embeddings: (b, num_possible_objects, dim)
        # decoder_layer_outputs: list of (b, num_possible_objects, dim)

        bboxes = self.bbox_mlp(object_embeddings)
        # (b, num_possible_objects, 6 + 1 + num_classes)

        if return_intermediates:
            if process_intermediates:
                for i in range(len(decoder_layer_outputs)):
                    decoder_layer_outputs[i] = self.bbox_mlp(decoder_layer_outputs[i])
            return bboxes, object_embeddings, decoder_layer_outputs, encoder_layer_outputs

        return bboxes

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

    def bipartite_matching_loss(
        self,
        pred: torch.Tensor,
        target: torch.Tensor | list[torch.Tensor],
        intermediate_preds: list[torch.Tensor] | None = None,
        classification_cost_weight: float = 1.0,
        bbox_l1_cost_weight: float = 1.0,
        bbox_giou_cost_weight: float = 1.0,
        matched_indices: tuple[list[int], list[int]] | None = None,
        update_class_prevalences: bool = True,
        reduction: str = "mean",
        return_matching: bool = False,
        return_loss_components: bool = False,
    ) -> (
        torch.Tensor  # just loss
        | list[torch.Tensor]  # main loss plus auxiliary losses
        | tuple[torch.Tensor, tuple]  # loss with matching
        | tuple[list[torch.Tensor], tuple]  # main+auxiliary losses with matching
        | tuple[torch.Tensor, dict[str, torch.Tensor]]  # loss with loss components
        | tuple[list[torch.Tensor], list[dict[str, torch.Tensor]]]  # main+auxiliary losses with loss components
        | tuple[torch.Tensor, tuple, dict[str, torch.Tensor]]  # loss, matching, loss components
        | tuple[list[torch.Tensor], tuple, list[dict[str, torch.Tensor]]]  # all losses, matching, all loss components
    ):
        """Bipartite matching loss for DETR. The classes are expected to optimize for a multi-class classification
        problem. Expects raw logits in class predictions, not probabilities. Use ``logits_to_scores_fn=None`` in the
        ``forward`` function to avoid applying any transformation.

        Args:
            pred: Predicted bounding boxes and class scores. It should be of shape
                `(B, num_objects, 6 + 1 + num_classes)`. Number of objects and number of classes will be inferred from
                here.
            target: Target bounding boxes and class scores. If provided as a list, each element should be a tensor for
                the corresponding batch element in ``pred`` and therefore should have a length of `B`. Each tensor
                should have less than or equal to the number of objects in `pred`. The number of classes can either be
                the exact same as in `pred`, or it should be 1 argmax (one-cold) decoding.
            intermediate_preds: A list of any intermediate decoder outputs to use for auxiliary losses. If provided,
                returned loss values become a list instead of single values. The matching used is the one used for the
                final prediction i.e. `pred`.
            classification_cost_weight: Weight for the classification cost in hungarian matching. Only used when
                matched_indices is None.
            bbox_l1_cost_weight: Weight for the bounding box L1 loss cost in hungarian matching. Only used when
                matched_indices is None.
            bbox_giou_cost_weight: Weight for the bounding box IoU cost in hungarian matching. Only used when
                matched_indices is None.
            matched_indices: Hungarian matching is done only if this is None, else this is used to match pred and target
                boxes. Useful when calculating auxiliary losses with intermediate layers of DETR3D.
            update_class_prevalences: Whether or not to update class prevalences in the class balanced cross entropy
                loss. Useful when calculating auxiliary losses with intermediate layers of DETR3D.
            reduction: Specifies the reduction to apply to the output.
            return_matching: Whether or not to return the matched indices from the bipartite matching.
            return_loss_components: Whether or not to return the individual loss components.

        Returns:
            The first element is always the main loss that has been calculated. If no `intermediate_preds` were
            provided, a tensor containing the bipartite matching loss with the shape depending on the `reduction`
            argument is present, else it is a list of the same for all preds provided with the first one being the main
            pred.

            If `return_matching` is True, also returns a list of tuples containing matched indices for predictions and
            targets. Each tuple is of the form `(pred_indices, target_indices)`, where `pred_indices` and
            `target_indices` are lists of indices for the matched predictions and targets, respectively.

            If `return_loss_components` is True, also returns a dict of each loss component reduced based on
            `reduction`. If `intermediate_preds` were also provided, it is a list of the same.
        """
        B = pred.shape[0]

        # Convert target to a list of tensors if not already
        if isinstance(target, torch.Tensor):
            target = list(target)

        # argmax encode the class labels if they are not already
        for i in range(len(target)):
            if target[i].shape[-1] > 7:  # 6 bbox + 1 class
                target[i] = torch.cat([target[i][:, :6], target[i][:, 6:].argmax(-1, keepdims=True)], dim=-1)

        # Perform hungarian matching
        if matched_indices is None:
            matched_indices = self.hungarian_matching(
                pred, target, classification_cost_weight, bbox_l1_cost_weight, bbox_giou_cost_weight
            )

        # Update class prevalences in class balanced cross entropy loss
        if update_class_prevalences:
            self._update_class_prevalences(target)

        # Get a combined list of all preds
        all_preds = [pred]
        if intermediate_preds is not None:
            all_preds += intermediate_preds

        # Calculate losses for all preds
        loss = []
        loss_components = []
        for pred in all_preds:  # pred is overwritten here
            batch_loss_components = {
                "classification_loss": [],
                "bbox_l1_loss": [],
                "bbox_giou_loss": [],
                "total_loss": [],
            }
            for i in range(B):
                batch_losses = {}

                pred_indices, target_indices = matched_indices[i]

                # Classification loss for ALL predictions
                all_class_labels = torch.zeros_like(pred[i][:, 6], dtype=torch.long)
                all_class_labels[pred_indices] = target[i][target_indices, 6].long()
                all_pred_classes = pred[i][..., 6:]
                if self.class_balanced_cross_entropy_loss is None:
                    batch_losses["classification_loss"] = F.cross_entropy(all_pred_classes, all_class_labels)
                else:
                    batch_losses["classification_loss"] = self.class_balanced_cross_entropy_loss(
                        all_pred_classes, all_class_labels, update_class_prevalences=False
                    )

                # Calculate all other losses only if prediction and target have objects
                if pred_indices != [] and target_indices != []:
                    matched_pred = pred[i][pred_indices]
                    matched_target = target[i][target_indices]

                    pred_bboxes = matched_pred[:, :6].clone()
                    target_bboxes = matched_target[:, :6].clone()

                    # Compute losses for matched pairs
                    # BBox L1 loss
                    batch_losses["bbox_l1_loss"] = F.l1_loss(pred_bboxes, target_bboxes)

                    # For IOU calculation, it does not matter if bboxes are in actual pixel lengths or normalized based on
                    # image size, metric value will be the same.

                    # BBox IOU loss
                    batch_losses["bbox_giou_loss"] = 1 - DETR3D.bbox_giou(pred_bboxes, target_bboxes).mean()

                # Total loss for this batch element
                total_loss = (
                    classification_cost_weight * batch_losses.get("classification_loss", 0.0)
                    + bbox_l1_cost_weight * batch_losses.get("bbox_l1_loss", 0.0)
                    + bbox_giou_cost_weight * batch_losses.get("bbox_giou_loss", 0.0)
                )

                # Add to losses
                batch_loss_components["classification_loss"].append(
                    batch_losses.get("classification_loss", torch.tensor(torch.nan, device=pred.device)),
                )
                batch_loss_components["bbox_l1_loss"].append(
                    batch_losses.get("bbox_l1_loss", torch.tensor(torch.nan, device=pred.device)),
                )
                batch_loss_components["bbox_giou_loss"].append(
                    batch_losses.get("bbox_giou_loss", torch.tensor(torch.nan, device=pred.device)),
                )
                batch_loss_components["total_loss"].append(total_loss)

            # Stack batch losses and apply reduction
            for key in batch_loss_components:
                batch_loss_component = torch.stack(batch_loss_components[key])

                if reduction == "mean":
                    batch_loss_component = batch_loss_component.nanmean()
                elif reduction == "sum":
                    batch_loss_component = batch_loss_component.nansum()
                elif reduction == "none" or reduction is None:
                    pass
                else:
                    raise ValueError(f"Invalid reduction mode: {reduction}")

                batch_loss_components[key] = batch_loss_component

            loss.append(batch_loss_components.pop("total_loss"))
            loss_components.append(batch_loss_components)

        # Finalize return value
        if len(loss) == 1:
            loss = loss[0]
            loss_components = loss_components[0]

        return_value = [loss]
        if return_matching:
            return_value.append(matched_indices)
        if return_loss_components:
            return_value.append(loss_components)

        if len(return_value) == 1:
            return_value = return_value[0]
        else:
            return_value = tuple(return_value)

        return return_value

    @torch.no_grad()
    def hungarian_matching(
        self,
        pred: torch.Tensor,
        target: list[torch.Tensor],
        classification_cost_weight: float = 1.0,
        bbox_l1_cost_weight: float = 1.0,
        bbox_giou_cost_weight: float = 1.0,
    ) -> list[tuple[list[int], list[int]]]:
        """Hungarian matching between predictions and targets.

        Args:
            pred: Predicted bounding boxes and class scores. It should be of shape
                `(B, num_objects, 6 + 1 + num_classes)`. Number of objects and number of classes will be inferred from
                here.
            target: Target bounding boxes and class scores. This is expected in argmax encoding or one-hot encoding.
            classification_cost_weight: Weight for the classification cost.
            bbox_l1_cost_weight: Weight for the bounding box L1 loss cost.
            bbox_giou_cost_weight: Weight for the bounding box IoU cost.

        Returns:
            A list of tuples containing matched indices for predictions and targets. Each tuple is of the form
            `(pred_indices, target_indices)`, where `pred_indices` and `target_indices` are lists of indices for the
            matched predictions and targets, respectively.
        """
        B = pred.shape[0]

        matched_indices = []
        for i in range(B):
            pred_bboxes = pred[i, :, :6]  # (num_objects, 6)
            target_bboxes = target[i][:, :6]  # (<=num_objects, 6)

            # If either prediction or target has no objects, skip matching
            if pred_bboxes.shape[0] == 0 or target_bboxes.shape[0] == 0:
                matched_indices.append(([], []))
                continue

            pred_class_logits = pred[i, :, 6:]  # (num_objects, num_classes)
            target_class_labels = target[i][:, 6].long()  # (<=num_objects,) this is in argmax encoding

            # ----- Cost matrix calculation -----

            # Classification cost
            if self.config.classification_cost_fn == "softmax":
                pred_class_probabilities = F.softmax(pred_class_logits, dim=-1)
                # (num_objects, num_classes)
                classification_cost = -pred_class_probabilities[:, target_class_labels]
                # (num_objects, <=num_objects)
            elif self.config.classification_cost_fn == "log_softmax":
                pred_class_log_probabilities = F.log_softmax(pred_class_logits, dim=-1)
                # (num_objects, num_classes)
                classification_cost = -pred_class_log_probabilities[:, target_class_labels]
                # (num_objects, <=num_objects)

            # Weigh by class-balancing weights (to align matching with the training objective)
            if self.class_balanced_cross_entropy_loss is not None:
                class_weights = self.class_balanced_cross_entropy_loss.get_class_weights(pred.device)  # (C,)
                target_weights = class_weights[target_class_labels]  # (<=num_objects,)
                classification_cost = classification_cost * target_weights  # broadcast across rows

            # L1 loss for bounding boxes
            bbox_l1_cost = torch.cdist(pred_bboxes, target_bboxes, p=1)
            # (num_objects, <=num_objects)

            # IOU cost for bounding boxes
            bbox_giou_cost = 1 - DETR3D.pairwise_bbox_giou(pred_bboxes, target_bboxes)
            # (num_objects, <=num_objects)

            # Total cost matrix
            cost_matrix = (
                classification_cost_weight * classification_cost
                + bbox_l1_cost_weight * bbox_l1_cost
                + bbox_giou_cost_weight * bbox_giou_cost
            )
            # (num_objects, <=num_objects)

            # Hungarian matching
            pred_indices_element, target_indices_element = linear_sum_assignment(cost_matrix.detach().cpu().numpy())

            matched_indices.append((list(pred_indices_element), list(target_indices_element)))

        return matched_indices

    @staticmethod
    @populate_docstring
    def bbox_iou(
        pred_bboxes: torch.Tensor,
        target_bboxes: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the IoU between two matched sets of bounding boxes.

        Args:
            pred_bboxes: Predicted bounding boxes of shape `(num_boxes, 6)`. {BOUNDING_BOXES_FORMAT_DOC}
            target_bboxes: Target bounding boxes of shape `(num_boxes, 6)`. {BOUNDING_BOXES_FORMAT_DOC}

        Returns:
            A tensor containing the IoU.
        """
        pred_corners = convert_box_mode(pred_bboxes, "cccwhd", "xyzxyz")
        target_corners = convert_box_mode(target_bboxes, "cccwhd", "xyzxyz")

        # Monai does not have a function to calculate pairwise IOUs

        # Intersection corners
        max_min = torch.max(pred_corners[:, :3], target_corners[:, :3])
        min_max = torch.min(pred_corners[:, 3:], target_corners[:, 3:])
        inter_dims = (min_max - max_min).clamp(min=0)
        inter_vol = inter_dims.prod(dim=1)

        # Volumes
        pred_dims = pred_corners[:, 3:] - pred_corners[:, :3]
        target_dims = target_corners[:, 3:] - target_corners[:, :3]
        pred_vol = pred_dims.prod(dim=1)
        target_vol = target_dims.prod(dim=1)
        union_vol = pred_vol + target_vol - inter_vol

        # IOU
        ious = inter_vol / union_vol.clamp(min=1e-5)

        return ious

    @staticmethod
    @populate_docstring
    def pairwise_bbox_iou(
        pred_bboxes: torch.Tensor,
        target_bboxes: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the IoUs between all combinations of predicted and target bounding boxes.

        Args:
            pred_bboxes: Predicted bounding boxes of shape `(num_objects, 6)`. {BOUNDING_BOXES_FORMAT_DOC}
            target_bboxes: Target bounding boxes of shape `(<=num_objects, 6)`. {BOUNDING_BOXES_FORMAT_DOC}

        Returns:
            A tensor containing the IoU losses of all combinations.
        """
        # Monai requires input in corner coordinates format
        pred_corners = convert_box_mode(pred_bboxes, "cccwhd", "xyzxyz")
        target_corners = convert_box_mode(target_bboxes, "cccwhd", "xyzxyz")

        # Calculate GIOUs
        gious = box_iou(pred_corners, target_corners)

        return gious

    @staticmethod
    @populate_docstring
    def bbox_giou(
        pred_bboxes: torch.Tensor,
        target_bboxes: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the Generalized IoU between two matched sets of bounding boxes.

        Args:
            pred_bboxes: Predicted bounding boxes of shape `(num_boxes, 6)`. {BOUNDING_BOXES_FORMAT_DOC}
            target_bboxes: Target bounding boxes of shape `(num_boxes, 6)`. {BOUNDING_BOXES_FORMAT_DOC}

        Returns:
            A tensor containing the IoU loss.
        """
        # Monai requires input in corner coordinates format
        pred_corners = convert_box_mode(pred_bboxes, "cccwhd", "xyzxyz")
        target_corners = convert_box_mode(target_bboxes, "cccwhd", "xyzxyz")

        # Calculate GIOUs
        gious = box_pair_giou(pred_corners, target_corners)

        return gious

    @staticmethod
    @populate_docstring
    def pairwise_bbox_giou(
        pred_bboxes: torch.Tensor,
        target_bboxes: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the Generalized IoUs between all combinations of predicted and target bounding boxes.

        Args:
            pred_bboxes: Predicted bounding boxes of shape `(num_objects, 6)`. {BOUNDING_BOXES_FORMAT_DOC}
            target_bboxes: Target bounding boxes of shape `(<=num_objects, 6)`. {BOUNDING_BOXES_FORMAT_DOC}

        Returns:
            A tensor containing the IoU losses of all combinations.
        """
        # Monai requires input in corner coordinates format
        pred_corners = convert_box_mode(pred_bboxes, "cccwhd", "xyzxyz")
        target_corners = convert_box_mode(target_bboxes, "cccwhd", "xyzxyz")

        # Calculate GIOUs
        gious = box_giou(pred_corners, target_corners)

        return gious

    def _update_class_prevalences(self, target: list[torch.Tensor]):
        """Update the class prevalences based on the classes present in the ground truth

        Args:
            target: A list of ground truth tensors, each of shape (num_objects, 7) where the last dimension contains the
                bounding box class.
        """
        if self.class_balanced_cross_entropy_loss is not None:
            # Note that target classes will not be in order here. But since counts are all that matters, this doesn't
            # affect the results
            target_classes = pad_sequence(target, batch_first=True)[..., -1].long()
            if target_classes.shape[-1] > self.config.num_objects:
                # If there are more than self.config.num_objects targets, only the first self.config.num_objects will be
                # considered
                target_classes = target_classes[..., : self.config.num_objects]
            elif target_classes.shape[-1] < self.config.num_objects:
                # If there are less than self.config.num_objects targets, pad with background class
                target_classes = F.pad(target_classes, (0, self.config.num_objects - target_classes.shape[-1]), value=0)
            self.class_balanced_cross_entropy_loss.update_class_prevalences(target_classes)


# TODO: Implement mAP and mAR calculation
