"""Ettin-68M (jhu-clsp/ettin-encoder-68m, MIT) — ModernBERT-энкодер, hidden 512, byte-level BPE.

Взят из релиза Agate. Из файла релиза оставлен только Ettin: вариантов T5 и Hash-заглушки из
FCDM/T2I-ветки (и функции для разморозки top-N блоков T5) в проекте нет — модель заморожена.

Держит регистр, пунктуацию, диакритику и нелатинские названия, которые SentencePiece у Flan-T5
превращает в <unk>. Пустая строка кодируется как [CLS][SEP] — это безусловный контекст для CFG.
"""
from __future__ import annotations

import torch


def _load_modernbert(path: str):
    """AutoModel.from_pretrained, или — там, где transformers отказывается читать .bin на torch < 2.6
    (Ettin поставляется только с pytorch_model.bin) — сборка из конфига и загрузка state_dict с
    weights_only=True. Тензоры MLM-головы отбрасываются."""
    from transformers import AutoConfig, AutoModel
    try:
        return AutoModel.from_pretrained(path, torch_dtype=torch.float32)
    except ValueError:
        import os
        m = AutoModel.from_config(AutoConfig.from_pretrained(path))
        f = os.path.join(path, "pytorch_model.bin") if os.path.isdir(path) else None
        if f is None:
            from huggingface_hub import hf_hub_download
            f = hf_hub_download(path, "pytorch_model.bin")
        sd = {k.removeprefix("model."): v for k, v in torch.load(f, map_location="cpu", weights_only=True).items()}
        missing, _ = m.load_state_dict(sd, strict=False)
        if missing:
            raise KeyError(f"{path}: encoder tensors missing from the checkpoint: {missing[:5]}")
        return m


class EttinTextEncoder:
    """ModernBERT-энкодер Agate. train_blocks != 0 тренирует весь энкодер (опция «верхние блоки»
    бессмысленна: он маленький, и пространство обусловливания всё равно учится заново).

    Масок паддинга две: attention_mask в модели и mask наружу (1 = настоящий токен), по нему
    кросс-внимание мышления игнорирует паддинг."""

    def __init__(self, path: str, device, max_len: int = 128, train_blocks: int = 0, grad_ckpt: bool = False):
        from transformers import AutoTokenizer
        self.device, self.max_len = torch.device(device), max_len
        self.tok = AutoTokenizer.from_pretrained(path)
        self.model = _load_modernbert(path).to(self.device)
        self.model.eval().requires_grad_(bool(train_blocks))
        if grad_ckpt and train_blocks:
            # HF checkpoint'ит только в train()-режиме; у Ettin все dropout = 0.0 (в конфиге), поэтому
            # численно это ничего не меняет. Пересчёт слоёв в backward: обучаемые активации Ettin
            # занимали ~22 ГБ/GPU при 128 x 416 токенах без компиляции (OOM 2026-09-25).
            if any(getattr(self.model.config, k, 0) for k in ("attention_dropout", "mlp_dropout", "embedding_dropout")):
                raise ValueError("Ettin checkpointing needs train() mode, but this config has dropout")
            self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
            self.model.train()
        self.dim = self.model.config.hidden_size
        self.trainable = [(n, p) for n, p in self.model.named_parameters() if p.requires_grad]

    def tokenize(self, captions: list[str], pad_to: int = 1) -> tuple[torch.Tensor, torch.Tensor]:
        t = self.tok(captions, padding="longest", truncation=True, max_length=self.max_len,
                     return_tensors="pt")
        ids, mask = t["input_ids"], t["attention_mask"]
        extra = (-ids.shape[1]) % pad_to
        if extra:
            ids = torch.nn.functional.pad(ids, (0, extra), value=self.tok.pad_token_id)
            mask = torch.nn.functional.pad(mask, (0, extra))
        return ids, mask

    def encode(self, ids: torch.Tensor, mask: torch.Tensor, grad: bool = False
               ) -> tuple[torch.Tensor, torch.Tensor]:
        ids, mask = ids.to(self.device, non_blocking=True), mask.to(self.device, non_blocking=True)
        with torch.set_grad_enabled(grad and bool(self.trainable)), \
                torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
            h = self.model(input_ids=ids, attention_mask=mask).last_hidden_state
        return h.float(), mask

    def __call__(self, captions: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
        return self.encode(*self.tokenize(captions))