"""
Mage-ViT — the standalone vision encoder of the Mage-VL family.

The implementation is adapted from the vision tower of Mage-VL (``modeling_mage_vl.py``):

    * fused ``qkv`` / ``proj`` self-attention,
    * 3D (T,H,W) rotary position embeddings with a 4:6:6 split (``VisionRotaryEmbedding``),
    * SigLIP-style MLP blocks and a multi-head attention pooling head.

It is written to be transformers-version agnostic (works with both ``transformers>=5``
and ``transformers==4.57.x``): attention is dispatched internally across
``sdpa`` / ``flash_attention_2`` / ``eager`` without relying on any transformers-5-only
symbols, and it does not require ``flash-attn`` to be installed (``sdpa`` is the default).
"""

from typing import Callable, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
from transformers.modeling_utils import PreTrainedModel
from transformers.models.siglip.modeling_siglip import SiglipMLP
from transformers.utils import logging

from .configuration_mage_vit import MageViTConfig


try:
    from flash_attn import flash_attn_func

    _flash_attn_available = True
except ImportError:
    _flash_attn_available = False

logger = logging.get_logger(__name__)


# ---------------------------------------------------------------------------
# Helper Functions & Layers
# ---------------------------------------------------------------------------


def get_norm_layer(config):
    if config.layer_norm_type == "rms_norm":
        return nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
    else:
        return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)


def rotate_half(x):
    """
    Interleaved rotation matching the training-time implementation.
    (x1, x2, x3, x4) -> (-x2, x1, -x4, x3)
    """
    x_even = x[..., ::2]
    x_odd = x[..., 1::2]
    return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)


def apply_rotary_pos_emb(q, k, freqs):
    # q, k: (B, H, L, D); freqs: (B, L, D) or (1, L, D)
    orig_q_dtype = q.dtype
    orig_k_dtype = k.dtype
    q, k = q.float(), k.float()
    cos = freqs.cos().unsqueeze(1).float()  # (B, 1, L, D)
    sin = freqs.sin().unsqueeze(1).float()

    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed.to(orig_q_dtype), k_embed.to(orig_k_dtype)


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: float,
    dropout: float = 0.0,
):
    """Eager attention; query/key/value are expected as ``(B, H, L, D)``."""
    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask
    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = F.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value)
    attn_output = attn_output.transpose(1, 2).contiguous()  # (B, L, H, D)
    return attn_output, attn_weights


class VisionRotaryEmbedding(nn.Module):
    """
    3D (T,H,W) Rotary frequency constructor with a 4:6:6 split for T:H:W.
    """

    def __init__(self, config: MageViTConfig):
        super().__init__()
        head_dim = config.hidden_size // config.num_attention_heads
        base = config.rope_theta

        assert head_dim % 2 == 0, "head_dim must be even for rotary."
        assert head_dim % 16 == 0, "head_dim must be divisible by 16."
        half = head_dim // 2
        assert half % 16 == 0, "head_dim//2 must also be divisible by 16 to split into 4:6:6."

        self.head_dim = head_dim
        self.half = half
        self.base = base

        unit = half // 16
        self.t_size = 4 * unit
        self.h_size = 6 * unit
        self.w_size = 6 * unit

        self.register_buffer(
            "inv_freq_t",
            1.0 / (base ** (torch.arange(self.t_size, dtype=torch.float32) / self.t_size)),
            persistent=False,
        )
        self.register_buffer(
            "inv_freq_h",
            1.0 / (base ** (torch.arange(self.h_size, dtype=torch.float32) / self.h_size)),
            persistent=False,
        )
        self.register_buffer(
            "inv_freq_w",
            1.0 / (base ** (torch.arange(self.w_size, dtype=torch.float32) / self.w_size)),
            persistent=False,
        )

    def forward_with_thw(self, t: int, h: int, w: int, device=None) -> torch.Tensor:
        """Build RoPE frequencies for a full (t, h, w) patch grid -> [t*h*w, half]."""
        if device is None:
            device = self.inv_freq_t.device

        inv_t = self.inv_freq_t.to(device=device)
        inv_h = self.inv_freq_h.to(device=device)
        inv_w = self.inv_freq_w.to(device=device)

        ft = torch.outer(torch.arange(t, device=device, dtype=torch.float32), inv_t)
        fh = torch.outer(torch.arange(h, device=device, dtype=torch.float32), inv_h)
        fw = torch.outer(torch.arange(w, device=device, dtype=torch.float32), inv_w)

        t_ids = torch.arange(t, device=device).repeat_interleave(h * w)
        h_ids = torch.arange(h, device=device).repeat_interleave(w).repeat(t)
        w_ids = torch.arange(w, device=device).repeat(h).repeat(t)

        return torch.cat([ft[t_ids], fh[h_ids], fw[w_ids]], dim=-1)

    def forward_from_positions(self, patch_positions: torch.Tensor) -> torch.Tensor:
        """
        Build RoPE frequencies from explicit patch positions.

        Args:
            patch_positions: [seq_len, 3] or [batch_size, seq_len, 3] with [t, h, w] per patch.

        Returns:
            freqs with a leading batch dim: [batch_size, seq_len, half].
        """
        if patch_positions.dim() == 2:
            patch_positions = patch_positions.unsqueeze(0)

        device = patch_positions.device
        inv_t = self.inv_freq_t.to(device=device)
        inv_h = self.inv_freq_h.to(device=device)
        inv_w = self.inv_freq_w.to(device=device)

        t_pos = patch_positions[..., 0].float()  # [B, L]
        h_pos = patch_positions[..., 1].float()
        w_pos = patch_positions[..., 2].float()

        ft = torch.einsum("bs,d->bsd", t_pos, inv_t)
        fh = torch.einsum("bs,d->bsd", h_pos, inv_h)
        fw = torch.einsum("bs,d->bsd", w_pos, inv_w)

        return torch.cat([ft, fh, fw], dim=-1)


