# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
# Based on https://github.com/IDEA-Research/GroundingDINO
from __future__ import annotations

import torch
from torch import nn

from ultralytics.nn.modules.utils import _get_clones

from .model_misc import get_valid_ratio


class TransformerEncoderLayer(nn.Module):
    """Transformer encoder layer that performs self-attention followed by cross-attention.

    This layer was previously called TransformerDecoderLayer but was renamed to better reflect its role in the
    architecture. It processes input sequences through self-attention and then cross-attention with another input
    (typically image features).

    The layer supports both pre-norm and post-norm configurations, as well as positional encoding at different stages of
    the attention mechanism.
    """

    def __init__(
        self,
        d_model: int,
        dim_feedforward: int,
        dropout: float,
        pos_enc_at_attn: bool,
        pos_enc_at_cross_attn_keys: bool,
        pos_enc_at_cross_attn_queries: bool,
        pre_norm: bool,
        self_attention: nn.Module | None = None,
        cross_attention: nn.Module | None = None,
    ):
        """Initialize a transformer encoder layer.

        Args:
            d_model (int): Model dimension/hidden size.
            dim_feedforward (int): Dimension of the feedforward network.
            dropout (float): Dropout probability.
            pos_enc_at_attn (bool): Whether to add positional encodings at self-attention.
            pos_enc_at_cross_attn_keys (bool): Whether to add positional encodings to keys in cross-attention.
            pos_enc_at_cross_attn_queries (bool): Whether to add positional encodings to queries in cross-attention.
            pre_norm (bool): Whether to use pre-norm (True) or post-norm (False) architecture.
            self_attention (nn.Module | None): Self-attention module. If None, a default 8-head, 256-dim
                nn.MultiheadAttention is used.
            cross_attention (nn.Module | None): Cross-attention module for attending to prompt/memory features. If None,
                a default 8-head, 256-dim nn.MultiheadAttention is used.
        """
        super().__init__()
        self.d_model = d_model
        self.dim_feedforward = dim_feedforward
        self.dropout_value = dropout
        self.self_attn = self_attention or nn.MultiheadAttention(num_heads=8, dropout=0.1, embed_dim=256)
        self.cross_attn_image = cross_attention or nn.MultiheadAttention(num_heads=8, dropout=0.1, embed_dim=256)

        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.activation = nn.ReLU()
        self.pre_norm = pre_norm

        self.pos_enc_at_attn = pos_enc_at_attn
        self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries
        self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys

        self.layer_idx = None

    def forward_post(
        self,
        tgt: torch.Tensor,
        memory: torch.Tensor,
        tgt_mask: torch.Tensor | None = None,
        memory_mask: torch.Tensor | None = None,
        tgt_key_padding_mask: torch.Tensor | None = None,
        memory_key_padding_mask: torch.Tensor | None = None,
        pos: torch.Tensor | None = None,
        query_pos: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        """Forward pass for post-norm architecture.

        In post-norm architecture, normalization is applied after attention and feedforward operations.

        Args:
            tgt (torch.Tensor): Input tensor to be processed.
            memory (torch.Tensor): Memory tensor for cross-attention.
            tgt_mask (torch.Tensor | None): Mask for self-attention.
            memory_mask (torch.Tensor | None): Mask for cross-attention.
            tgt_key_padding_mask (torch.Tensor | None): Key padding mask for self-attention.
            memory_key_padding_mask (torch.Tensor | None): Key padding mask for cross-attention.
            pos (torch.Tensor | None): Positional encoding for memory.
            query_pos (torch.Tensor | None): Positional encoding for query.
            **kwargs (Any): Additional keyword arguments (e.g. dac), ignored in post-norm mode.

        Returns:
            (torch.Tensor): Processed tensor with the same shape as tgt.
        """
        q = k = tgt + query_pos if self.pos_enc_at_attn else tgt

        # Self attention
        tgt2 = self.self_attn(
            q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask, need_weights=False
        )[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)

        # Cross attention to image
        tgt2 = self.cross_attn_image(
            query=tgt + query_pos if self.pos_enc_at_cross_attn_queries else tgt,
            key=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
            value=memory,
            attn_mask=memory_mask,
            key_padding_mask=memory_key_padding_mask,
            need_weights=False,
        )[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)

        # FFN
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt

    def forward_pre(
        self,
        tgt: torch.Tensor,
        memory: torch.Tensor,
        dac: bool = False,
        tgt_mask: torch.Tensor | None = None,
        memory_mask: torch.Tensor | None = None,
        tgt_key_padding_mask: torch.Tensor | None = None,
        memory_key_padding_mask: torch.Tensor | None = None,
        pos: torch.Tensor | None = None,
        query_pos: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Forward pass for pre-norm architecture.

        In pre-norm architecture, normalization is applied before attention and feedforward operations.

        Args:
            tgt (torch.Tensor): Input tensor to be processed.
            memory (torch.Tensor): Memory tensor for cross-attention.
            dac (bool): Whether to use Divide-and-Conquer attention (only apply self-attention to the first half).
            tgt_mask (torch.Tensor | None): Mask for self-attention.
            memory_mask (torch.Tensor | None): Mask for cross-attention.
            tgt_key_padding_mask (torch.Tensor | None): Key padding mask for self-attention.
            memory_key_padding_mask (torch.Tensor | None): Key padding mask for cross-attention.
            pos (torch.Tensor | None): Positional encoding for memory.
            query_pos (torch.Tensor | None): Positional encoding for query.

        Returns:
            (torch.Tensor): Processed tensor with the same shape as tgt.
        """
        if dac:
            # we only apply self attention to the first half of the queries
            assert tgt.shape[0] % 2 == 0
            other_tgt = tgt[tgt.shape[0] // 2 :]
            tgt = tgt[: tgt.shape[0] // 2]
        tgt2 = self.norm1(tgt).contiguous()
        q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        if dac:
            # Recombine
            tgt = torch.cat((tgt, other_tgt), dim=0)
        tgt2 = self.norm2(tgt)
        memory = memory.to(tgt2.dtype).contiguous()
        tgt2 = self.cross_attn_image(
            query=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2,
            key=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
            value=memory,
            attn_mask=memory_mask,
            key_padding_mask=memory_key_padding_mask,
        )[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt

    def forward(
        self,
        tgt: torch.Tensor,
        memory: torch.Tensor,
        dac: bool = False,
        tgt_mask: torch.Tensor | None = None,
        memory_mask: torch.Tensor | None = None,
        tgt_key_padding_mask: torch.Tensor | None = None,
        memory_key_padding_mask: torch.Tensor | None = None,
        pos: torch.Tensor | None = None,
        query_pos: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Forward pass for the transformer encoder layer.

        Args:
            tgt (torch.Tensor): Input tensor to be processed.
            memory (torch.Tensor): Memory tensor (e.g., prompt features) for cross-attention.
            dac (bool): Whether to use Divide-and-Conquer attention (only apply self-attention to first half). Only used
                in pre-norm mode.
            tgt_mask (torch.Tensor | None): Mask for self-attention.
            memory_mask (torch.Tensor | None): Mask for cross-attention.
            tgt_key_padding_mask (torch.Tensor | None): Key padding mask for self-attention.
            memory_key_padding_mask (torch.Tensor | None): Key padding mask for cross-attention.
            pos (torch.Tensor | None): Positional encoding for memory.
            query_pos (torch.Tensor | None): Positional encoding for query.

        Returns:
            (torch.Tensor): Processed tensor after self-attention, cross-attention, and feedforward network.
        """
        fwd_fn = self.forward_pre if self.pre_norm else self.forward_post
        return fwd_fn(
            tgt,
            memory,
            dac=dac,
            tgt_mask=tgt_mask,
            memory_mask=memory_mask,
            tgt_key_padding_mask=tgt_key_padding_mask,
            memory_key_padding_mask=memory_key_padding_mask,
            pos=pos,
            query_pos=query_pos,
            # attn_bias=attn_bias,
            # **kwds,
        )


class TransformerEncoder(nn.Module):
    """Transformer encoder that processes multi-level features.

    This encoder takes multi-level features (e.g., from a backbone network) and processes them through a stack of
    transformer encoder layers. It supports features from multiple levels (e.g., different resolutions) and can apply
    activation checkpointing for memory efficiency during training.

    Args:
        layer (nn.Module): The encoder layer to be stacked multiple times.
        num_layers (int): Number of encoder layers to stack.
        d_model (int): Model dimension/hidden size.
        num_feature_levels (int): Number of feature levels to process.
        frozen (bool): Whether to freeze the parameters of this module.
        use_act_checkpoint (bool): Whether to use activation checkpointing during training.
    """

    def __init__(
        self,
        layer: nn.Module,
        num_layers: int,
        d_model: int,
        num_feature_levels: int,
        frozen: bool = False,
        use_act_checkpoint: bool = False,
    ):
        """Initialize the transformer encoder."""
        super().__init__()
        self.layers = _get_clones(layer, num_layers)
        self.num_layers = num_layers

        self.num_feature_levels = num_feature_levels
        self.level_embed = None
        if num_feature_levels > 1:
            self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model))

        if frozen:
            for p in self.parameters():
                p.requires_grad_(False)

        self.use_act_checkpoint = use_act_checkpoint

        # assign layer index to each layer so that some layers can decide what to do
        # based on which layer index they are (e.g. cross attention to memory bank only
        # in selected layers)
        for layer_idx, encoder_layer in enumerate(self.layers):
            encoder_layer.layer_idx = layer_idx

    def _prepare_multilevel_features(self, srcs, masks, pos_embeds):
        """Prepare multi-level features for transformer encoder."""
        assert len(srcs) == self.num_feature_levels, "mismatch between expected and received # of feature levels"

        src_flatten = []
        mask_flatten = []
        lvl_pos_embed_flatten = []
        has_mask = masks is not None and masks[0] is not None
        for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)):
            src = src.flatten(2).transpose(1, 2)  # bs, hw, c
            if has_mask:
                mask = mask.flatten(1)
            pos_embed = pos_embed.flatten(2).transpose(1, 2)  # bs, hw, c
            if self.level_embed is not None:
                lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)
            else:
                lvl_pos_embed = pos_embed
            lvl_pos_embed_flatten.append(lvl_pos_embed)
            src_flatten.append(src)
            if has_mask:
                mask_flatten.append(mask)
        src_flatten = torch.cat(src_flatten, 1)  # bs, \sum{hxw}, c
        mask_flatten = torch.cat(mask_flatten, 1) if has_mask else None  # bs, \sum{hxw}
        lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)  # bs, \sum{hxw}, c
        if has_mask:
            valid_ratios = torch.stack([get_valid_ratio(m) for m in masks], 1)
        else:
            valid_ratios = torch.ones(
                (src_flatten.shape[0], self.num_feature_levels, 2),
                device=src_flatten.device,
                dtype=src_flatten.dtype,
            )

        return (
            src_flatten,
            mask_flatten,
            lvl_pos_embed_flatten,
            valid_ratios,
        )

    def forward(
        self,
        src: list[torch.Tensor],
        src_key_padding_masks: list[torch.Tensor] | None = None,
        pos: list[torch.Tensor] | None = None,
        prompt: torch.Tensor | None = None,
        prompt_key_padding_mask: torch.Tensor | None = None,
        encoder_extra_kwargs: dict | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Process multi-level features through the transformer encoder.

        Args:
            src (list[torch.Tensor]): List of multi-level features, each with shape (batch_size, channels, height,
                width).
            src_key_padding_masks (list[torch.Tensor] | None): List of padding masks for each feature level, each with
                shape (batch_size, height, width).
            pos (list[torch.Tensor] | None): List of positional embeddings for each feature level, each with shape
                (batch_size, channels, height, width).
            prompt (torch.Tensor | None): Text/prompt features to attend to, with shape (batch_size, seq_len, d_model)
                since the encoder layers operate batch-first.
            prompt_key_padding_mask (torch.Tensor | None): Padding mask for prompt, with shape (batch_size, seq_len).
            encoder_extra_kwargs (dict | None): Additional arguments to pass to each encoder layer.

        Returns:
            output (torch.Tensor): Processed features with shape (sum(H*W), batch_size, d_model).
            key_padding_masks_flatten (torch.Tensor | None): Flattened padding masks with shape (sum(H*W), batch_size).
            lvl_pos_embed_flatten (torch.Tensor): Flattened positional embeddings with shape (sum(H*W), batch_size,
                d_model).
            valid_ratios (torch.Tensor): Valid ratios for each feature level with shape (batch_size, num_feature_levels,
                2).
        """
        assert len(src) == self.num_feature_levels, "must be equal to num_feature_levels"
        if src_key_padding_masks is not None:
            assert len(src_key_padding_masks) == self.num_feature_levels
        if pos is not None:
            assert len(pos) == self.num_feature_levels
        # Flatten multilevel feats and add level pos embeds
        (
            src_flatten,
            key_padding_masks_flatten,
            lvl_pos_embed_flatten,
            valid_ratios,
        ) = self._prepare_multilevel_features(src, src_key_padding_masks, pos)

        output = src_flatten
        for layer in self.layers:
            layer_kwargs = {}

            assert isinstance(layer, TransformerEncoderLayer)
            layer_kwargs["memory"] = prompt
            layer_kwargs["memory_key_padding_mask"] = prompt_key_padding_mask
            layer_kwargs["query_pos"] = lvl_pos_embed_flatten
            layer_kwargs["tgt"] = output
            layer_kwargs["tgt_key_padding_mask"] = key_padding_masks_flatten

            if self.training:
                assert self.use_act_checkpoint, "activation ckpt not enabled in encoder"
            if encoder_extra_kwargs is not None:
                layer_kwargs.update(encoder_extra_kwargs)
            output = layer(**layer_kwargs)
        # return as seq first
        return (
            output.transpose(0, 1),
            (key_padding_masks_flatten.transpose(0, 1) if key_padding_masks_flatten is not None else None),
            lvl_pos_embed_flatten.transpose(0, 1),
            valid_ratios,
        )


class TransformerEncoderFusion(TransformerEncoder):
    """Transformer encoder that fuses text and image features.

    This encoder extends TransformerEncoder to handle both text and image features, with the ability to add pooled text
    features to image features for better cross-modal fusion. It supports torch.compile for performance optimization.

    Args:
        layer (nn.Module): The encoder layer to be stacked multiple times.
        num_layers (int): Number of encoder layers to stack.
        d_model (int): Model dimension/hidden size.
        num_feature_levels (int): Number of feature levels to process.
        add_pooled_text_to_img_feat (bool): Whether to add pooled text features to image features.
        pool_text_with_mask (bool): Whether to use the mask when pooling text features.
        compile_mode (str | None): Mode for torch.compile, or None to disable compilation.
        **kwargs (Any): Additional arguments to pass to the parent class.
    """

    def __init__(
        self,
        layer: nn.Module,
        num_layers: int,
        d_model: int,
        num_feature_levels: int,
        add_pooled_text_to_img_feat: bool = True,
        pool_text_with_mask: bool = False,
        compile_mode: str | None = None,
        **kwargs,
    ):
        """Initialize the transformer encoder with text-image fusion."""
        super().__init__(
            layer,
            num_layers,
            d_model,
            num_feature_levels,
            **kwargs,
        )
        self.add_pooled_text_to_img_feat = add_pooled_text_to_img_feat
        if self.add_pooled_text_to_img_feat:
            self.text_pooling_proj = nn.Linear(d_model, d_model)
        self.pool_text_with_mask = pool_text_with_mask
        if compile_mode is not None:
            self.forward = torch.compile(self.forward, mode=compile_mode, fullgraph=True)

    def forward(
        self,
        src: list[torch.Tensor],
        prompt: torch.Tensor,
        src_key_padding_mask: list[torch.Tensor] | None = None,
        src_pos: list[torch.Tensor] | None = None,
        prompt_key_padding_mask: torch.Tensor | None = None,
        feat_sizes: list[tuple[int, int]] | None = None,
        encoder_extra_kwargs: dict | None = None,
    ):
        """Fuse image features with prompt features through the transformer encoder.

        Args:
            src (list[torch.Tensor]): Multi-level image features, each with shape (H*W, bs, C) if feat_sizes is given,
                otherwise (bs, C, H, W).
            prompt (torch.Tensor): Prompt features with shape (seq_len, bs, d_model).
            src_key_padding_mask (list[torch.Tensor] | None): Padding masks for each feature level.
            src_pos (list[torch.Tensor] | None): Positional embeddings for each feature level, matching src.
            prompt_key_padding_mask (torch.Tensor | None): Padding mask for prompt with shape (bs, seq_len), where True
                marks padding.
            feat_sizes (list[tuple[int, int]] | None): Spatial sizes (H, W) used to reshape sequence-first src.
            encoder_extra_kwargs (dict | None): Additional arguments to pass to each encoder layer.

        Returns:
            (dict): Dictionary with keys "memory", "padding_mask", "pos_embed", "memory_text", and "valid_ratios".
        """
        # Restore spatial shapes of vision
        bs = src[0].shape[1]  # seq first
        if feat_sizes is not None:
            assert len(feat_sizes) == len(src)
            if src_key_padding_mask is None:
                src_key_padding_mask = [None] * len(src)
            for i, (h, w) in enumerate(feat_sizes):
                src[i] = src[i].reshape(h, w, bs, -1).permute(2, 3, 0, 1)
                src_pos[i] = src_pos[i].reshape(h, w, bs, -1).permute(2, 3, 0, 1)
                src_key_padding_mask[i] = (
                    src_key_padding_mask[i].reshape(h, w, bs).permute(2, 0, 1)
                    if src_key_padding_mask[i] is not None
                    else None
                )
        else:
            assert all(x.dim() == 4 for x in src), "expected list of (bs, c, h, w) tensors"

        if self.add_pooled_text_to_img_feat:
            # Fusion: Add mean pooled text to image features
            pooled_text = pool_text_feat(prompt, prompt_key_padding_mask, self.pool_text_with_mask)
            pooled_text = self.text_pooling_proj(pooled_text)[..., None, None]  # prompt is seq first
            src = [x.add_(pooled_text) for x in src]

        (
            out,
            key_padding_masks_flatten,
            lvl_pos_embed_flatten,
            valid_ratios,
        ) = super().forward(
            src,
            src_key_padding_masks=src_key_padding_mask,
            pos=src_pos,
            prompt=prompt.transpose(0, 1),
            prompt_key_padding_mask=prompt_key_padding_mask,
            encoder_extra_kwargs=encoder_extra_kwargs,
        )

        return {
            "memory": out,
            "padding_mask": key_padding_masks_flatten,
            "pos_embed": lvl_pos_embed_flatten,
            "memory_text": prompt,
            "valid_ratios": valid_ratios,
        }


def pool_text_feat(prompt, prompt_mask, pool_with_mask):
    """Mean-pool prompt embeddings over the sequence dimension, optionally using only valid tokens.

    Args:
        prompt (torch.Tensor): Prompt embeddings with shape (seq_len, bs, dim).
        prompt_mask (torch.Tensor | None): Padding mask with shape (bs, seq_len), where True marks padding.
        pool_with_mask (bool): Whether to average only over valid (non-padding) tokens.

    Returns:
        (torch.Tensor): Pooled prompt embeddings with shape (bs, dim).
    """
    # prompt has shape (seq, bs, dim)
    if not pool_with_mask:
        return prompt.mean(dim=0)

    # prompt_mask has shape (bs, seq), where False is valid and True is padding
    assert prompt_mask.dim() == 2
    # is_valid has shape (seq, bs, 1), where 1 is valid and 0 is padding
    is_valid = (~prompt_mask).float().permute(1, 0)[..., None]
    # num_valid has shape (bs, 1)
    num_valid = torch.clamp(torch.sum(is_valid, dim=0), min=1.0)

    # mean pool over all the valid tokens
    pooled_text = (prompt * is_valid).sum(dim=0) / num_valid
    return pooled_text