"""LFM2 backbone with bidirectional attention + non-causal short-conv, for retrieval/embedding use.

Wired into the HF repo via `auto_map` in config.json so that

    AutoModel.from_pretrained(repo, trust_remote_code=True)
    SentenceTransformer(repo, trust_remote_code=True)

both return a model with the encoder-style patches already applied.

Supports `attn_implementation` in {"eager", "sdpa", "flash_attention_2"}:
eager/sdpa consume a 4D additive pad-only mask and reproduce the exact
training-time behavior; flash_attention_2 receives the 2D padding mask (or
None) and runs the kernel non-causally via `Lfm2Attention.is_causal = False`,
yielding outputs equivalent to the unpadded forward.

Repos may set `"disable_flash_attention": true` in config.json to reject
flash_attention_2 at load time (used for ColBERT, where PyLate query expansion
tokens — attention_mask=0 but scored in MaxSim — are incompatible with FA2
unpadding and severely degrade retrieval quality).
"""

from typing import Optional

import torch
import torch.nn.functional as F
from transformers.models.lfm2 import modeling_lfm2 as _lfm2_mod
from transformers.models.lfm2.modeling_lfm2 import (
    Lfm2Attention,
    Lfm2Model,
    Lfm2ShortConv,
    apply_mask_to_padding_states,
)


def _bidirectional_mask(config, **kwargs) -> Optional[torch.Tensor]:
    # transformers has renamed the embeds kwarg across versions
    # (input_embeds <-> inputs_embeds); accept either to stay forward-compatible.
    embeds = kwargs.get("inputs_embeds")
    if embeds is None:
        embeds = kwargs.get("input_embeds")
    attention_mask = kwargs.get("attention_mask")
    past_key_values = kwargs.get("past_key_values")

    if config._attn_implementation == "flash_attention_2":
        # FA2 only uses the 2D padding mask to unpad sequences; causality is
        # controlled by `Lfm2Attention.is_causal` (set to False below).
        if attention_mask is not None and not attention_mask.all():
            return attention_mask
        return None

    device = embeds.device
    dtype = embeds.dtype
    bsz, q_len = embeds.shape[:2]
    past = past_key_values.get_seq_length() if past_key_values is not None else 0
    kv_len = past + q_len

    mask = torch.zeros((bsz, 1, q_len, kv_len), device=device, dtype=dtype)
    if attention_mask is not None:
        cur_len = attention_mask.size(-1)
        key_pad_flags = (attention_mask == 0).to(device=device, dtype=torch.float32)
        pad_vec = torch.zeros((bsz, kv_len), device=device, dtype=torch.float32)
        if cur_len > 0:
            pad_vec[:, past:past + cur_len] = key_pad_flags * -1e9
        mask = mask + pad_vec.to(dtype)[:, None, None, :]
    return mask


def _noncausal_shortconv_forward(
    self,
    hidden_states: torch.Tensor,
    past_key_values=None,
    cache_position=None,
    attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    # eager/sdpa pass the 4D additive mask, on which this is a no-op — matching
    # the behavior the checkpoints were trained with. FA2 passes the 2D pad
    # mask, so pads are zeroed before the conv (closest match to the unpadded
    # forward, since FA2 cannot reproduce the padded sdpa pad-state evolution).
    x = apply_mask_to_padding_states(hidden_states, attention_mask)

    BCx = self.in_proj(x).transpose(-1, -2)
    B, C, x = BCx.chunk(3, dim=-2)
    Bx = B * x

    k = self.conv.weight.shape[-1]
    pad = k // 2
    conv_out = F.conv1d(
        Bx, weight=self.conv.weight, bias=self.conv.bias,
        stride=1, padding=pad, dilation=1, groups=Bx.shape[1],
    )
    if conv_out.shape[-1] > Bx.shape[-1]:
        conv_out = conv_out[..., :Bx.shape[-1]]
    elif conv_out.shape[-1] < Bx.shape[-1]:
        conv_out = F.pad(conv_out, (0, Bx.shape[-1] - conv_out.shape[-1]))

    y = C * conv_out
    y = y.transpose(-1, -2).contiguous()
    return self.out_proj(y)


def _shortconv_forward(self, *args, **kwargs):
    return self.slow_forward(*args, **kwargs)


_PATCHED = False


def _install_patches() -> None:
    global _PATCHED
    if _PATCHED:
        return
    _lfm2_mod.create_causal_mask = _bidirectional_mask
    Lfm2ShortConv.slow_forward = _noncausal_shortconv_forward
    Lfm2ShortConv.forward = _shortconv_forward
    _PATCHED = True


_install_patches()


class Lfm2BidirectionalModel(Lfm2Model):
    """LFM2 patched for encoder-style use: full bidirectional attention + non-causal short-conv."""

    def __init__(self, config):
        if (
            getattr(config, "_attn_implementation", None) == "flash_attention_2"
            and getattr(config, "disable_flash_attention", False)
        ):
            raise ValueError(
                "flash_attention_2 is disabled for this model: query expansion "
                "tokens (attention_mask=0 but scored in MaxSim) are incompatible "
                "with FA2 unpadding and severely degrade retrieval quality. "
                "Load with attn_implementation='sdpa' (default) or 'eager'."
            )
        _install_patches()
        super().__init__(config)
        for module in self.modules():
            if isinstance(module, Lfm2Attention):
                module.is_causal = False