class Siglip2MultiheadAttentionPoolingHead(nn.Module):
    """Multi-Head Attention Pooling with a learned probe (PMA-style)."""

    def __init__(self, config: MageViTConfig):
        super().__init__()
        self.embed_dim = config.hidden_size
        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
        self.attention = nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
        self.norm = nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.mlp = SiglipMLP(config)

    def forward(self, hidden_states):
        batch_size = hidden_states.shape[0]
        probe = self.probe.repeat(batch_size, 1, 1)

        attn_output, _ = self.attention(probe, hidden_states, hidden_states)

        residual = attn_output
        attn_output = self.norm(attn_output)
        attn_output = residual + self.mlp(attn_output)

        return attn_output[:, 0]


# ---------------------------------------------------------------------------
# Modeling Components
# ---------------------------------------------------------------------------


class MageViTEmbeddings(nn.Module):
    """Conv2d patch embedding for pixel-value inputs (4D image or 5D video)."""

    def __init__(self, config: MageViTConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.image_size = config.image_size
        self.patch_size = config.patch_size

        self.patch_embedding = nn.Conv2d(
            in_channels=config.num_channels,
            out_channels=self.embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
            bias=False,
        )

    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
        # Handle 4D (B, C, H, W) or 5D (B, C, T, H, W) inputs
        if pixel_values.dim() == 4:
            pixel_values = pixel_values.unsqueeze(2)  # (B, C, 1, H, W)

        batch_size, channels, t_frames, height, width = pixel_values.shape
        target_dtype = self.patch_embedding.weight.dtype

        # Merge time into batch for Conv2d
        x_2d = pixel_values.permute(0, 2, 1, 3, 4).reshape(batch_size * t_frames, channels, height, width)

        embeddings = self.patch_embedding(x_2d.to(dtype=target_dtype))  # (B*T, C, Hp, Wp)
        embeddings = embeddings.flatten(2).transpose(1, 2)  # (B*T, L_frame, C)

        total_patches = t_frames * (height // self.patch_size) * (width // self.patch_size)
        embeddings = embeddings.reshape(batch_size, total_patches, self.embed_dim)

        return embeddings


class MageViTAttention(nn.Module):
    """
    Multi-headed attention with fused ``qkv`` / ``proj`` projections and RoPE support.

    Attention is dispatched internally across ``sdpa`` (default) / ``flash_attention_2`` /
    ``eager`` based on ``config._attn_implementation``.
    """

    def __init__(self, config: MageViTConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and "
                f"`num_heads`: {self.num_heads})."
            )

        self.scale = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.qkv = nn.Linear(self.embed_dim, self.embed_dim * 3)
        self.proj = nn.Linear(self.embed_dim, self.embed_dim)

    def _attn_impl(self) -> str:
        impl = getattr(self.config, "_attn_implementation", None) or "sdpa"
        if impl == "flash_attention_2" and not _flash_attn_available:
            logger.warning_once("flash-attn is not installed; falling back to `sdpa` attention.")
            impl = "sdpa"
        return impl

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        batch_size, q_len, _ = hidden_states.size()

        # (B, L, 3*H*D) -> (B, L, 3, H, D) -> 3 x (B, H, L, D)
        q, k, v = (
            self.qkv(hidden_states)
            .reshape(batch_size, q_len, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 1, 3, 4)
            .unbind(0)
        )
        query_states = q.transpose(1, 2)
        key_states = k.transpose(1, 2)
        value_states = v.transpose(1, 2)

        if rotary_pos_emb is not None:
            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, rotary_pos_emb)

        dropout = self.attention_dropout if self.training else 0.0
        impl = self._attn_impl()
        attn_weights = None

        if output_attentions or impl == "eager":
            attn_output, attn_weights = eager_attention_forward(
                self, query_states, key_states, value_states, attention_mask, self.scale, dropout
            )
        elif impl == "flash_attention_2":
            # flash-attn expects (B, L, H, D)
            attn_output = flash_attn_func(
                query_states.transpose(1, 2),
                key_states.transpose(1, 2),
                value_states.transpose(1, 2),
                dropout_p=dropout,
                softmax_scale=self.scale,
                causal=False,
            )  # (B, L, H, D)
        else:  # sdpa
            attn_output = F.scaled_dot_product_attention(
                query_states,
                key_states,
                value_states,
                attn_mask=attention_mask,
                dropout_p=dropout,
                scale=self.scale,
            )  # (B, H, L, D)
            attn_output = attn_output.transpose(1, 2)  # (B, L, H, D)

        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
        attn_output = self.proj(attn_output)

        return attn_output, attn_weights if output_attentions else None


class MageViTEncoderLayer(nn.Module):
    """Vision encoder layer with pre-norm attention and MLP."""

    def __init__(self, config: MageViTConfig):
        super().__init__()
        self.embed_dim = config.hidden_size
        self.self_attn = MageViTAttention(config)
        self.layer_norm1 = get_norm_layer(config)
        self.mlp = SiglipMLP(config)
        self.layer_norm2 = get_norm_layer(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        residual = hidden_states
        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            rotary_pos_emb=rotary_pos_emb,
            output_attentions=output_attentions,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states, attn_weights) if output_attentions else (hidden_states,)
        return outputs


class MageViTEncoder(nn.Module):
    def __init__(self, config: MageViTConfig):
        super().__init__()
        self.config = config
        self.layers = nn.ModuleList([MageViTEncoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ) -> Union[tuple, BaseModelOutput]:
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None

        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    layer.__call__,
                    hidden_states,
                    attention_mask,
                    rotary_pos_emb,
                    output_attentions,
                )
            else:
                layer_outputs = layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    rotary_pos_emb=rotary_pos_emb,
                    output_attentions=output_attentions,
                )

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)

        return BaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
        )


