"""V-SPLADE document encoder for Hugging Face Transformers.

Wraps the ModernVBERT backbone (``transformers>=5.3.0``) together with the
V-SPLADE MLM head so that this repository loads directly with
``AutoModelForMaskedLM.from_pretrained(..., trust_remote_code=True)``.

The module tree deliberately mirrors the checkpoint layout of the V-SPLADE
export (``encoder.encoder.model.*`` for the backbone, ``encoder.mlm_head.*``
for the sparse head), so ``model.safetensors`` loads without any key
remapping. The ``query_encoder.*`` tensors hold the inference-free Li-LSR
query lookup (used by the Sentence Transformers integration) and are not part
of the document encoder, so they are ignored here.

The returned ``logits`` are the SPLADE term logits: MLM logits scaled by
``hidden_size ** -0.25`` with special tokens masked out, exactly as in
https://github.com/naver/v-splade (``UnifiedRetriever._apply_sparse_head``).
A sparse document embedding is obtained via ``log1p(relu(logits))`` followed
by a max-pool over the sequence dimension (see the README).
"""

from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import nn
from transformers.modeling_outputs import MaskedLMOutput

try:
    from transformers.models.modernvbert.configuration_modernvbert import ModernVBertConfig
    from transformers.models.modernvbert.modeling_modernvbert import (
        ModernVBertModel,
        ModernVBertPreTrainedModel,
    )
except ImportError as exc:
    raise ImportError(
        "V-SPLADE requires the ModernVBERT architecture, which is available in "
        "transformers>=5.3.0. Please upgrade with `pip install -U transformers`."
    ) from exc

# Special tokens that are masked out of the sparse representation:
# [UNK], [CLS], [SEP], [PAD], [MASK]
SPECIAL_TOKEN_IDS = [50280, 50281, 50282, 50283, 50284]


class VSPLADEDecoupledEmbedding(nn.Embedding):
    """Word embeddings split into the base vocabulary and the added vision tokens.

    Matches the V-SPLADE export layout: ``weight`` holds the base (MLM) vocabulary
    and ``additional_embedding.weight`` holds the extra tokens appended for the
    vision chat format (``<image>``, ``<end_of_utterance>``, tile markers, ...).
    """

    def __init__(self, num_embeddings: int, num_additional_embeddings: int, embedding_dim: int, **kwargs) -> None:
        super().__init__(num_embeddings, embedding_dim, **kwargs)
        self.num_additional_embeddings = num_additional_embeddings
        self.additional_embedding = nn.Embedding(num_additional_embeddings, embedding_dim)

    def forward(self, input_ids: torch.LongTensor) -> torch.Tensor:
        input_ids = input_ids.clone()
        additional_vocab_indices = torch.where(input_ids >= self.num_embeddings)
        additional_embeddings = self.additional_embedding(input_ids[additional_vocab_indices] - self.num_embeddings)
        input_ids[additional_vocab_indices] = 0
        full_vector = F.embedding(input_ids, self.weight)
        full_vector[additional_vocab_indices] = additional_embeddings
        return full_vector


class VSPLADEModalityProjection(nn.Module):
    """Vision-to-text projection stored as ``modality_projection.proj`` in the export."""

    def __init__(self, input_size: int, output_size: int) -> None:
        super().__init__()
        self.proj = nn.Linear(input_size, output_size, bias=False)

    @property
    def weight(self) -> torch.Tensor:
        # ModernVBertPreTrainedModel._init_weights initializes ``modality_projection.weight``
        return self.proj.weight

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.proj(hidden_states)


class VSPLADEMLMHead(nn.Module):
    """V-SPLADE MLM head: dense -> GELU -> LayerNorm -> decoder (base vocabulary)."""

    def __init__(self, hidden_size: int, vocab_size: int) -> None:
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.norm = nn.LayerNorm(hidden_size)
        self.decoder = nn.Linear(hidden_size, vocab_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.decoder(self.norm(F.gelu(self.dense(hidden_states))))


class _Wrapper(nn.Module):
    """Empty container used to mirror the checkpoint's key prefixes."""


class VSPLADEForMaskedLM(ModernVBertPreTrainedModel):
    config_class = ModernVBertConfig
    _keys_to_ignore_on_load_unexpected = [r"query_encoder\..*"]

    def __init__(self, config: ModernVBertConfig) -> None:
        super().__init__(config)
        main_vocab_size = config.text_config.vocab_size - config.additional_vocab_size

        backbone = ModernVBertModel(config)
        # The export stores the connector projection under an extra ``proj`` level; mirror that.
        backbone.connector.modality_projection = VSPLADEModalityProjection(
            input_size=config.vision_config.hidden_size * (config.pixel_shuffle_factor**2),
            output_size=config.text_config.hidden_size,
        )
        # The export splits the embedding into base + additional tokens; mirror that.
        backbone.text_model.set_input_embeddings(
            VSPLADEDecoupledEmbedding(
                num_embeddings=main_vocab_size,
                num_additional_embeddings=config.additional_vocab_size,
                embedding_dim=config.text_config.hidden_size,
                padding_idx=getattr(config, "pad_token_id", None),
            )
        )

        self.encoder = _Wrapper()
        self.encoder.encoder = _Wrapper()
        self.encoder.encoder.model = backbone
        self.encoder.mlm_head = VSPLADEMLMHead(config.text_config.hidden_size, main_vocab_size)

        self.logit_scale = config.text_config.hidden_size**-0.25

        self.post_init()

    def get_input_embeddings(self):
        return self.encoder.encoder.model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.encoder.encoder.model.set_input_embeddings(value)

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        pixel_values: torch.FloatTensor | None = None,
        pixel_attention_mask: torch.BoolTensor | None = None,
        image_hidden_states: torch.FloatTensor | None = None,
        return_dict: bool | None = None,
    ) -> MaskedLMOutput:
        outputs = self.encoder.encoder.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            pixel_values=pixel_values,
            pixel_attention_mask=pixel_attention_mask,
            image_hidden_states=image_hidden_states,
            return_dict=True,
        )
        logits = self.encoder.mlm_head(outputs.last_hidden_state) * self.logit_scale
        # Zero out special tokens so they never activate in the sparse representation
        # (log1p(relu(0)) == 0), matching the reference special_token_mask.
        # Built on the fly: buffers created in __init__ do not survive meta-device loading.
        special_token_ids = torch.tensor(SPECIAL_TOKEN_IDS, dtype=torch.long, device=logits.device)
        logits = logits.index_fill(-1, special_token_ids, 0.0)

        return MaskedLMOutput(
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


__all__ = ["VSPLADEForMaskedLM"]