"""SigLIP2-Giant + E5-Small-v2 gated-fusion model for animal identification.

This mirrors, one-for-one, the `FaceRecognizer` wrapper the checkpoint was trained
with (see the model card): the tensor names in `model.safetensors` are exactly the
attribute names used here, so the published weights load unchanged.
"""

from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn.functional as F
from torch import nn
from transformers import BertModel, SiglipModel
from transformers.modeling_outputs import ModelOutput
from transformers.modeling_utils import PreTrainedModel

from .configuration_avito_gated import AvitoGatedFusionConfig


@dataclass
class AvitoGatedFusionOutput(ModelOutput):
    """
    Args:
        embeds: L2-normalised fused embedding, `(batch, embedding_dim)`.
        image_embeds: image embedding after `proj_img`, `(batch, embedding_dim)`.
        text_embeds: text embedding after `proj_text`, `(batch, embedding_dim)`.
        gate_weights: gate output, `(batch, 2)`, ordered `[text, image]`.
    """

    embeds: Optional[torch.FloatTensor] = None
    image_embeds: Optional[torch.FloatTensor] = None
    text_embeds: Optional[torch.FloatTensor] = None
    gate_weights: Optional[torch.FloatTensor] = None


class AvitoGatedFusionModel(PreTrainedModel):
    config_class = AvitoGatedFusionConfig
    # Deliberately not a prefix that appears in the checkpoint, so that the
    # `clip.` / `text_encoder.` keys are loaded verbatim.
    base_model_prefix = "avito_gated_fusion"
    main_input_name = "pixel_values"
    _supports_sdpa = True
    supports_gradient_checkpointing = True

    def __init__(self, config: AvitoGatedFusionConfig):
        super().__init__(config)
        self.clip = SiglipModel(config.siglip_config)
        self.text_encoder = BertModel(config.text_config)
        self.proj_img = nn.Linear(config.image_embed_dim, config.embedding_dim)
        self.proj_text = nn.Linear(config.text_embed_dim, config.embedding_dim)
        self.gate = nn.Sequential(
            nn.Linear(config.embedding_dim * 2, config.gate_hidden_dim),
            nn.ReLU(),
            nn.Linear(config.gate_hidden_dim, 2),
            nn.Softmax(dim=-1),
        )
        self.post_init()

    def _init_weights(self, module):
        std = getattr(self.config, "initializer_range", 0.02)
        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):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        elif isinstance(module, nn.Parameter):
            module.data.normal_(mean=0.0, std=std)

    @staticmethod
    def average_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

    def get_image_features(self, pixel_values: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
        """Raw SigLIP2 image embedding, `(batch, image_embed_dim)`."""
        return self.clip.get_image_features(pixel_values=pixel_values, **kwargs)

    def get_text_features(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor,
        token_type_ids: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> torch.FloatTensor:
        """Mean-pooled E5 text embedding, `(batch, text_embed_dim)`."""
        outputs = self.text_encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            **kwargs,
        )
        return self.average_pool(outputs.last_hidden_state, attention_mask)

    def forward(
        self,
        pixel_values: torch.FloatTensor,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor,
        token_type_ids: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        img_emb = self.get_image_features(pixel_values)
        text_emb = self.get_text_features(input_ids, attention_mask, token_type_ids)

        img_proj = self.proj_img(img_emb)
        text_proj = self.proj_text(text_emb)

        fused = torch.cat([text_proj, img_proj], dim=-1)
        w = self.gate(fused)
        fused_emb = w[:, 0:1] * text_proj + w[:, 1:2] * img_proj
        embeds = F.normalize(fused_emb, dim=1)

        if not return_dict:
            return (embeds, img_proj, text_proj, w)
        return AvitoGatedFusionOutput(
            embeds=embeds, image_embeds=img_proj, text_embeds=text_proj, gate_weights=w
        )


__all__ = ["AvitoGatedFusionModel", "AvitoGatedFusionOutput"]