# ---------------------------------------------------------------------------
# Main Models
# ---------------------------------------------------------------------------


class MageViTPreTrainedModel(PreTrainedModel):
    config_class = MageViTConfig
    base_model_prefix = "mage_vit"
    supports_gradient_checkpointing = True
    _no_split_modules = ["MageViTEncoderLayer"]
    _supports_flash_attn_2 = True
    _supports_flash_attn = True
    _supports_sdpa = True

    def _init_weights(self, module):
        """Initialize the weights."""
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, (nn.LayerNorm, nn.RMSNorm)):
            module.weight.data.fill_(1.0)
            if hasattr(module, "bias") and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, VisionRotaryEmbedding):
            # inv_freq buffers are registered with persistent=False, so they are not in the
            # checkpoint. When `from_pretrained` materializes the model from meta tensors these
            # buffers would otherwise stay uninitialized; re-fill them so RoPE is correct post-load.
            base = module.base
            with torch.no_grad():
                inv_t = 1.0 / (base ** (torch.arange(module.t_size, dtype=torch.float32) / module.t_size))
                inv_h = 1.0 / (base ** (torch.arange(module.h_size, dtype=torch.float32) / module.h_size))
                inv_w = 1.0 / (base ** (torch.arange(module.w_size, dtype=torch.float32) / module.w_size))
                module.inv_freq_t.copy_(inv_t.to(module.inv_freq_t.device))
                module.inv_freq_h.copy_(inv_h.to(module.inv_freq_h.device))
                module.inv_freq_w.copy_(inv_w.to(module.inv_freq_w.device))


