"""Sentence Transformers module for the V-SPLADE inference-free query encoder.

Referenced from ``router_config.json``: the "query" route uses
:class:`VSPLADEStaticEmbedding`, whose weights are the precomputed Li-LSR
lookup table ``softplus(projection(embedding))`` extracted from the
``query_encoder.*`` tensors in ``model.safetensors`` (with the special tokens
[UNK]/[CLS]/[SEP]/[PAD]/[MASK] zeroed out).
"""

from __future__ import annotations

import torch

try:
    # sentence-transformers >= 5.6
    from sentence_transformers.sparse_encoder.modules import SparseStaticEmbedding
except ImportError:
    from sentence_transformers.sparse_encoder.models import SparseStaticEmbedding


class VSPLADEStaticEmbedding(SparseStaticEmbedding):
    """Inference-free Li-LSR query encoder for V-SPLADE.

    Behaves like :class:`SparseStaticEmbedding` with two differences, matching
    ``InferenceFreeQueryEncoder.encode_with_lookup`` from
    https://github.com/naver/v-splade:

    * repeated query tokens accumulate their weight (scatter-add) instead of
      being counted once;
    * token ids outside the lookup table (the 40 added vision tokens, e.g.
      ``<image>``) contribute nothing instead of raising an index error;
    * the lookup table covers the base (MLM) vocabulary (50368 entries), which
      is smaller than the full tokenizer vocabulary, so its size is stored in
      the module config (``num_dimensions``) for loading.
    """

    config_keys: list[str] = ["frozen", "num_dimensions"]

    def __init__(self, tokenizer, weight: torch.Tensor | None = None, frozen: bool = False, num_dimensions: int | None = None):
        if weight is None and num_dimensions is not None:
            weight = torch.zeros(num_dimensions)
        super().__init__(tokenizer=tokenizer, weight=weight, frozen=frozen)

    def forward(self, features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        input_ids = features["input_ids"]
        attention_mask = features["attention_mask"]

        valid = (input_ids < self.num_dimensions) & (attention_mask > 0)
        safe_ids = input_ids.clamp(max=self.num_dimensions - 1)
        scores = self.weight[safe_ids] * valid.to(self.weight.dtype)

        embeddings = torch.zeros(
            input_ids.size(0), self.num_dimensions, device=input_ids.device, dtype=self.weight.dtype
        )
        embeddings.scatter_add_(1, safe_ids, scores)

        features["sentence_embedding"] = embeddings
        return features


__all__ = ["VSPLADEStaticEmbedding"]