"""Word-level embeddings from a subword transformer encoder.

Port of AllenNLP's `pretrained_transformer_mismatched` (indexer + embedder pair): tokenize
each word into subwords, run the transformer, mean-pool subwords back to one vector per word
(see training/README.md: word tokenization is fixed upstream, so the model always sees
word-level spans). Long documents are split into non-overlapping, word-aligned chunks of at
most `max_length` transformer positions (the medbit arm's "context folding": medBIT-r3-plus's
native window is 512 tokens, most reports are longer).

Deliberate deviation from v1: AllenNLP folds the flat *wordpiece* stream into fixed windows,
so a word's subwords can occasionally land split across two windows, encoded with no shared
context. Here chunks are packed at word granularity, so a word's subwords are always encoded
together -- strictly better representation quality at chunk boundaries, same architecture.

Each chunk is re-tokenized and special-token-wrapped through the tokenizer's own public
`__call__(..., is_split_into_words=True)` API rather than hand-rolling special-token insertion:
transformers 5.x dropped the old `build_inputs_with_special_tokens` / `prepare_for_model`
helpers from its fast-tokenizer wrapper (see training_v2/README.md), and going through the
stable public API is more robust across tokenizer families anyway.
"""
import logging
from typing import List

import torch
from torch import nn
from transformers import AutoConfig, AutoModel, AutoTokenizer

logger = logging.getLogger(__name__)


class MismatchedEmbedder(nn.Module):
    def __init__(self, model_name: str, max_length: int, train_parameters: bool = True):
        super().__init__()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        # Built from config, not `AutoModel.from_pretrained`: this HF-packaged copy is only ever
        # constructed either directly (fresh pretrained-encoder init, matching training_v2's
        # behavior) or via `RadgraphModel.from_pretrained(...)`, which builds the whole module
        # tree under a `torch.device("meta")` context and then loads our fine-tuned
        # model.safetensors over it -- a nested `from_pretrained` call in that path hits
        # transformers' "meta device anti-pattern" guard. Building from config sidesteps that
        # and is equivalent here since the encoder weights get overwritten by the checkpoint
        # load either way.
        self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(model_name))
        self.max_length = max_length
        self.output_dim = self.encoder.config.hidden_size
        if not train_parameters:
            for p in self.encoder.parameters():
                p.requires_grad = False

        num_specials = self.tokenizer.num_special_tokens_to_add(pair=False)
        budget = max_length - num_specials
        if budget <= 0:
            raise ValueError(f"max_length={max_length} too small for {model_name}'s "
                              f"{num_specials} special tokens")
        self._budget = budget

    def get_output_dim(self) -> int:
        return self.output_dim

    def _word_subword_counts(self, words: List[str]) -> List[int]:
        enc = self.tokenizer(words, is_split_into_words=True, add_special_tokens=False)
        counts = [0] * len(words)
        for w in enc.word_ids():
            if w is not None:
                counts[w] += 1
        return counts

    def _pack_chunks(self, counts: List[int]) -> List[List[int]]:
        """Greedily group word indices into chunks whose total subword count fits the budget."""
        chunks: List[List[int]] = []
        current: List[int] = []
        current_len = 0
        for w, n in enumerate(counts):
            n = max(n, 1)
            if n > self._budget:
                logger.warning("word %d has %d subwords > budget %d; it will be truncated "
                                "by the tokenizer", w, n, self._budget)
            if current and current_len + n > self._budget:
                chunks.append(current)
                current, current_len = [], 0
            current.append(w)
            current_len += n
        if current:
            chunks.append(current)
        return chunks

    def forward(self, words: List[str]) -> torch.Tensor:
        """Returns (num_words, hidden_size)."""
        device = next(self.encoder.parameters()).device
        counts = self._word_subword_counts(words)
        chunks = self._pack_chunks(counts)

        word_embeddings: List[torch.Tensor] = [None] * len(words)  # type: ignore[list-item]
        for chunk_word_ixs in chunks:
            chunk_words = [words[i] for i in chunk_word_ixs]
            enc = self.tokenizer(chunk_words, is_split_into_words=True, add_special_tokens=True,
                                  truncation=True, max_length=self.max_length, return_tensors="pt")
            input_ids = enc["input_ids"].to(device)
            attention_mask = enc["attention_mask"].to(device)
            hidden = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[0]

            word_ids = enc.word_ids(batch_index=0)
            positions_by_word: dict = {}
            for pos, w in enumerate(word_ids):
                if w is not None:
                    positions_by_word.setdefault(w, []).append(pos)
            for local_w, positions in positions_by_word.items():
                word_embeddings[chunk_word_ixs[local_w]] = hidden[positions].mean(dim=0)

        for i, emb in enumerate(word_embeddings):
            if emb is None:  # word tokenized to zero subwords (rare; e.g. an unknown char)
                word_embeddings[i] = hidden.new_zeros(self.output_dim)

        return torch.stack(word_embeddings, dim=0)