class MageViTModel(MageViTPreTrainedModel):
    """Mage-ViT vision transformer encoder."""

    def __init__(self, config: MageViTConfig):
        super().__init__(config)
        self.config = config

        self.embeddings = MageViTEmbeddings(config)
        self.layernorm_pre = get_norm_layer(config)
        self.encoder = MageViTEncoder(config)
        self.video_rope = VisionRotaryEmbedding(config)

        if config.use_head:
            self.layernorm_post = get_norm_layer(config)
            self.head = Siglip2MultiheadAttentionPoolingHead(config)
        else:
            self.layernorm_post = None
            self.head = None

        self.post_init()

    def forward(
        self,
        pixel_values: torch.Tensor,
        visible_indices: Optional[torch.Tensor] = None,
        patch_positions: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[tuple, BaseModelOutputWithPooling]:
        r"""
        Examples:

        ```python
        >>> from transformers import AutoModel, AutoImageProcessor
        >>> from PIL import Image

        >>> model = AutoModel.from_pretrained("microsoft/Mage-ViT", trust_remote_code=True)
        >>> preprocessor = AutoImageProcessor.from_pretrained("microsoft/Mage-ViT", trust_remote_code=True)
        >>> image = Image.open("path/to/your/image.jpg")
        >>> pixel_values = preprocessor(images=image, return_tensors="pt")["pixel_values"]
        >>> outputs = model(pixel_values)
        >>> last_hidden_states = outputs.last_hidden_state
        >>> pooled_output = outputs.pooler_output
        ```
        """
        output_attentions = (
            output_attentions if output_attentions is not None else getattr(self.config, "output_attentions", False)
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else getattr(self.config, "output_hidden_states", False)
        )
        return_dict = True if return_dict is None else return_dict

        # Determine grid dimensions for RoPE (pixel_values may be 4D or 5D)
        if pixel_values.dim() == 5:
            t_frames = (
                self.config.rope_temporal_size if self.config.rope_temporal_size is not None else pixel_values.shape[2]
            )
            height, width = pixel_values.shape[3], pixel_values.shape[4]
        else:
            t_frames = 1
            height, width = pixel_values.shape[2], pixel_values.shape[3]

        # 1. Embeddings
        hidden_states = self.embeddings(pixel_values)
        batch_size, total_patches, _ = hidden_states.shape

        # 2. Visible-index handling (defaults to all patches)
        if visible_indices is None:
            visible_indices = torch.arange(total_patches, device=pixel_values.device).unsqueeze(0).expand(
                batch_size, -1
            )

        # 3. RoPE construction
        if patch_positions is not None:
            freqs_visible = self.video_rope.forward_from_positions(patch_positions)  # (B, L, half)
        else:
            freqs_full = self.video_rope.forward_with_thw(
                t=t_frames,
                h=height // self.config.patch_size,
                w=width // self.config.patch_size,
                device=pixel_values.device,
            )
            freqs_visible = freqs_full[visible_indices]  # (B, L, half)

        # Concatenate half + half -> head_dim
        freqs_visible = torch.cat([freqs_visible, freqs_visible], dim=-1)

        # 4. Pre-norm & encoder
        hidden_states = self.layernorm_pre(hidden_states)

        # Sparse mode: gather only visible patches to match freqs_visible
        if visible_indices.shape[1] != total_patches:
            hidden_states = hidden_states.gather(
                1, visible_indices.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1])
            )

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=None,
            rotary_pos_emb=freqs_visible,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        sequence_output = encoder_outputs.last_hidden_state

        if self.layernorm_post is not None:
            sequence_output = self.layernorm_post(sequence_output)

        pooled_output = self.head(sequence_output) if self.head is not None else None

        if not return_dict:
            outputs = (sequence_output, pooled_output)
            if output_hidden_states:
                outputs = outputs + (encoder_outputs.hidden_states,)
            if output_attentions:
                outputs = outputs + (encoder_outputs.attentions,)
            return outputs

        return BaseModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )


__all__ = ["MageViTModel", "MageViTPreTrainedModel"]